Neural ODEs & Continuous-Depth Networks

Neural Ordinary Differential Equations (Neural ODEs) parameterize the derivative of a hidden state, enabling continuous-depth networks with constant memory cost. This module covers the ODE formulation, adjoint method for efficient backpropagation, variants, and connections to normalizing flows.

1. The ODE Formulation

1.1 Continuous Dynamics

A Neural ODE defines a continuous transformation:

where is the hidden state at "time" and is a neural network.

1.2 Initial Value Problem (IVP)

Given initial condition , the output is:

1.3 Connection to ResNets

A discrete ResNet with layers:

is the Euler discretization of the ODE with step size .

1.4 Neural ODE as a Layer

The full model:

where enc/dec are standard neural network layers.

2. Adjoint Method for Backpropagation

2.1 The Problem

Computing gradients through the ODE solver requires storing all intermediate states for backpropagation, which is memory-intensive.

2.2 Adjoint State

Define the adjoint state , which tracks the gradient of the loss with respect to the hidden state.

2.3 Adjoint ODE

The adjoint satisfies:

with terminal condition .

2.4 Parameter Gradients

2.5 Time Gradients (for varying T)

2.6 Implementation of Adjoint

def ode_adjoint(run, x, T, adjoint_method='dopri5'):
    # Forward pass
    h_T = run.odeint(x, T)
    
    # Backward pass using adjoint
    a_T = torch.autograd.grad(loss, h_T)
    
    # Solve adjoint ODE backward
    def adjoint_dynamics(t, h, a):
        h.requires_grad_(True)
        f = run.f(h, t)
        dfdh = torch.autograd.grad(f, h, grad_outputs=a, create_graph=True)[0]
        return -dfdh
    
    a_0 = odeint(adjoint_dynamics, a_T, [T, 0])
    
    # Parameter gradients
    def param_dynamics(t, h, a):
        h.requires_grad_(True)
        f = run.f(h, t)
        dfdtheta = torch.autograd.grad(f, run.parameters(), grad_outputs=a, create_graph=True)
        return dfdtheta
    
    grad_theta = quad(param_dynamics, h_T, a_T, [0, T])
    
    return a_0, grad_theta

3. ODE Solvers

3.1 Runge-Kutta Methods

RK4 (4th-order Runge-Kutta):

3.2 Adaptive Step-Size Solvers

Dormand-Prince (DOPRI5): 5th-order method with 4th-order error estimate.

Error tolerance:

Step size control:

3.3 Fixed-Step Solvers

For simpler problems, fixed-step Euler or RK4 may be sufficient:

4. Neural ODE Variants

4.1 Latent Neural ODE

Encode data to latent space, apply ODE, decode:

Training: Use ELBO for VAE-style training.

4.2 Augmented Neural ODE

Augment state space to handle non-homeomorphic flows:

where are auxiliary dimensions. This allows flows that cannot be represented by volume-preserving ODEs.

4.3 FFJORD: Free-Form Jacobian of Reversible Dynamics

FFJORD (Grathwohl et al., 2018) estimates the trace of the Jacobian using Hutchinson's estimator:

where .

Log-likelihood:

4.4 Neural ODE with Time-Varying Parameters

Allow to depend explicitly on time:

This enables learning time-varying dynamics.

5. Continuous Normalizing Flows

5.1 Change of Variables Formula

For a continuous transformation :

5.2 Instantaneous Change of Variables

This is a continuous version of the discrete change-of-variables formula.

5.3 Trace Estimation

The exact trace is . Hutchinson's estimator:

where or . Cost: .

5.4 Training Continuous Normalizing Flows

Maximum likelihood:

6. Flow Matching

6.1 Probability Flow ODE

Given a noise schedule and data distribution :

6.2 Conditional Flow Matching

Learn a velocity field that transports to :

where is the target velocity field.

6.3 Optimal Transport Flow

The optimal transport velocity field:

