Neural Ordinary Differential Equations (Neural ODEs)

ニューラル常微分方程式（Neural ODEs）は、固定された層の積み重ねではなく、隠れ状態が時間とともに連続的にどのように変化するかを定義するニューラルネットワークです。小規模なネットワークfが導関数dh/dt = f(h, t)を学習し、数値ODEソルバーが初期状態から出力が必要な時点まで導関数を積分します。Ricky Chen、Yulia Rubanova、Jesse Bettencourt、David Duvenaudは、2018年の論文でこの考え方を提案しました。この論文はNeurIPSで最優秀論文賞を受賞しました。

残差ブロックから連続深度へ#

残差ネットワークは、h(t + 1) = h(t) + f(h(t))として状態を更新します。これは、ステップサイズを1としてODEを解くオイラー法の1ステップに相当します。Neural ODEは無限に小さなステップを無限回行う極限を取り、積分をソルバーに任せます。SciPyのsolve_ivpなどの従来の数値計算ソフトウェアも、手書きの方程式について同じ種類の初期値問題を解きます。

この手法では「深さ」の意味が変わります。Dormand-Prince法など、Runge-Kutta系の適応型ソルバーは、ダイナミクスが急速に変化する場所では小さなステップを、滑らかな場所では大きなステップを取ります。順伝播のコストはソルバーが必要とする関数評価の回数であり、これは入力、学習したダイナミクス、誤差許容度によって異なります。推論時に許容度を緩めれば、再学習をせずに精度と引き換えに推論レイテンシを低減できます。

随伴法によるトレーニング#

ソルバーを通した通常の誤差逆伝播では、すべての中間状態を保存するため、メモリ使用量はソルバーのステップ数に応じて増加します。随伴感度法は、勾配を運ぶ2つ目のODEを時間を逆向きにして解き、その過程で状態を保存するのではなく再計算することで、この問題を回避します。深さに対するメモリコストは一定に保たれますが、計算量が増え、順方向と逆方向の軌跡に小さな数値差が生じます。

torchdiffeqパッケージは、PyTorch用の微分可能なソルバーを提供し、直接バックプロパゲーションにはodeintを、随伴法にはodeint_adjointを使用します。

import torch
from torch import nn
from torchdiffeq import odeint

torch.manual_seed(0)

# ネットワークは2D隠れ状態の導関数dh/dtを予測します
dynamics = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 2))

# センサーの断続的な読み取り値のように、不規則なタイムスタンプで状態を取得します
initial_state = torch.tensor([[1.0, 0.0]])
times = torch.tensor([0.0, 0.3, 0.35, 1.1, 2.0])
trajectory = odeint(lambda t, h: dynamics(h), initial_state, times)

# 勾配はソルバーを通ってネットワークの重みに逆伝播します
loss = trajectory[-1].square().sum()
loss.backward()
print(trajectory.squeeze(1))
print(dynamics[0].weight.grad.norm())

軌跡には、要求した各時点の状態が保持されます。間隔の狭い0.3と0.35から、1.1までの長い間隔までが含まれますが、内部のステップはソルバーが独自に選択します。ゼロではない勾配ノルムは、最終時点で評価した損失関数を使い、ソルバーを通してネットワークを学習できることを示しています。odeintをodeint_adjointに置き換えると、勾配のメモリ使用量が一定の方式に切り替わります。ただし、ソルバーがパラメーターを見つけられるよう、ダイナミクス関数はnn.Moduleである必要があります。

Neural ODEの適用先#

不規則にサンプリングされた時系列が、最も適した用途です。2019年にRubanovaらが提案した潜在ODEモデルは、観測シーケンスを隠れ状態にエンコードし、それを連続的に発展させるため、間隔が不均一な医療記録やセンサーログを固定グリッドに無理に合わせずにモデル化できます。リカレントニューラルネットワークは離散ステップでのみ更新されるため、欠落または不規則なタイムスタンプを処理するには追加の仕組みが必要です。これは時系列分析における主要な課題です。

トラッキングデータも同じ形状です。Ultralytics YOLO26のTrackモードは、各フレームの各トラックIDについて位置を出力します。物体がオクルージョンされた場合や見失われた場合は間隔が空きます。デフォルトのトラッカーは、失われた各トラックをtrack_bufferフレーム（デフォルトでは30フレーム）先まで保持し、1フレームごとに1回進む固定線形ODEである、一定速度の運動モデルを持つカルマンフィルターを使います。同じ位置から学習したNeural ODEは、運動自体を連続軌跡として学習します。マルチオブジェクトトラッキングでは、フレーム間や途切れた区間をまたいで軌跡を問い合わせることができます。

生成モデリングは、もう1つの主要な用途です。元のNeural ODE論文で開発され、FFJORDで拡張された連続正規化フローは、学習したベクトル場を積分して単純な分布を複雑な分布に変換します。その後のフローマッチングでは、学習時にODEを解かず、ベクトル場を直接学習することで、これらのモデルの学習コストが大幅に低下しました。Stable Diffusion 3などのRectified Flow画像生成モデルは、このようなODEを積分してサンプリングします。また、拡散モデルは確率フローODEを積分することで決定論的にサンプリングできます。

制限事項と関連する連続時間モデル#

ODEの軌跡は交差できないため、入力と同じ次元のNeural ODEでは、単純な写像の一部を表現できません。Emilien Dupontらが2019年に提案した拡張Neural ODEは、状態に次元を追加することでこの問題を解決します。スティッフなダイナミクスや、学習時に増加する関数評価回数により、Neural ODEは固定深度のネットワークより遅くなることがあります。

ほかのアーキテクチャも、異なるトレードオフで連続時間ダイナミクスを利用しています。リキッドニューラルネットワークは、入力に応じて時定数が変化するODEで定義された再帰セルです。S4やMambaなどの状態空間モデルは、線形連続時間システムx'(t) = Ax(t) + Bu(t)から始め、それを再帰形式に離散化します。この形式は長い畳み込みや並列スキャンとしても計算できます。そのため、適応型ソルバーを必要とする一般的なNeural ODEよりも、長い系列を高速に処理できます。