Skip to content

Maximum Likelihood Objective

Trains the same NeuralODEFlow by maximizing exact log-likelihood, using the Jacobian-trace term rather than velocity regression.

maximum_likelihood

Maximum likelihood objective for neural ODE flows.

MaximumLikelihoodObjective

Bases: Module

Negative log likelihood objective for NeuralODEFlow.

The objective uses the flow's exact log probability computation:

\[\mathcal{L} = -\mathbb{E}[\log p_\theta(x \mid c)]\]

By default, the negative log likelihood is divided by the flattened target dimension so training logs report NLL per dimension instead of total NLL.

Source code in flowpde/objectives/maximum_likelihood.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
class MaximumLikelihoodObjective(nn.Module):
    r"""
    Negative log likelihood objective for `NeuralODEFlow`.

    The objective uses the flow's exact log probability computation:

    $$\mathcal{L} = -\mathbb{E}[\log p_\theta(x \mid c)]$$

    By default, the negative log likelihood is divided by the flattened target
    dimension so training logs report NLL per dimension instead of total NLL.
    """

    def __init__(
        self,
        flow: NeuralODEFlow,
        regularization: float = 0.0,
        normalize_by_dim: bool = True,
        target_key: Optional[str] = None,
        condition_key: Optional[str] = None,
    ):
        super().__init__()
        self.flow = flow
        self.model = flow.model
        self.regularization = regularization
        self.normalize_by_dim = normalize_by_dim
        self.target_key = target_key or flow.target_key
        self.condition_key = condition_key or flow.condition_key

    def compute_loss(
        self,
        batch: Dict[str, Tensor],
        target_key: Optional[str] = None,
        condition_key: Optional[str] = None,
        **kwargs,
    ) -> Tensor:
        x, condition = self.flow._extract_target_condition(
            batch,
            target_key=target_key or self.target_key,
            condition_key=condition_key or self.condition_key,
        )
        target_dim = x.flatten(start_dim=1).shape[1]
        self.flow.set_target_dim(target_dim)
        log_px = self.flow.log_prob(x, condition)
        loss = -log_px.mean()
        if self.normalize_by_dim:
            loss = loss / target_dim

        if self.regularization > 0.0:
            t = torch.rand(x.shape[0], 1, device=self.flow.model_device)
            velocity = self.model(x.detach(), condition, t)
            loss = loss + self.regularization * velocity.pow(2).mean()

        return loss

    def get_config(self) -> Dict[str, Any]:
        return {
            "objective": "maximum_likelihood",
            "flow": self.flow.get_config(),
            "regularization": self.regularization,
            "normalize_by_dim": self.normalize_by_dim,
            "target_key": self.target_key,
            "condition_key": self.condition_key,
        }