Augmented Neural ODEs

Augmenting a neural ODE with extra state dimensions removes a fundamental expressiveness limit caused by trajectories that cannot cross.

A topological obstruction

A neural ODE transforms its input by flowing it continuously along a learned vector field. A basic property of such flows is that trajectories cannot intersect: two distinct starting points always map to distinct points at every later time, and the mapping is a continuous, invertible deformation of space, a homeomorphism. This is elegant but restrictive, because some functions require the input space to be folded or torn in ways a homeomorphism cannot achieve.

A concrete failure

Consider a one-dimensional task where points inside a small interval should map to one class and points outside to another, so the desired function is not monotonic and would require trajectories to cross. A plain neural ODE cannot represent this and must resort to a very stiff, hard-to-integrate vector field that only approximates it, driving up the number of solver steps as training proceeds.

Adding dimensions

Augmented neural ODEs solve this by appending extra zero-initialized dimensions to the state before integration. The ODE then flows in a higher-dimensional space where the required separation is easy, since trajectories that would have to cross in the original space can pass one another by moving through the added dimensions. After integration, the model reads out from the augmented state.

• Higher-dimensional flows can realize functions the original dimension forbids
• Trajectories become smoother, so adaptive solvers need fewer steps
• Training is often faster and more stable than the unaugmented model
• The extra dimensions add modest cost relative to the expressiveness gained

# augment: pad state with zeros before solving
h_aug = torch.cat([h0, torch.zeros(batch, aug_dim)], dim=-1)
h_T = odeint(f_aug, h_aug, t)[-1]
out = readout(h_T)

Why it matters

Augmented neural ODEs show that the continuous-depth framework's limits are not fundamental to the idea but to the dimensionality it is given. The fix mirrors a familiar trick elsewhere in machine learning: lift data into a higher-dimensional space to make a hard separation easy. It preserves the constant-memory adjoint training of the base neural ODE while widening the class of functions it can learn.