torchdyn

The torchdyn library provides a unified and flexible API for implementing continuous and implicit learning models. The torchdyn.models module and associated tutorials include implementations of several key architectures and strategies:

Core Architectures

• Neural Ordinary Differential Equations (Neural ODE)
• Galerkin Neural ODE
• Neural Stochastic Differential Equations (Neural SDE)
• Graph Neural ODEs
• Hamiltonian Neural Networks

Sequence and Hybrid Models

• ODE-RNN: A recurrent or 'hybrid' version designed for sequences.

Numerical Methods and Augmentation

• Hypersolvers: Neural numerical methods.
• Augmentation Strategies: Designed to increase expressivity and reduce computational burden on numerical solvers, including:
  • ANODE (0-augmentation)
  • Input-layer augmentation
  • Higher-order augmentation

Sensitivity Algorithms

• Integral loss adjoint variants.

The notebook demonstrates that torchdyn's gradient computation (using the adjoint method) is consistent with torchdiffeq. When calculating the gradient of a loss with respect to the time span T or the start time t0, the difference between torchdyn and torchdiffeq.odeint_adjoint should be negligible (near zero).

# torchdyn gradient
t_span = torch.cat([t0, T])
t_eval, traj = prob(x, t_span)
l = ((t_span[-1:] - torch.tensor([5]))**2).mean()
dldt_torchdyn = grad(l, T)[0]

# torchdiffeq gradient (requires wrapping vector field)
class VectorField(nn.Module):
    def __init__(self, f):
        super().__init__()
        self.f = f
    def forward(self, t, x):
        return self.f(x)

sys = VectorField(f)
t_span = torch.cat([t0, T])
t_eval_diff, traj_diff = torchdiffeq.odeint_adjoint(sys, x, t_span, method='dopri5', atol=1e-4, rtol=1e-4), None # simplified
l_diff = ((t_span[-1:] - torch.tensor([5]))**2).mean()
dldt_torchdiffeq = grad(l_diff, T)[0]

# Check consistency
print(dldt_torchdyn - dldt_torchdiffeq)

This variant combines stacked Neural ODEs with learned discrete transition maps (jumps) at the end of each interval. The state is updated by $z^+ = g(z(t), \omega_i)$ at the boundary $t = t_{i+1}$.

To implement this:

1. Define a sequence of NeuralODE flows.
2. Define a sequence of transition layers (e.g., nn.Linear).
3. In the forward pass, iterate through the flows, and after each flow, apply the corresponding jump layer to the last state of the trajectory before passing it to the next flow.

from torchdyn.core import NeuralODE
from torchdyn.nn import DataControl
import torch.nn as nn

num_pieces = 5

# stacked depth-invariant Neural ODEs
nde = []
for i in range(num_pieces):
    nde.append(NeuralODE(nn.Sequential(DataControl(),
                                      nn.Linear(4, 4), 
                                      nn.Tanh(), 
                                      nn.Linear(4, 2)), solver='rk4'))

# State "jump" parametrized by a linear layer
jumps = nn.Sequential(*[nn.Linear(2, 2) for _ in range(num_pieces)]).to(device)
flows = nn.Sequential(*nde).to(device)

t_span = torch.linspace(0, 1, 5)

class PCDST_NeuralODE(nn.Module):
    def __init__(self):
        super().__init__()
        self.flows = flows
        self.jumps = jumps
    def forward(self, x, t_span):
        for k, node in enumerate(self.flows):
            t_eval, traj = node(x, t_span)
            x = self.jumps[k](traj[-1])
        return t_eval, traj

model = PCDST_NeuralODE()

Stacked Neural ODEs approximate depth-varying weights by using piece-wise constant weights over $N$ intervals. This is equivalent to stacking $N$ depth-invariant Neural ODEs sequentially, where each ODE is solved over its own sub-interval $\Delta t_i$.

To implement this:

1. Create a list of NeuralODE instances (one for each piece).
2. Wrap them in an nn.Sequential or a custom nn.Module.
3. In the forward pass, iterate through the nodes, passing the output of one node as the input to the next.

Note: Since individual nodes are depth-invariant, you can solve them in the interval $[0, \Delta t_i]$.

from torchdyn.core import NeuralODE
from torchdyn.nn import DataControl
import torch.nn as nn

num_pieces = 5

# Stacked depth-invariant Neural ODEs
nde = []
for i in range(num_pieces):
    nde.append(NeuralODE(nn.Sequential(DataControl(),
                                      nn.Linear(4, 4), 
                                      nn.Tanh(), 
                                      nn.Linear(4, 2)), solver='rk4')) 
pc_node = nn.Sequential(*nde).to(device)

class PC_NeuralODE(nn.Module):
    def __init__(self):
        super().__init__()
        self.pc_node = pc_node
    def forward(self, x, t_span):
        for node in self.pc_node:
            t_eval, traj = node(x, t_span)
            x = traj[-1]
        return t_eval, traj

model = PC_NeuralODE()