ガウス過程を用いた常微分方程式のベイズ学習 – Variational Multiple Shooting法

常微分方程式(ODE: Ordinary Differential Equations)は、非確率的な(決定論的)なシステムを記述するのに非常に強力な手法です。しかし、実際の物理現象や生物学的システムでは、ダイナミクスの関数形が未知であることが多く、観測データからダイナミクスを推定する必要があります。

本記事では、ガウス過程(GP)を用いて未知のODEを学習するベイズ的手法について解説します。

本記事の内容

  • ODE学習問題の定式化
  • ガウス過程を用いたODEのベイジアンモデリング
  • 補助変数法(Inducing Variables)による近似
  • Multiple Shooting法による変分推論

問題設定

以下のような微分方程式を考えます。

$$ \begin{align} \bm{x}(t) &= \bm{x}_0 + \int_{0}^{t} \bm{f}(\bm{x}(\gamma)) \, d\gamma \\ \dot{\bm{x}}(t) &= \frac{d\bm{x}(t)}{dt} = \bm{f}(\bm{x}(t)) \end{align} $$

ここで、$\bm{x}(t) \in \mathbb{R}^D$ は時間発展する潜在変数であり、$\bm{f}: \mathbb{R}^D \to \mathbb{R}^D$ は未知の関数です。

私たちの目的は、観測値から微分関数 $\bm{f}$ を推定することです。

観測値 $\bm{Y} = (\bm{y}_1, \bm{y}_2, \dots, \bm{y}_N)^T \in \mathbb{R}^{N \times D}$ が与えられたとき、各要素 $\bm{y}_i \in \mathbb{R}^D$ は、時刻 $t_i$ における真の状態 $\bm{x}(t_i)$ にノイズが加わった観測値であるとします。

$$ \bm{y}_i = \bm{x}(t_i) + \bm{\varepsilon}_i, \quad \bm{\varepsilon}_i \sim \mathcal{N}(\bm{0}, \sigma^2 \bm{I}) $$

ガウス過程を用いたODEのベイジアンモデリング

未知のダイナミクス $\bm{f}$ に対して、ガウス過程の事前分布を置きます。

$$ \bm{f}(\bm{x}) \sim \mathcal{GP}(\bm{0}, K(\bm{x}, \bm{x}’)) $$

この分布は、カーネル関数 $K(\bm{x}, \bm{x}’)$ によって $\bm{f}$ の滑らかさや特性が決まります。代表的なカーネルとしてRBF(Radial Basis Function)カーネルがあります。

$$ K(\bm{x}, \bm{x}’) = \sigma_f^2 \exp\left( -\frac{\|\bm{x} – \bm{x}’\|^2}{2l^2} \right) $$

補助変数法(Inducing Variables Method)

計算量を削減するために、補助変数法を導入します。行列 $\bm{U} \in \mathbb{R}^{M \times D}$ と $\bm{Z} \in \mathbb{R}^{M \times D}$ を導入し、以下の関係を持たせます。

$$ \bm{u}_m = \bm{f}(\bm{z}_m) $$

この拡張により、以下の分布が導かれます。

$$ p(\bm{U}) = \mathcal{N}(\bm{U} \mid \bm{0}, \bm{K}_{zz}) $$

$$ p(\bm{f} \mid \bm{U}; \bm{Z}) = \mathcal{N}(\bm{f} \mid \bm{A} \, \text{vec}(\bm{U}), \bm{K}_{xx} – \bm{A}\bm{K}_{zz}\bm{A}^T) $$

ここで $\bm{A} = \bm{K}_{xz}\bm{K}_{zz}^{-1}$ です。

Multiple Shooting法

ODEの数値積分では、初期値の誤差が時間とともに蓄積するため、長い時系列では勾配消失や発散の問題が発生します。

Multiple Shooting法は、時系列をいくつかのセグメントに分割し、各セグメントの初期値を独立に最適化することで、この問題を軽減します。

時系列を $S$ 個のセグメントに分割し、各セグメント $s$ の初期状態 $\bm{s}_s$ を導入します。

$$ \bm{x}_{s}(t) = \bm{s}_s + \int_{t_s}^{t} \bm{f}(\bm{x}_s(\gamma)) \, d\gamma, \quad t \in [t_s, t_{s+1}] $$

各セグメントの終了状態と次のセグメントの初期状態の間の連続性は、ペナルティ項として導入します。

$$ \mathcal{L}_{\text{continuity}} = \sum_{s=1}^{S-1} \|\bm{x}_s(t_{s+1}) – \bm{s}_{s+1}\|^2 $$

Pythonでの概念実装

ガウス過程を用いたODE学習の概念を、簡単な1次元の例で実装してみましょう。

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from scipy.spatial.distance import cdist

# 真のダイナミクス: dx/dt = -0.5 * sin(x)
def true_dynamics(t, x):
    return -0.5 * np.sin(x)

# 観測データの生成
np.random.seed(42)
x0 = 3.0
t_span = (0, 10)
t_eval = np.linspace(0, 10, 50)

sol = solve_ivp(true_dynamics, t_span, [x0], t_eval=t_eval, method='RK45')
x_true = sol.y[0]

# ノイズ付き観測
sigma_obs = 0.15
y_obs = x_true + np.random.normal(0, sigma_obs, len(x_true))

# ガウス過程回帰でダイナミクスを学習
# 有限差分でdx/dtを近似
dt = t_eval[1] - t_eval[0]
dxdt_approx = np.diff(y_obs) / dt
x_mid = (y_obs[:-1] + y_obs[1:]) / 2  # 中間点

