Neural ODE（Neural Ordinary Differential Equation）は、ResNet（残差ネットワーク）の離散的な残差接続を常微分方程式（ODE）の連続極限として捉え直す画期的なアイデアです。Chen et al.（2018）による “Neural Ordinary Differential Equations” で提案され、NeurIPS 2018のBest Paperに選出されました。
Neural ODEは、ニューラルネットワークの層を「連続時間における状態の発展」として定式化し、ODEソルバーで順伝播を行います。さらに、随伴法（adjoint method）を用いることで、中間状態を保持せずにメモリ効率の良い勾配計算が可能になります。
本記事の内容
- ResNetの離散残差接続からODE解釈への発展
- Neural ODEの順伝播（ODEソルバーによる状態の積分）
- 随伴法による効率的勾配計算の導出
- 連続正規化フロー（CNF）への応用
- Pythonでスパイラルデータの学習・軌道の可視化
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。
ResNetからODEへ
ResNetの残差接続
ResNet（He et al., 2016）の各層は、以下の残差接続で定義されます。
$$ \bm{h}_{t+1} = \bm{h}_t + f(\bm{h}_t, \bm{\theta}_t) $$
ここで $\bm{h}_t$ は $t$ 番目の層の隠れ状態、$f$ は残差ブロック（畳み込み層+活性化関数など）、$\bm{\theta}_t$ は $t$ 番目の層のパラメータです。
この更新式は、オイラー法（ステップ幅 $\Delta t = 1$）の離散化に他なりません。
$$ \bm{h}_{t+\Delta t} = \bm{h}_t + \Delta t \cdot f(\bm{h}_t, \bm{\theta}_t), \quad \Delta t = 1 $$
連続極限
$\Delta t \to 0$ の極限を取ると、離散的な層の更新は連続時間のODEになります。
$$ \frac{d\bm{h}(t)}{dt} = f(\bm{h}(t), t, \bm{\theta}) $$
ここで重要な変化があります。
- 離散的な層番号 $t = 0, 1, 2, \dots$ が連続時間 $t \in [0, T]$ になる
- 各層ごとの独立なパラメータ $\bm{\theta}_t$ が、時間に依存する共有パラメータ $\bm{\theta}$ になる
- 層の数 $L$ が連続的な積分区間 $[0, T]$ になる
Neural ODEでは、$f$ をニューラルネットワーク（パラメータ $\bm{\theta}$）で表現します。入力 $\bm{h}(0)$ が与えられたとき、出力 $\bm{h}(T)$ は以下の初期値問題の解として得られます。
$$ \bm{h}(T) = \bm{h}(0) + \int_0^T f(\bm{h}(t), t, \bm{\theta}) \, dt $$
この積分は解析的には求められないので、ODEソルバー（例: Runge-Kutta法、Dormand-Prince法）を使って数値的に計算します。
Neural ODEの順伝播
Neural ODEの順伝播は以下のステップで行われます。
- 入力データ $\bm{x}$ を初期状態 $\bm{h}(t_0) = \bm{x}$ とする（必要に応じて線形射影で次元を変換）
- ODEソルバーで $\bm{h}(t_0)$ から $\bm{h}(t_1)$ を計算する
- $\bm{h}(t_1)$ を出力層に通して最終予測を得る
$$ \bm{h}(t_1) = \text{ODESolve}(\bm{h}(t_0), f, \bm{\theta}, t_0, t_1) $$
ODEソルバーとして4次ルンゲ＝クッタ法（RK4）を使う場合、1ステップの計算は以下です。
$$ \begin{align} \bm{k}_1 &= f(\bm{h}_n, t_n, \bm{\theta}) \\ \bm{k}_2 &= f\left(\bm{h}_n + \frac{\Delta t}{2}\bm{k}_1, t_n + \frac{\Delta t}{2}, \bm{\theta}\right) \\ \bm{k}_3 &= f\left(\bm{h}_n + \frac{\Delta t}{2}\bm{k}_2, t_n + \frac{\Delta t}{2}, \bm{\theta}\right) \\ \bm{k}_4 &= f(\bm{h}_n + \Delta t \bm{k}_3, t_n + \Delta t, \bm{\theta}) \\ \bm{h}_{n+1} &= \bm{h}_n + \frac{\Delta t}{6}(\bm{k}_1 + 2\bm{k}_2 + 2\bm{k}_3 + \bm{k}_4) \end{align} $$
適応的ステップ幅制御（Dormand-Prince法など）を使えば、ODEの解の滑らかさに応じて自動的にステップ幅が調整されます。解が急変する領域では細かいステップ、滑らかな領域では粗いステップが使われ、精度と効率のバランスが取れます。
ResNetとの比較
| 性質 | ResNet | Neural ODE |
|---|---|---|
| 層数 | 固定（$L$ 層） | 連続（ODEソルバーが適応的に決定） |
| メモリ | $O(L)$（全中間層の保持） | $O(1)$（随伴法使用時） |
| パラメータ | 層ごとに独立 | 全時刻で共有 |
| 評価精度 | 固定 | ODEソルバーの許容誤差で制御 |
逆伝播の問題
通常のニューラルネットワークの逆伝播では、損失関数 $L$ の勾配 $\frac{\partial L}{\partial \bm{\theta}}$ を計算するために、順伝播中のすべての中間状態を保持する必要があります。
Neural ODEで同様のアプローチを取ると、ODEソルバーの各ステップの中間状態を保持する必要があり、メモリ消費量がソルバーのステップ数 $N_{\text{steps}}$ に比例します。
$$ \text{メモリ消費量} \propto N_{\text{steps}} \times \dim(\bm{h}) $$
適応的ソルバーではステップ数が予測不可能であり、長い時間区間を積分する場合にメモリが不足する問題があります。
随伴法（Adjoint Method）による効率的勾配計算
随伴法はこの問題を解決する手法で、一定のメモリ量で勾配を計算できます。
随伴状態の定義
損失関数 $L(\bm{h}(t_1))$ のパラメータ $\bm{\theta}$ に対する勾配を計算したいとします。まず、随伴状態（adjoint state）を以下のように定義します。
$$ \bm{a}(t) = \frac{\partial L}{\partial \bm{h}(t)} $$
これは時刻 $t$ における隠れ状態 $\bm{h}(t)$ に対する損失の感度を表すベクトルです。
随伴状態のODE導出
$\bm{a}(t)$ が満たすODEを導出します。$\bm{h}(t)$ は以下のODEに従います。
$$ \frac{d\bm{h}}{dt} = f(\bm{h}(t), t, \bm{\theta}) $$
連鎖律を適用して、$\bm{a}(t)$ の時間微分を計算します。微小時間 $\epsilon > 0$ を考えると、$\bm{h}(t + \epsilon)$ は以下のように近似されます。
$$ \bm{h}(t + \epsilon) = \bm{h}(t) + \epsilon \cdot f(\bm{h}(t), t, \bm{\theta}) + O(\epsilon^2) $$
ここで、$\bm{a}(t)$ を $\bm{h}(t+\epsilon)$ を介して表すと、
$$ \bm{a}(t) = \frac{\partial L}{\partial \bm{h}(t)} = \frac{\partial L}{\partial \bm{h}(t + \epsilon)} \cdot \frac{\partial \bm{h}(t + \epsilon)}{\partial \bm{h}(t)} $$
右辺の第2因子を計算します。
$$ \frac{\partial \bm{h}(t + \epsilon)}{\partial \bm{h}(t)} = \bm{I} + \epsilon \frac{\partial f(\bm{h}(t), t, \bm{\theta})}{\partial \bm{h}(t)} + O(\epsilon^2) $$
したがって、
$$ \bm{a}(t) = \bm{a}(t + \epsilon) \left(\bm{I} + \epsilon \frac{\partial f}{\partial \bm{h}}\bigg|_{t}\right) + O(\epsilon^2) $$
$$ \bm{a}(t) = \bm{a}(t + \epsilon) + \epsilon \cdot \bm{a}(t + \epsilon) \frac{\partial f}{\partial \bm{h}}\bigg|_{t} + O(\epsilon^2) $$
両辺から $\bm{a}(t + \epsilon)$ を引いて $\epsilon$ で割り、$\epsilon \to 0$ の極限を取ると、
$$ \frac{\bm{a}(t) – \bm{a}(t + \epsilon)}{\epsilon} = \bm{a}(t+\epsilon) \frac{\partial f}{\partial \bm{h}}\bigg|_t + O(\epsilon) $$
$$ -\frac{d\bm{a}(t)}{dt} = \bm{a}(t) \frac{\partial f}{\partial \bm{h}}\bigg|_t $$
よって、随伴状態のODEは以下になります。
$$ \boxed{\frac{d\bm{a}(t)}{dt} = -\bm{a}(t) \frac{\partial f(\bm{h}(t), t, \bm{\theta})}{\partial \bm{h}}} $$
終端条件は $\bm{a}(t_1) = \frac{\partial L}{\partial \bm{h}(t_1)}$ です。
このODEは時間を逆方向に（$t_1$ から $t_0$ へ）解く必要があります。
パラメータに対する勾配
パラメータ $\bm{\theta}$ に対する勾配は、随伴状態を使って以下のように計算されます。
$$ \frac{dL}{d\bm{\theta}} = -\int_{t_1}^{t_0} \bm{a}(t) \frac{\partial f(\bm{h}(t), t, \bm{\theta})}{\partial \bm{\theta}} \, dt $$
この式の導出を確認しましょう。損失 $L$ は $\bm{h}(t_1)$ を通じて $\bm{\theta}$ に依存しますが、$\bm{h}(t_1)$ は $\bm{\theta}$ の関数です。全微分を連鎖律で書くと、
$$ \frac{dL}{d\bm{\theta}} = \frac{\partial L}{\partial \bm{h}(t_1)} \frac{d\bm{h}(t_1)}{d\bm{\theta}} $$
$\frac{d\bm{h}(t_1)}{d\bm{\theta}}$ を計算するために、ODE $\frac{d\bm{h}}{dt} = f(\bm{h}, t, \bm{\theta})$ の両辺を $\bm{\theta}$ で微分します。
$$ \frac{d}{dt}\frac{d\bm{h}}{d\bm{\theta}} = \frac{\partial f}{\partial \bm{h}}\frac{d\bm{h}}{d\bm{\theta}} + \frac{\partial f}{\partial \bm{\theta}} $$
$\bm{s}(t) = \frac{d\bm{h}(t)}{d\bm{\theta}}$ と置くと、$\bm{s}(t)$ は以下のODEを満たします（感度方程式）。
$$ \frac{d\bm{s}}{dt} = \frac{\partial f}{\partial \bm{h}} \bm{s} + \frac{\partial f}{\partial \bm{\theta}}, \quad \bm{s}(t_0) = \bm{0} $$
この感度方程式を直接解くこともできますが、$\bm{s}$ は $\dim(\bm{h}) \times \dim(\bm{\theta})$ の行列なので、パラメータ数が多い場合にはメモリ効率が悪くなります。
随伴法では、感度方程式を解く代わりに、随伴ODEを逆方向に解きながら勾配を累積します。具体的には、以下の拡張状態を同時に逆方向に積分します。
$$ \frac{d}{dt}\begin{bmatrix} \bm{h}(t) \\ \bm{a}(t) \\ \frac{dL}{d\bm{\theta}}(t) \end{bmatrix} = \begin{bmatrix} f(\bm{h}(t), t, \bm{\theta}) \\ -\bm{a}(t) \frac{\partial f}{\partial \bm{h}} \\ -\bm{a}(t) \frac{\partial f}{\partial \bm{\theta}} \end{bmatrix} $$
この拡張ODEを $t_1$ から $t_0$ まで逆方向に1回解くだけで、$\bm{h}(t_0)$（逆方向に再構成された状態）、$\bm{a}(t_0)$（入力に対する勾配）、$\frac{dL}{d\bm{\theta}}$（パラメータ勾配）がすべて得られます。
メモリ効率
随伴法のメモリ消費量は $O(\dim(\bm{h}) + \dim(\bm{\theta}))$ であり、ODEソルバーのステップ数に依存しません。これは通常の逆伝播の $O(N_{\text{steps}} \times \dim(\bm{h}))$ と比較して大きな改善です。
ただし、逆方向に $\bm{h}(t)$ を再構成する必要があるため、数値誤差の蓄積に注意が必要です。実際にはチェックポイント法を併用して精度を保つことが多いです。
連続正規化フロー（CNF）
Neural ODEの重要な応用の一つが連続正規化フロー（Continuous Normalizing Flow, CNF）です。
通常の正規化フローは、単純な基底分布 $p_0(\bm{z}_0)$（例: 標準正規分布）を可逆な変換 $T$ で変形してデータ分布を表現します。変数変換の公式から、
$$ \log p_1(\bm{z}_1) = \log p_0(\bm{z}_0) – \log \left|\det \frac{\partial \bm{z}_1}{\partial \bm{z}_0}\right| $$
離散的な正規化フローでは、ヤコビアンの行列式の計算が $O(d^3)$ で計算コストが高くなります。
Neural ODEを使った連続正規化フローでは、状態の変化をODEで定義します。
$$ \frac{d\bm{z}(t)}{dt} = f(\bm{z}(t), t, \bm{\theta}) $$
このとき、対数確率密度の変化は瞬時変化率の公式（instantaneous change of variables formula）で記述されます。
$$ \frac{d \log p(\bm{z}(t))}{dt} = -\text{tr}\left(\frac{\partial f}{\partial \bm{z}}\right) $$
この式を導出しましょう。連続性方程式（continuity equation）から出発します。確率の保存則より、
$$ \frac{\partial p(\bm{z}, t)}{\partial t} = -\nabla \cdot (p(\bm{z}, t) f(\bm{z}, t)) $$
右辺を展開すると、
$$ \frac{\partial p}{\partial t} = -p \nabla \cdot f – f \cdot \nabla p $$
両辺を $p$ で割ると、
$$ \frac{1}{p}\frac{\partial p}{\partial t} = -\nabla \cdot f – \frac{f \cdot \nabla p}{p} $$
左辺は $\frac{\partial \log p}{\partial t}$ であり、右辺の第2項は $f \cdot \nabla \log p$ です。ラグランジュ微分（流れに沿った微分）を使うと、
$$ \frac{d \log p(\bm{z}(t))}{dt} = \frac{\partial \log p}{\partial t} + f \cdot \nabla \log p = -\nabla \cdot f = -\text{tr}\left(\frac{\partial f}{\partial \bm{z}}\right) $$
ここで $\nabla \cdot f = \sum_i \frac{\partial f_i}{\partial z_i} = \text{tr}\left(\frac{\partial f}{\partial \bm{z}}\right)$ です。
重要な点は、完全なヤコビアン行列の行列式 $O(d^3)$ ではなく、トレースのみ $O(d)$ で計算できることです。さらに、ハッチンソンのトレース推定量を使えば、$O(1)$ で不偏推定が可能です。
$$ \text{tr}\left(\frac{\partial f}{\partial \bm{z}}\right) = \mathbb{E}_{\bm{\epsilon}}\left[\bm{\epsilon}^\top \frac{\partial f}{\partial \bm{z}} \bm{\epsilon}\right], \quad \bm{\epsilon} \sim \mathcal{N}(\bm{0}, \bm{I}) $$
Pythonでの実装
ここからは、Neural ODEをPyTorchで実装し、スパイラルデータの学習と軌道の可視化を行います。
必要なライブラリ
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
# 再現性
torch.manual_seed(42)
np.random.seed(42)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
スパイラルデータの生成
def generate_spiral_data(n_samples=1000, noise=0.5):
"""2つのスパイラル軌道からのサンプルを生成"""
t = np.linspace(0, 4 * np.pi, n_samples // 2)
# スパイラル1
x1 = t * np.cos(t) + noise * np.random.randn(len(t))
y1 = t * np.sin(t) + noise * np.random.randn(len(t))
# スパイラル2（反転）
x2 = -t * np.cos(t) + noise * np.random.randn(len(t))
y2 = -t * np.sin(t) + noise * np.random.randn(len(t))
# 連結
X = np.vstack([np.column_stack([x1, y1]),
np.column_stack([x2, y2])])
labels = np.hstack([np.zeros(len(t)), np.ones(len(t))])
return X.astype(np.float32), labels.astype(np.int64)
X_data, y_data = generate_spiral_data(n_samples=1000, noise=0.3)
# データの可視化
plt.figure(figsize=(8, 8))
plt.scatter(X_data[y_data == 0, 0], X_data[y_data == 0, 1],
c='blue', s=10, alpha=0.6, label='Class 0')
plt.scatter(X_data[y_data == 1, 0], X_data[y_data == 1, 1],
c='red', s=10, alpha=0.6, label='Class 1')
plt.xlabel('$x_1$', fontsize=14)
plt.ylabel('$x_2$', fontsize=14)
plt.title('Spiral Data', fontsize=14)
plt.legend(fontsize=12)
plt.axis('equal')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
ODEソルバーの実装（RK4法）
def rk4_step(f, h, t, dt, **kwargs):
"""4次ルンゲ＝クッタ法の1ステップ"""
k1 = f(h, t, **kwargs)
k2 = f(h + dt / 2 * k1, t + dt / 2, **kwargs)
k3 = f(h + dt / 2 * k2, t + dt / 2, **kwargs)
k4 = f(h + dt * k3, t + dt, **kwargs)
return h + dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
def ode_solve(f, h0, t_span, n_steps=20, **kwargs):
"""ODEを指定区間で解く（RK4法）"""
t0, t1 = t_span
dt = (t1 - t0) / n_steps
h = h0
t = t0
trajectory = [h]
times = [t]
for _ in range(n_steps):
h = rk4_step(f, h, t, dt, **kwargs)
t = t + dt
trajectory.append(h)
times.append(t)
return torch.stack(trajectory), times
Neural ODEモデルの定義
class ODEFunc(nn.Module):
"""ODE右辺のニューラルネットワーク: dh/dt = f(h, t)"""
def __init__(self, hidden_dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(2, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, 2),
)
# 重み初期化
for m in self.net:
if isinstance(m, nn.Linear):
nn.init.orthogonal_(m.weight)
nn.init.zeros_(m.bias)
def forward(self, h, t):
return self.net(h)
class NeuralODE(nn.Module):
"""Neural ODEモデル"""
def __init__(self, ode_func, t_span=(0, 1), n_steps=20):
super().__init__()
self.ode_func = ode_func
self.t_span = t_span
self.n_steps = n_steps
# 分類用の出力層
self.classifier = nn.Linear(2, 2)
def forward(self, x, return_trajectory=False):
# ODEソルバーで順伝播
trajectory, times = ode_solve(
self.ode_func, x, self.t_span, self.n_steps
)
h_final = trajectory[-1] # 最終状態
# 分類
logits = self.classifier(h_final)
if return_trajectory:
return logits, trajectory, times
return logits
学習ループ
# データの準備
X_tensor = torch.tensor(X_data, dtype=torch.float32).to(device)
y_tensor = torch.tensor(y_data, dtype=torch.long).to(device)
# モデルの初期化
ode_func = ODEFunc(hidden_dim=64).to(device)
model = NeuralODE(ode_func, t_span=(0, 1), n_steps=20).to(device)
# 損失関数とオプティマイザ
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)
# ミニバッチ学習
from torch.utils.data import DataLoader, TensorDataset
dataset = TensorDataset(X_tensor, y_tensor)
loader = DataLoader(dataset, batch_size=128, shuffle=True)
# 学習ループ
epochs = 200
loss_history = []
acc_history = []
for epoch in range(1, epochs + 1):
model.train()
epoch_loss = 0
correct = 0
total = 0
for batch_x, batch_y in loader:
optimizer.zero_grad()
logits = model(batch_x)
loss = criterion(logits, batch_y)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
epoch_loss += loss.item() * batch_x.size(0)
_, predicted = logits.max(1)
correct += predicted.eq(batch_y).sum().item()
total += batch_y.size(0)
scheduler.step()
avg_loss = epoch_loss / total
accuracy = correct / total
loss_history.append(avg_loss)
acc_history.append(accuracy)
if epoch % 40 == 0:
print(f'Epoch {epoch}/{epochs}, Loss: {avg_loss:.4f}, Accuracy: {accuracy:.4f}')
学習曲線の可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
axes[0].plot(loss_history, 'b-', linewidth=1.5)
axes[0].set_xlabel('Epoch', fontsize=12)
axes[0].set_ylabel('Loss', fontsize=12)
axes[0].set_title('Training Loss', fontsize=14)
axes[0].grid(True, alpha=0.3)
axes[1].plot(acc_history, 'g-', linewidth=1.5)
axes[1].set_xlabel('Epoch', fontsize=12)
axes[1].set_ylabel('Accuracy', fontsize=12)
axes[1].set_title('Training Accuracy', fontsize=14)
axes[1].grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
決定境界の可視化
# 格子点での予測
x_min, x_max = X_data[:, 0].min() - 2, X_data[:, 0].max() + 2
y_min, y_max = X_data[:, 1].min() - 2, X_data[:, 1].max() + 2
xx, yy = np.meshgrid(np.linspace(x_min, x_max, 200),
np.linspace(y_min, y_max, 200))
grid = torch.tensor(np.column_stack([xx.ravel(), yy.ravel()]),
dtype=torch.float32).to(device)
model.eval()
with torch.no_grad():
logits_grid = model(grid)
probs = torch.softmax(logits_grid, dim=1)[:, 1].cpu().numpy()
preds = logits_grid.argmax(dim=1).cpu().numpy()
plt.figure(figsize=(8, 8))
plt.contourf(xx, yy, probs.reshape(xx.shape), levels=50,
cmap='RdBu_r', alpha=0.5)
plt.colorbar(label='P(Class 1)')
plt.scatter(X_data[y_data == 0, 0], X_data[y_data == 0, 1],
c='blue', s=10, alpha=0.6, edgecolors='none')
plt.scatter(X_data[y_data == 1, 0], X_data[y_data == 1, 1],
c='red', s=10, alpha=0.6, edgecolors='none')
plt.xlabel('$x_1$', fontsize=14)
plt.ylabel('$x_2$', fontsize=14)
plt.title('Neural ODE Decision Boundary', fontsize=14)
plt.axis('equal')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
状態空間の軌道可視化
# サンプルデータのODE軌道を可視化
model.eval()
n_viz = 50
indices = np.random.choice(len(X_data), n_viz, replace=False)
X_viz = torch.tensor(X_data[indices], dtype=torch.float32).to(device)
y_viz = y_data[indices]
with torch.no_grad():
_, trajectory, times = model(X_viz, return_trajectory=True)
# 軌道をNumPyに変換
traj_np = torch.stack([tr.cpu() for tr in trajectory]).numpy()
# traj_np shape: (n_steps+1, n_viz, 2)
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
# t=0 （初期状態）
ax = axes[0]
ax.scatter(traj_np[0, y_viz == 0, 0], traj_np[0, y_viz == 0, 1],
c='blue', s=30, label='Class 0')
ax.scatter(traj_np[0, y_viz == 1, 0], traj_np[0, y_viz == 1, 1],
c='red', s=30, label='Class 1')
ax.set_title('t = 0 (Input)', fontsize=14)
ax.set_xlabel('$h_1$', fontsize=12)
ax.set_ylabel('$h_2$', fontsize=12)
ax.legend(fontsize=10)
ax.grid(True, alpha=0.3)
ax.set_aspect('equal')
# t=0.5 （中間状態）
mid_idx = len(traj_np) // 2
ax = axes[1]
ax.scatter(traj_np[mid_idx, y_viz == 0, 0], traj_np[mid_idx, y_viz == 0, 1],
c='blue', s=30, label='Class 0')
ax.scatter(traj_np[mid_idx, y_viz == 1, 0], traj_np[mid_idx, y_viz == 1, 1],
c='red', s=30, label='Class 1')
ax.set_title(f't = {times[mid_idx]:.1f} (Intermediate)', fontsize=14)
ax.set_xlabel('$h_1$', fontsize=12)
ax.set_ylabel('$h_2$', fontsize=12)
ax.legend(fontsize=10)
ax.grid(True, alpha=0.3)
ax.set_aspect('equal')
# t=1 （最終状態）
ax = axes[2]
ax.scatter(traj_np[-1, y_viz == 0, 0], traj_np[-1, y_viz == 0, 1],
c='blue', s=30, label='Class 0')
ax.scatter(traj_np[-1, y_viz == 1, 0], traj_np[-1, y_viz == 1, 1],
c='red', s=30, label='Class 1')
ax.set_title('t = 1 (Output)', fontsize=14)
ax.set_xlabel('$h_1$', fontsize=12)
ax.set_ylabel('$h_2$', fontsize=12)
ax.legend(fontsize=10)
ax.grid(True, alpha=0.3)
ax.set_aspect('equal')
plt.suptitle('State Space Trajectories at Different Times', fontsize=15)
plt.tight_layout()
plt.show()
個別軌道のアニメーション的な可視化
# 軌道の流れを矢印付きで表示
fig, ax = plt.subplots(1, 1, figsize=(10, 10))
for i in range(n_viz):
traj_i = traj_np[:, i, :] # (n_steps+1, 2)
color = 'blue' if y_viz[i] == 0 else 'red'
alpha = 0.3
# 軌道を描画
ax.plot(traj_i[:, 0], traj_i[:, 1], color=color, alpha=alpha, linewidth=0.8)
# 始点
ax.scatter(traj_i[0, 0], traj_i[0, 1], c=color, s=30, marker='o',
alpha=0.7, edgecolors='none')
# 終点
ax.scatter(traj_i[-1, 0], traj_i[-1, 1], c=color, s=30, marker='x',
alpha=0.7)
ax.set_xlabel('$h_1$', fontsize=14)
ax.set_ylabel('$h_2$', fontsize=14)
ax.set_title('Neural ODE Trajectories (o: start, x: end)', fontsize=14)
ax.grid(True, alpha=0.3)
ax.set_aspect('equal')
plt.tight_layout()
plt.show()
ODE解法の比較: オイラー法 vs RK4法
def euler_step(f, h, t, dt, **kwargs):
"""オイラー法の1ステップ"""
return h + dt * f(h, t, **kwargs)
def ode_solve_method(f, h0, t_span, n_steps, method='rk4', **kwargs):
"""指定した方法でODEを解く"""
t0, t1 = t_span
dt = (t1 - t0) / n_steps
h = h0
t = t0
trajectory = [h]
step_fn = rk4_step if method == 'rk4' else euler_step
for _ in range(n_steps):
h = step_fn(f, h, t, dt, **kwargs)
t = t + dt
trajectory.append(h)
return torch.stack(trajectory)
# 真の軌道（RK4, 高解像度）を基準にする
sample_x = X_viz[:5]
with torch.no_grad():
traj_rk4_fine = ode_solve_method(
ode_func, sample_x, (0, 1), n_steps=100, method='rk4'
)
traj_rk4_coarse = ode_solve_method(
ode_func, sample_x, (0, 1), n_steps=10, method='rk4'
)
traj_euler = ode_solve_method(
ode_func, sample_x, (0, 1), n_steps=10, method='euler'
)
# 比較の可視化
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
methods = [
('Euler (10 steps)', traj_euler),
('RK4 (10 steps)', traj_rk4_coarse),
('RK4 (100 steps, reference)', traj_rk4_fine),
]
for ax, (name, traj) in zip(axes, methods):
traj_cpu = traj.cpu().numpy()
for i in range(5):
color = 'blue' if y_viz[i] == 0 else 'red'
ax.plot(traj_cpu[:, i, 0], traj_cpu[:, i, 1],
color=color, linewidth=1.5, alpha=0.7)
ax.scatter(traj_cpu[0, i, 0], traj_cpu[0, i, 1],
c=color, s=50, marker='o', zorder=5)
ax.scatter(traj_cpu[-1, i, 0], traj_cpu[-1, i, 1],
c=color, s=50, marker='x', zorder=5)
ax.set_title(name, fontsize=14)
ax.set_xlabel('$h_1$', fontsize=12)
ax.set_ylabel('$h_2$', fontsize=12)
ax.grid(True, alpha=0.3)
ax.set_aspect('equal')
plt.suptitle('Comparison of ODE Solvers', fontsize=15)
plt.tight_layout()
plt.show()
# 終点の誤差比較
ref = traj_rk4_fine[-1].cpu().numpy()
euler_error = np.mean(np.linalg.norm(traj_euler[-1].cpu().numpy() - ref, axis=1))
rk4_error = np.mean(np.linalg.norm(traj_rk4_coarse[-1].cpu().numpy() - ref, axis=1))
print(f'Euler (10 steps) endpoint error: {euler_error:.6f}')
print(f'RK4 (10 steps) endpoint error: {rk4_error:.6f}')
まとめ
本記事では、Neural ODEの理論を数式の導出とともに解説し、Pythonでスパイラルデータの分類を実装しました。
- ResNetの残差接続 $\bm{h}_{t+1} = \bm{h}_t + f(\bm{h}_t, \bm{\theta})$ はオイラー法の離散化であり、連続極限で $\frac{d\bm{h}}{dt} = f(\bm{h}(t), t, \bm{\theta})$ というODEになる
- 順伝播はODEソルバー（RK4法など）で $\bm{h}(0)$ から $\bm{h}(T)$ を計算する
- 随伴法により、中間状態を保持せず $O(1)$ メモリで勾配計算が可能。随伴状態のODEは $\frac{d\bm{a}}{dt} = -\bm{a}^\top \frac{\partial f}{\partial \bm{h}}$ で、時間逆方向に解く
- 連続正規化フローでは、ヤコビアンの行列式の代わりにトレースのみで対数確率密度の変化を追跡でき、$O(d)$ から $O(1)$（ハッチンソン推定量）に計算量を削減できる
- スパイラルデータの実装では、Neural ODEがデータを連続的に変形しながら分類可能な表現に変換する様子を確認した
次のステップとして、以下の記事も参考にしてください。