solver.compute_likelihood with enable_grad=True returns a trajectory and logprob whose gradients are set to None

[jbarraquand]

I tried to use solver.compute_likelihood() with enable_grad=True to use your library to test a Continous Normalising Flow loss (logprob) instead of a Flow Matching objective. The gradients of the trajectory and logprob returned are both set to None. The maximum likelihood training objective cannot be trained.

[itaigat]

Can you please attach an example and output?

[jbarraquand]

Here is a modified version of your 2D flow matching example adapted to use the negative log likelihood computed with solver.compute_likelihood, as well as the error log. The next comment is the error log, the following one is the code

[jbarraquand]

/usr/local/bin/python3.10 /Users/jerome/src/pycharm/mindwave/generative/MetaCNF2Dexample.py
Using cpu.
Traceback (most recent call last):
File "/Users/jerome/src/pycharm/mindwave/generative/MetaCNF2Dexample.py", line 151, in
loss.backward() # backward
File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/torch/_tensor.py", line 522, in backward
torch.autograd.backward(
File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/torch/autograd/init.py", line 266, in backward
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

Process finished with exit code 1

[jbarraquand]

import time
import torch

from torch import nn, Tensor

# flow_matching
from flow_matching.path.scheduler import CondOTScheduler
from flow_matching.path import AffineProbPath
from flow_matching.solver import Solver, ODESolver
from flow_matching.utils import ModelWrapper

from torch.distributions import Independent, Normal


# visualization
import matplotlib.pyplot as plt

from matplotlib import cm


# To avoide meshgrid warning
import warnings

warnings.filterwarnings("ignore", category=UserWarning, module='torch')

if torch.cuda.is_available():
    device = 'cuda:0'
    print('Using gpu')
else:
    device = 'cpu'
    print('Using cpu.')

torch.manual_seed(42)


def inf_train_gen(batch_size: int = 200, device: str = "cpu"):
    x1 = torch.rand(batch_size, device=device) * 4 - 2
    x2_ = torch.rand(batch_size, device=device) - torch.randint(high=2, size=(batch_size,), device=device) * 2
    x2 = x2_ + (torch.floor(x1) % 2)

    data = 1.0 * torch.cat([x1[:, None], x2[:, None]], dim=1) / 0.45

    #print(f"data shape {data.float().shape} data = {data.float()}")

    return data.float()


# Activation class
class Swish(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: Tensor) -> Tensor:
        return torch.sigmoid(x) * x


# Model class
class MLP(nn.Module):
    def __init__(self, input_dim: int = 2, time_dim: int = 1, hidden_dim: int = 128):
        super().__init__()

        self.input_dim = input_dim
        self.time_dim = time_dim
        self.hidden_dim = hidden_dim

        self.main = nn.Sequential(
            nn.Linear(input_dim + time_dim, hidden_dim),
            Swish(),
            nn.Linear(hidden_dim, hidden_dim),
            Swish(),
            nn.Linear(hidden_dim, hidden_dim),
            Swish(),
            nn.Linear(hidden_dim, hidden_dim),
            Swish(),
            nn.Linear(hidden_dim, input_dim),
        )

    def forward(self, x: Tensor, t: Tensor) -> Tensor:
        sz = x.size()
        x = x.reshape(-1, self.input_dim)
        t = t.reshape(-1, self.time_dim).float()

        t = t.reshape(-1, 1).expand(x.shape[0], 1)
        h = torch.cat([x, t], dim=1)
        output = self.main(h)

        return output.reshape(*sz)

# training arguments
lr = 0.001
#batch_size = 4096
batch_size = 4096
#iterations = 20001
iterations = 2000
#print_every = 2000
print_every = 20
#hidden_dim = 512
hidden_dim = 64


class WrappedModel(ModelWrapper):
    def forward(self, x: torch.Tensor, t: torch.Tensor, **extras):
        return self.model(x, t)

# step size for ode solver
step_size = 0.05

# velocity field model init
vf = MLP(input_dim=2, time_dim=1, hidden_dim=hidden_dim).to(device)
wrapped_vf = WrappedModel(vf)

# instantiate an affine path object
#path = AffineProbPath(scheduler=CondOTScheduler())

#instantiate the solver to compute likelihood loss
solver = ODESolver(velocity_model=wrapped_vf)  # create an ODESolver class
gaussian_log_density = Independent(Normal(torch.zeros(2, device=device), torch.ones(2, device=device)), 1).log_prob

# init optimizer
optim = torch.optim.Adam(vf.parameters(), lr=lr)



# train
start_time = time.time()
for i in range(iterations):
    optim.zero_grad()

    # sample data (user's responsibility): in this case, (X_0,X_1) ~ pi(X_0,X_1) = N(X_0|0,I)q(X_1)
    x_1 = inf_train_gen(batch_size=batch_size, device=device) # sample data
    x_0 = torch.randn_like(x_1).to(device)

    # sample time (user's responsibility)
    t = torch.rand(x_1.shape[0]).to(device)

    # sample probability path
    #path_sample = path.sample(t=t, x_0=x_0, x_1=x_1)

    # flow matching l2 loss
    #loss = torch.pow(vf(path_sample.x_t, path_sample.t) - path_sample.dx_t, 2).mean()

    #Replace flow matching l2 loss with negative log likelihood loss

    with torch.set_grad_enabled(True):
        traj, log_p = solver.compute_likelihood(x_1=x_1, method='midpoint', step_size=step_size,
                                                     exact_divergence=True, log_p0=gaussian_log_density,
                                                     enable_grad=True)
        loss = -torch.mean(log_p)

    # optimizer step
    loss.backward()  # backward
    optim.step()  # update

    # log loss
    if (i + 1) % print_every == 0:
        elapsed = time.time() - start_time
        print('| iter {:6d} | {:5.2f} ms/step | loss {:8.3f} '
              .format(i + 1, elapsed * 1000 / print_every, loss.item()))
        start_time = time.time()




norm = cm.colors.Normalize(vmax=50, vmin=0)

#batch_size = 50000  # batch size
batch_size = 50000
eps_time = 1e-2

time1 = time.time()
print(f"start sampling ODE solver\n")
T = torch.linspace(0,1,10)  # sample times
T = T.to(device=device)

x_init = torch.randn((batch_size, 2), dtype=torch.float32, device=device)
solver = ODESolver(velocity_model=wrapped_vf)  # create an ODESolver class
sol = solver.sample(time_grid=T, x_init=x_init, method='midpoint', step_size=step_size, return_intermediates=True)  # sample from the model

time2 = time.time()
print(f"finished sampling ODE solver\n {time2-time1}")
sol = sol.cpu().numpy()
T = T.cpu()

#fig, axs = plt.subplots(1, 10, figsize=(20, 20))
fig, axs = plt.subplots(1, 10, figsize=(10, 3))

for i in range(10):
    H = axs[i].hist2d(sol[i, :, 0], sol[i, :, 1], 300, range=((-5, 5), (-5, 5)))

    cmin = 0.0
    cmax = torch.quantile(torch.from_numpy(H[0]), 0.99).item()

    norm = cm.colors.Normalize(vmax=cmax, vmin=cmin)

    _ = axs[i].hist2d(sol[i, :, 0], sol[i, :, 1], 300, range=((-5, 5), (-5, 5)), norm=norm)

    axs[i].set_aspect('equal')
    axs[i].axis('off')
    axs[i].set_title('t= %.2f' % (T[i]))

plt.tight_layout()
plt.show()

from torch.distributions import Independent, Normal

# sample with likelihood
print(f"start computing likelihood\n")
time1 = time.time()
T = torch.tensor([1., 0.])  # sample times
T = T.to(device=device)

#grid_size = 200
grid_size = 50
x_1 = torch.meshgrid(torch.linspace(-5, 5, grid_size), torch.linspace(-5, 5, grid_size))
x_1 = torch.stack([x_1[0].flatten(), x_1[1].flatten()], dim=1).to(device)
# source distribution is an isotropic gaussian
gaussian_log_density = Independent(Normal(torch.zeros(2, device=device), torch.ones(2, device=device)), 1).log_prob

# compute log likelihood with unbiased hutchinson estimator, average over num_acc
num_acc = 10
log_p_acc = 0

for i in range(num_acc):
    _, log_p = solver.compute_likelihood(x_1=x_1, method='midpoint', step_size=step_size, exact_divergence=False, log_p0=gaussian_log_density)
    log_p_acc += log_p

log_p_acc /= num_acc

# compute with exact divergence
_, exact_log_p = solver.compute_likelihood(x_1=x_1, method='midpoint', step_size=step_size, exact_divergence=True, log_p0=gaussian_log_density)


likelihood = torch.exp(log_p_acc).cpu().reshape(grid_size, grid_size).detach().numpy()
exact_likelihood = torch.exp(exact_log_p).cpu().reshape(grid_size, grid_size).detach().numpy()

time2 = time.time()
print(f"finished computing likelihood {time2-time1}")
fig, axs = plt.subplots(1, 2,figsize=(10,10))

cmin = 0.0
cmax = 1/32 # 1/32 is the gt likelihood value

norm = cm.colors.Normalize(vmax=cmax, vmin=cmin)

axs[0].imshow(likelihood, extent=(-5, 5, -5, 5), origin='lower', cmap='viridis', norm=norm)
axs[0].set_title('Model Likelihood, Hutchinson Estimator, #acc=%d' % num_acc)

axs[1].imshow(exact_likelihood, extent=(-5, 5, -5, 5), origin='lower', cmap='viridis', norm=norm)
axs[1].set_title('Exact Model Likelihood')

fig.colorbar(cm.ScalarMappable(norm=norm, cmap='viridis'), ax=axs, orientation='horizontal', label='density')
plt.show()

[yiboyang]

I'm also running into this exact same problem. For now I'm looking into the FFJORD implementation for likelihood-based training: https://github.com/rtqichen/ffjord/tree/master