# RBFカーネル
def rbf_kernel(X1, X2, length_scale=1.0, variance=1.0):
    dist = cdist(X1.reshape(-1, 1), X2.reshape(-1, 1), 'sqeuclidean')
    return variance * np.exp(-dist / (2 * length_scale**2))

# GP回帰
l, sf = 1.0, 1.0
sn = 0.5
K = rbf_kernel(x_mid, x_mid, l, sf) + sn**2 * np.eye(len(x_mid))
K_inv = np.linalg.inv(K)

# 予測点
x_pred = np.linspace(-1, 4, 200)
K_star = rbf_kernel(x_pred, x_mid, l, sf)
K_ss = rbf_kernel(x_pred, x_pred, l, sf)

# GPの予測平均と分散
f_mean = K_star @ K_inv @ dxdt_approx
f_var = np.diag(K_ss - K_star @ K_inv @ K_star.T)
f_std = np.sqrt(np.maximum(f_var, 0))

# 真のダイナミクス
f_true = -0.5 * np.sin(x_pred)

# 可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# 左: 観測データと真の軌道
axes[0].plot(t_eval, x_true, 'g-', linewidth=2, label='True trajectory')
axes[0].scatter(t_eval, y_obs, c='steelblue', s=20, alpha=0.7, label='Observations')
axes[0].set_title("ODE Trajectory")
axes[0].set_xlabel("t")
axes[0].set_ylabel("x(t)")
axes[0].legend()
axes[0].grid(True, alpha=0.3)

# 右: 学習されたダイナミクス
axes[1].plot(x_pred, f_true, 'g-', linewidth=2, label='True f(x) = -0.5*sin(x)')
axes[1].plot(x_pred, f_mean, 'r-', linewidth=2, label='GP posterior mean')
axes[1].fill_between(x_pred, f_mean - 2*f_std, f_mean + 2*f_std,
                     alpha=0.2, color='red', label='2-sigma')
axes[1].scatter(x_mid, dxdt_approx, c='steelblue', s=20, alpha=0.5, label='Finite diff. approx.')
axes[1].set_title("Learned Dynamics f(x)")
axes[1].set_xlabel("x")
axes[1].set_ylabel("dx/dt")
axes[1].legend()
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

右のグラフで、GPの事後平均(赤線)が真のダイナミクス(緑線)をよく近似していることが確認できます。不確実性(赤い帯)はデータが少ない領域で大きくなっています。

Multiple Shootingの効果の可視化

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp

def true_dynamics(t, x):
    return -0.5 * np.sin(x)

np.random.seed(42)
x0 = 3.0
t_span = (0, 10)
t_eval = np.linspace(0, 10, 100)

sol = solve_ivp(true_dynamics, t_span, [x0], t_eval=t_eval)
x_true = sol.y[0]

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# 左: Single Shooting(初期値からの一括積分)
axes[0].plot(t_eval, x_true, 'g-', linewidth=2, label='True')
# 初期値に誤差がある場合
x0_noisy = x0 + 0.5
sol_noisy = solve_ivp(true_dynamics, t_span, [x0_noisy], t_eval=t_eval)
axes[0].plot(t_eval, sol_noisy.y[0], 'r--', linewidth=2, label='Single shooting (noisy IC)')
axes[0].set_title("Single Shooting: error propagates")
axes[0].set_xlabel("t")
axes[0].set_ylabel("x(t)")
axes[0].legend()
axes[0].grid(True, alpha=0.3)

# 右: Multiple Shooting(セグメント分割)
n_segments = 5
segment_length = len(t_eval) // n_segments
colors = plt.cm.Set1(np.linspace(0, 1, n_segments))

axes[1].plot(t_eval, x_true, 'g-', linewidth=2, alpha=0.5, label='True')

for s in range(n_segments):
    start = s * segment_length
    end = min((s + 1) * segment_length + 1, len(t_eval))
    t_seg = t_eval[start:end]
    # 各セグメントの初期値を真の値に近い値にする
    x0_seg = x_true[start] + np.random.normal(0, 0.1)
    sol_seg = solve_ivp(true_dynamics, (t_seg[0], t_seg[-1]), [x0_seg], t_eval=t_seg)
    axes[1].plot(t_seg, sol_seg.y[0], '-', color=colors[s], linewidth=2)
    axes[1].axvline(x=t_seg[0], color='gray', linestyle=':', alpha=0.3)

axes[1].set_title("Multiple Shooting: segments are independent")
axes[1].set_xlabel("t")
axes[1].set_ylabel("x(t)")
axes[1].legend()
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

Single Shooting(左)では初期値の誤差が時間とともに累積しますが、Multiple Shooting(右)では各セグメントが独立に積分されるため、誤差の伝播が抑制されます。

まとめ

本記事では、ガウス過程を用いた常微分方程式のベイズ学習手法について解説しました。

  • GP-ODE: 未知のODEダイナミクス $\bm{f}$ にガウス過程の事前分布を置いて学習する手法
  • 補助変数法: 計算量を $O(N^3)$ から $O(NM^2)$ に削減する近似手法
  • Multiple Shooting法: 時系列をセグメントに分割し、長期の勾配消失問題を回避する
  • GP-ODEはダイナミクスの不確実性を自然に推定できるベイズ的な手法
  • モーション解析や物理シミュレーション、生物学的モデリングなど幅広い応用がある