Neural ODEs: Continuous-Depth Networks

Module: Machine Learning | Difficulty: Advanced

Neural ODE

Adjoint Method

where satisfies

Connection to ResNet

as step size .

Memory Cost

| Method | Memory | Time | |--------|--------|------| | Backprop | layers | | | Adjoint | | adjoint solves |

import torch
import torchdiffeq

class NeuralODE(nn.Module):
    def __init__(self, f, method='dopri5'):
        super().__init__()
        self.f = f; self.method = method
    def forward(self, x, t_span):
        return torchdiffeq.odeint(self.f, x, t_span, method=self.method)

class ODEFunc(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, 64), nn.Tanh(),
            nn.Linear(64, dim))
    def forward(self, t, x):
        return self.net(x)

Research Insight: Neural ODEs are the continuous analogue of ResNets. The key advantage is memory efficiency (constant in depth via the adjoint method), but they are slower to train due to ODE solver overhead. They naturally model residual connections.