Edifice.Energy.NeuralODE (Edifice v0.2.0)

Copy Markdown View Source

Neural ODE - Continuous-Depth Residual Networks.

Neural ODEs parameterize the continuous dynamics of hidden states using a neural network. Instead of discrete residual layers h_{t+1} = h_t + f(h_t), Neural ODEs solve the continuous-time ODE:

dh/dt = f(h(t), t; theta)

This provides:

  • Constant memory cost (via adjoint sensitivity method)
  • Adaptive computation (solver controls accuracy)
  • Continuous-depth networks (no fixed number of layers)

Architecture

Input [batch, input_size]
      |
      v
+--------------------------------------+
| Input Projection                     |
+--------------------------------------+
      |
      v
+--------------------------------------+
| ODE Solve: dh/dt = f(h, t)          |
|   f = MLP(h, t)                     |
|   Euler integration for num_steps    |
+--------------------------------------+
      |
      v
+--------------------------------------+
| Output Projection                    |
+--------------------------------------+
      |
      v
Output [batch, hidden_size]

Usage

model = NeuralODE.build(
  input_size: 256,
  hidden_size: 256,
  num_steps: 10,
  step_size: 0.1
)

References

Summary

Types

Options for build/1.

Functions

Build a Neural ODE model.

Build a Neural ODE with shared dynamics network.

Get the output size of a Neural ODE model.

Types

build_opt()

@type build_opt() ::
  {:hidden_size, pos_integer()}
  | {:input_size, pos_integer()}
  | {:num_steps, pos_integer()}
  | {:output_size, pos_integer()}
  | {:step_size, pos_integer()}

Options for build/1.

Functions

build(opts \\ [])

@spec build([build_opt()]) :: Axon.t()

Build a Neural ODE model.

Options

  • :input_size - Input feature dimension (required)
  • :hidden_size - Hidden state dimension (default: 256)
  • :num_steps - Number of Euler integration steps (default: 10)
  • :step_size - Integration step size (default: 0.1)
  • :output_size - Output dimension, nil to use hidden_size (default: nil)

Returns

An Axon model: [batch, input_size] -> [batch, output_size]

build_shared(opts)

@spec build_shared(keyword()) :: Axon.t()

Build a Neural ODE with shared dynamics network.

Delegates to build/1 which now always uses shared dynamics weights (the same f(h; theta) is re-evaluated at every integration step). Kept for backward compatibility.

Options

Same as build/1.

output_size(opts \\ [])

@spec output_size(keyword()) :: pos_integer()

Get the output size of a Neural ODE model.