where and are paired samples.

6.4 Rectified Flow

Iteratively distill a flow to fewer steps:

7. Implementation

import torch
import torch.nn as nn
from torchdiffeq import odeint_adjoint as odeint

class NeuralODE(nn.Module):
    def __init__(self, dynamics, solver='dopri5', 
                 rtol=1e-3, atol=1e-4):
        super().__init__()
        self.dynamics = dynamics
        self.solver = solver
        self.rtol = rtol
        self.atol = atol
    
    def forward(self, x, t_span):
        return odeint(
            self.dynamics, x, t_span,
            method=self.solver,
            rtol=self.rtol,
            atol=self.atol
        )

class ODEDynamics(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(hidden_dim, 128),
            nn.Tanh(),
            nn.Linear(128, 128),
            nn.Tanh(),
            nn.Linear(128, hidden_dim)
        )
    
    def forward(self, t, h):
        return self.net(h)

class LatentODE(nn.Module):
    def __init__(self, input_dim, latent_dim, hidden_dim):
        super().__init__()
        self.encoder = nn.Linear(input_dim, latent_dim)
        self.dynamics = ODEDynamics(latent_dim)
        self.decoder = nn.Linear(latent_dim, input_dim)
        self.ode = NeuralODE(self.dynamics)
    
    def forward(self, x):
        z0 = self.encoder(x)
        t_span = torch.tensor([0.0, 1.0])
        z_T = self.ode(z0, t_span)[-1]
        return self.decoder(z_T)

class FFJORD(nn.Module):
    def __init__(self, dynamics, trace_estimator='hutchinson'):
        super().__init__()
        self.dynamics = dynamics
        self.trace_estimator = trace_estimator
    
    def forward(self, x):
        def dynamics_with_trace(t, state):
            h, _ = state
            h.requires_grad_(True)
            f = self.dynamics(t, h)
            
            if self.trace_estimator == 'hutchinson':
                v = torch.randn_like(h)
                trace = v.mul(
                    torch.autograd.grad(f, h, v, create_graph=True)[0]
                ).sum(dim=-1)
            else:
                trace = torch.diagonal(
                    torch.autograd.functional.jacobian(
                        lambda h: self.dynamics(t, h), h
                    )
                ).sum()
            
            return f, trace
        
        log_prob = 0
        t_span = torch.linspace(0, 1, 10)
        states = odeint(dynamics_with_trace, (x, torch.zeros(x.shape[0])), t_span)
        
        z_T, log_det = states[0][-1], states[1][-1].sum()
        log_p_z = -0.5 * (z_T ** 2 + torch.log(2 * torch.pi)).sum(dim=-1)
        
        return log_p_z + log_det

8. SVG: Neural ODE Continuous Depth

9. SVG: Adjoint Method Illustration

10. Comparison of Methods

11. Open Problems

• Stiff ODEs: Handling dynamics with multiple time scales
• Long-time integration: Improving accuracy for long ODE trajectories
• Discrete data: Bridging continuous dynamics with discrete observations
• Scaling: Scaling Neural ODEs to high-dimensional data
• Theoretical guarantees: Convergence and expressiveness bounds

References

1. Chen, R. T. Q., Rubanova, Y., Bettencourt, J., & Duvenaud, D. (2018). Neural Ordinary Differential Equations. NeurIPS.
2. Grathwohl, W., Chen, R. T. Q., Bettencourt, J., Sutskever, I., & Duvenaud, D. (2018). FFJORD: Free-Form Continuous Dynamics for Scalable Reversible Generative Models. ICML.
3. Li, X., Chen, R. T. Q., & Duvenaud, D. (2020). Improved Variational Inference with Inverse Autoregressive Flow. NeurIPS.
4. Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., & Le, M. (2023). Flow Matching for Generative Modeling. ICLR.
5. Liu, X., et al. (2023). Rectified Flow: A Distillation Approach for Fast Sampling. ICML.