大気汚染や気候のシミュレーションでは、数十〜数千の化学種が数千の反応でつながった「化学反応メカニズム」を常微分方程式（ODE）として解きます。このメカニズムは長年、Kinetic PreProcessor（KPP）という専用言語で書かれ、KPP が Fortran や C のコードを生成して計算してきました。ただし、生成されたコードは PyTorch や JAX のような「微分可能プログラミング」の世界とはつながっていないため、ニューラルネットワーク（NN）を化学モデルの中に入れて一緒に学習させることが難しいという問題がありました。
2026年10月に arXiv に投稿された論文「KPPax: preprocess differentiable kinetic ODE for machine learning in atmospheric chemistry」は、KPP 言語で書かれたメカニズムから JAX のコードを生成し、硬い（stiff）化学 ODE の解を丸ごと自動微分できるようにする ツールを提案しています。
この記事では、論文の数式を追いながら、次の 3 点を解説します。
- 速度則を「対数空間」で書き直すと、なぜ自動微分できるようになるのか
- 自動微分で得た感度から、フィッシャー情報行列（FIM）を使って「どの反応速度定数がデータから決まるか」をフィッティング前に予測する方法
- NN を化学 ODE の中に入れて学習する「オンライン学習」が、従来の 2 段階の方法より何を取り戻せるのか
あわせて、論文のアイデアを数十行で再現したおもちゃのコードを JAX と diffrax で動かした結果も載せます。
対象読者：機械学習や自動微分の基本（勾配、ヤコビアン、最小二乗）を知っていて、科学計算の ODE と機械学習を組み合わせる話に興味がある方。大気化学の予備知識は不要です。
本論文は 2026年10月27日に arXiv に公開されたプレプリントで、査読前のものです。記事中の数値は論文に書かれている値で、筆者が論文の実験を再現したものではありません。
論文の概要
| 項目 | 内容 |
|---|---|
| タイトル | KPPax: preprocess differentiable kinetic ODE for machine learning in atmospheric chemistry |
| 著者 | Wenqing Peng, Michael Boy, Zhi-Song Liu（University of Helsinki, LUT University） |
| 公開 | arXiv:2610.37655（2026年10月27日、cs.AI） |
| 何をしたか | KPP 言語のメカニズムから、微分可能・JIT コンパイル可能・ベクトル化可能な JAX の化学コードを生成する |
| 検証 | POLLU ベンチマーク（20 化学種・25 反応）で、(1) 反応速度定数の識別可能性解析と推定、(2) NN による排出インベントリ補正の学習 |
主な結果は次の 2 つです。
- 自動微分で求めたヤコビアンから FIM を計算すると、25 個の反応速度定数のうちどれがデータから決まるかを、フィッティングの前に 予測できました。識別可能と予測された 9 個は、ノイズ入りの合成データからの推定誤差が予測と約 30% 以内で一致しました。
- 排出量の補正を NN で学習する問題で、化学モデルを通した「オンライン」のファインチューニングを加えると、植物由来（biogenic）の排出補正の誤差が 2.4〜5 分の 1 になりました。
化学反応を ODE として書く
まず、化学反応系がどう ODE になるかを整理します。化学種が $n$ 個、反応が $m$ 個あるとし、化学種 $i$ の濃度を $y_i$、濃度ベクトルを $\mathbf{y} = [y_1, \dots, y_n]^\top$ とします。
反応 $j$ で消費される化学種 $i$ の分子数を $s^-_{i,j}$、生成される分子数を $s^+_{i,j}$（化学量論係数）とすると、反応 $j$ は次のように書けます。
\sum_i s^-_{i,j}\, y_i \;\longrightarrow\; \sum_i s^+_{i,j}\, y_i
反応 $j$ の速さ $\omega_j$ は、質量作用の法則により、反応物の濃度を化学量論係数でべき乗して掛け合わせたものに比例します。比例定数 $k_j$ が反応速度定数です。
\omega_j(t, \mathbf{y}) = k_j(t) \prod_{i=1}^{n} y_i^{\,s^-_{i,j}}
$k_j$ は定数のこともあれば、光の強さで変わる光分解反応のように時間に依存することもあります。
各化学種の濃度は、すべての反応での生成と消費の合計で変化します。係数を行列 $S^-$、$S^+$ にまとめ、$S = S^+ - S^-$、$\boldsymbol{\omega} = [\omega_1, \dots, \omega_m]^\top$ とおくと、
\frac{d\mathbf{y}}{dt} = (S^+ - S^-)\,\boldsymbol{\omega}(t, \mathbf{y}) = S\,\boldsymbol{\omega}(t, \mathbf{y}) = \mathbf{f}(t, \mathbf{y})
という ODE になります。大気化学ではラジカルがミリ秒〜秒で反応する一方、長寿命の物質は数時間かけて変化するため、ヤコビアンの固有値が何桁にもわたる 硬い（stiff） 系になります。陽解法では安定性のために刻み幅を極端に小さくする必要があるので、陰解法が使われます。
速度則を対数空間で書き直す
疎な行列演算で速度則を計算する
KPPax は、メカニズムの読み込みに Python 製の pykpp のパーサーを使い、化学量論係数を COO 形式（行インデックス、列インデックス、値の 3 つの配列）の疎配列として保持します。GPU で速く計算するため、$S^-$ や $S$ を密行列にせず、インデックスで値を集める（gather）操作と、セグメントごとに集約する（segment reduction）操作だけで速度則を計算します。
素直に書くと、論文の Algorithm 1 のように「反応物の濃度のべき乗を反応ごとに掛け合わせる」処理になります。
- $p \leftarrow y[i_r]^{v_r}$：各反応物の濃度を集めて化学量論係数でべき乗する
- $\omega \leftarrow k \odot \mathrm{segment\_prod}(p, j_r)$：反応ごとに掛け合わせる
- $\dot{y} \leftarrow \mathrm{segment\_sum}(v_s \odot \omega[j_s], i_s)$：化学種ごとに足し合わせる
ここで $(i_r, j_r, v_r)$ は $S^-$ の、$(i_s, j_s, v_s)$ は $S$ の COO 表現です。
掛け算を足し算に変える
問題は 2 の segment_prod です。同じ反応インデックスに複数の値を掛け合わせる scatter-multiply には、JAX の微分規則が実装されていません。そこで論文は、速度則の対数をとります。
\log \omega_j = \log k_j + \sum_i s^-_{i,j} \log y_i
積が和になったので、segment_prod を segment_sum に置き換えられます（Algorithm 2）。
- $q \leftarrow v_r \odot \log y[i_r]$
- $l \leftarrow \mathrm{segment\_sum}(q, j_r)$（$= \log \prod_i y_i^{s^-_{i,j}}$）
- $\omega \leftarrow k \odot \exp(l)$
- $\dot{y} \leftarrow \mathrm{segment\_sum}(v_s \odot \omega[j_s], i_s)$
segment_sum には標準の微分規則があるので、計算コストと疎なインデックスのパターンはそのままで、$\mathbf{y}$ と $\mathbf{k}$ の両方について微分できるようになります。代わりに、$\log y_i$ を計算するために 濃度が常に正である 必要があります。論文では右辺を評価するたびに $\mathbf{y}$ を小さな正の値で下から切り（clip）、ソルバーが 0 や負の濃度を試さないようにしています。
手元の環境（JAX 0.12.0.1）で、この違いを確かめました。
import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)
y = jnp.array([0.05, 0.02, 0.06])
ir = jnp.array([0, 1, 2]) # 反応物の化学種
jr = jnp.array([0, 1, 1]) # どの反応か（反応 1 は 2 分子反応）
def omega_prod(y): # Algorithm 1: scatter-multiply
return jax.ops.segment_prod(y[ir], jr, num_segments=2).sum()
def omega_log(y): # Algorithm 2: log 空間で scatter-add
return jnp.exp(jax.ops.segment_sum(jnp.log(y[ir]), jr, num_segments=2)).sum()
print("log 空間:", jax.grad(omega_log)(y))
try:
print("segment_prod:", jax.grad(omega_prod)(y))
except Exception as e:
print("segment_prod:", type(e).__name__, str(e).splitlines()[0][:120])
log 空間: [1. 0.06 0.02]
segment_prod: NotImplementedError scatter_mul gradients are only implemented if `unique_indices=True`
$\omega = y_0 + y_1 y_2$ の勾配 $(1, y_2, y_1) = (1, 0.06, 0.02)$ が、対数空間の版では正しく得られています。
硬い ODE を陰解法で解く
時間積分には JAX ネイティブの微分可能な ODE ソルバーライブラリ diffrax を使い、陰的な Kvaerno 系のソルバー（対角陰的ルンゲ＝クッタ法）と、局所誤差に応じて刻み幅を変える PID コントローラーを組み合わせます。計算はすべて 64 ビット倍精度です。
ODE ソルバーを通して微分する
2 つの随伴法
ODE の解を微分するには、大きく 2 つの方法があります。
| 方法 | 考え方 | 得られるもの |
|---|---|---|
| optimize-then-discretize | 連続の方程式から随伴 ODE を導き、逆向きに別途積分する | 連続問題の勾配の近似 |
| discretize-then-optimize | 前向きの数値スキームそのものをステップごとに微分する | 数値解の厳密な勾配 |
KPP も両方を実装していますが、discretize-then-optimize の随伴は積分法ごとに手で導く必要があり、ヤコビアンの濃度微分（ヘッセテンソル）まで要るため、一部の積分法にしか用意されていません。
diffrax では、ソルバーの内部ステップにリバースモードの自動微分を直接かけることで、discretize-then-optimize の随伴が自動で得られます。すべてのステップの状態を保存するとメモリがステップ数に比例してしまうため、状態の一部だけを保存して残りを逆伝播時に再計算する 再帰的チェックポイント を使い、メモリをステップ数に対して劣線形に抑えます。論文はこれを「計算コストと引き換えに、積分法ごとの随伴を手で書く手間と誤りのリスクをなくすもの」と位置づけています。
フォワードモードとリバースモードの使い分け
論文の考察では、問題の形によって自動微分のモードを使い分けています。
- パラメータ数が観測数より少ない（反応速度定数 25 個、4D-Var の制御変数 50 個など）：フォワードモード。パラメータ 1 つにつき接線を 1 本積分するだけでヤコビアン全体が得られ、チェックポイントも不要
- 多数のパラメータについてスカラー損失を微分する（NN の 1 次の学習など）：リバースモード（再帰的チェックポイント付き）
diffrax では、この選択が生成コードではなく diffeqsolve の adjoint 引数で切り替えられます。
応用 1：どの反応速度定数がデータで決まるか
フィッシャー情報行列で識別可能性を調べる
反応速度定数を推定するとき、異なるパラメータの組み合わせが同じくらいデータに合ってしまい、値が一意に決まらないことがあります（非識別性）。論文は、推定の前にこれを調べます。
反応速度定数を $\mathbf{k} = \mathbf{k}_{\mathrm{ref}} \odot \exp(\boldsymbol{\theta})$ とパラメータ化します。$\mathbf{k}_{\mathrm{ref}}$ は基準値、$\boldsymbol{\theta}$ は対数の補正係数です。観測演算子 $h(\boldsymbol{\theta})$ は、観測したすべての化学種・時刻・実験の対数濃度を並べたベクトルで、観測値 $\tilde{\mathbf{y}}$ に対数空間で標準偏差 $\sigma$ のガウス誤差があるとき、残差とコストは
\mathbf{r}(\boldsymbol{\theta}) = \frac{h(\boldsymbol{\theta}) - \tilde{\mathbf{y}}}{\sigma}, \qquad
\chi^2(\boldsymbol{\theta}) = \mathbf{r}(\boldsymbol{\theta})^\top \mathbf{r}(\boldsymbol{\theta})
です。$h$ も $\boldsymbol{\theta}$ も対数なので、ヤコビアン $J = \partial h / \partial \boldsymbol{\theta}$ の成分は相対感度 $\partial \ln y / \partial \ln k$ になり、桁の違う化学種や反応どうしでも比べられます。
真のパラメータでのフィッシャー情報行列（FIM）とその固有値分解は
F = \frac{1}{\sigma^2} J^\top J = V \Lambda V^\top
です。固有ベクトル $\mathbf{v}_m$（$V$ の列）はパラメータ空間の方向、つまり「反応速度定数の組み合わせ」を表し、固有値 $\lambda_m$ はその組み合わせがデータでどれだけ強く制約されるかを表します。クラメール・ラオの下限から、個々のパラメータ $\theta_p$ の標準偏差の下限は
\sigma^{\mathrm{CR}}_p = \left[ (F^{-1})_{pp} \right]^{1/2} = \left( \sum_m \frac{V_{pm}^2}{\lambda_m} \right)^{1/2}
となります。論文は $\sigma^{\mathrm{CR}}_p < 0.1$（1σ の不確かさが約 10%）のものを「識別可能」としています。固有値が 0 に近い方向は、データで区別できない反応の組み合わせで、その成分 $V_{pm}$ を見るとどの反応が混ざっているかがわかります。
この解析に必要なのは 1 点でのヤコビアンだけ で、フィッティングの前に実行できます。
論文の結果
POLLU の 25 個の反応速度定数を、フローチューブ実験を模した設計（初期濃度を変えた 8 条件、60 分後に全 20 化学種を 1 回観測、計 160 観測、$\sigma = 0.05$）で調べています。
- 自動微分のヤコビアン（160×25）は中心差分と相対誤差 $8.7 \times 10^{-7}$ で一致し、計算時間は 0.77 秒（中心差分は順計算 50 回で 28.7 秒）
- FIM の固有値は 16 桁にわたり、$\sigma^{\mathrm{CR}}_p < 0.1$ の反応速度定数は 9 個
- 識別できない反応は、速い反応の 3 つのクラスター（NO–NO₂–O₃ の光定常サイクル、NO₃–NO₂–N₂O₅ の平衡、O₃–O(¹D)–O(³P) の相互変換）にまとまった。60 分の反応時間に比べて速すぎるため、観測濃度は 反応速度定数の比 だけで決まり、絶対値は決まらない
- ノイズ入り合成データからの推定（初期値をランダムに変えた 20 試行、レーベンバーグ・マーカート法）では、識別可能な 9 個の推定誤差が $\sigma^{\mathrm{CR}}_p$ と約 30% 以内で一致。識別不能なものは対数で最大 15（6 桁以上）ずれたが、クラスター内の比は保たれていた
一方で、FIM は「どの組み合わせが決まらないか」は正しく当てたものの、「どこまでずれるか」は過小評価しました。コスト関数の谷が曲がっていて、線形近似が成り立たないためです。
おもちゃのモデルで再現してみる
論文の考え方を確かめるため、6 化学種・6 反応の小さなメカニズムを自分で作り、同じ手順を JAX と diffrax で書きました。速い光定常サイクル（R1〜R3）と、速い平衡（R5 と R6）をわざと入れています。
| 反応 | 式 | $k_{\mathrm{ref}}$ |
|---|---|---|
| R1 | NO₂ → NO + O | 0.5 /min |
| R2 | O → O₃（O₂ は一定として省略） | 1×10⁴ /min |
| R3 | NO + O₃ → NO₂ | 20 /(ppm·min) |
| R4 | NO₂ + O₃ → NO₃ | 0.05 /(ppm·min) |
| R5 | NO₃ + NO₂ → N₂O₅ | 1×10⁴ /(ppm·min) |
| R6 | N₂O₅ → NO₃ + NO₂ | 100 /min |
import time
import jax
import jax.numpy as jnp
import diffrax
jax.config.update("jax_enable_x64", True)
# 化学種: 0:NO2 1:NO 2:O 3:O3 4:NO3 5:N2O5
SPECIES = ["NO2", "NO", "O", "O3", "NO3", "N2O5"]
N_SPEC = len(SPECIES)
# 反応: (反応物, 生成物)
REACTIONS = [
({0: 1}, {1: 1, 2: 1}), # R1: NO2 -> NO + O
({2: 1}, {3: 1}), # R2: O -> O3 (O2 は一定として省略)
({1: 1, 3: 1}, {0: 1}), # R3: NO + O3 -> NO2
({0: 1, 3: 1}, {4: 1}), # R4: NO2 + O3 -> NO3
({4: 1, 0: 1}, {5: 1}), # R5: NO3 + NO2 -> N2O5
({5: 1}, {4: 1, 0: 1}), # R6: N2O5 -> NO3 + NO2
]
N_REAC = len(REACTIONS)
K_REF = jnp.array([0.5, 1.0e4, 20.0, 0.05, 1.0e4, 1.0e2]) # 1/min, 1/(ppm min)
# S^- と S = S^+ - S^- を COO 形式 (行, 列, 値) で持つ
ir, jr, vr, is_, js, vs = [], [], [], [], [], []
for j, (reac, prod) in enumerate(REACTIONS):
for i, s in reac.items():
ir.append(i); jr.append(j); vr.append(s)
for i in set(reac) | set(prod):
net = prod.get(i, 0) - reac.get(i, 0)
if net:
is_.append(i); js.append(j); vs.append(net)
ir, jr, is_, js = map(jnp.array, (ir, jr, is_, js))
vr, vs = jnp.array(vr, float), jnp.array(vs, float)
def rhs(t, y, k):
"""対数空間での疎な速度則（論文の Algorithm 2 に相当）"""
y = jnp.maximum(y, 1e-30) # log のため正に保つ
q = vr * jnp.log(y[ir])
l = jax.ops.segment_sum(q, jr, num_segments=N_REAC)
omega = k * jnp.exp(l)
return jax.ops.segment_sum(vs * omega[js], is_, num_segments=N_SPEC)
def simulate(k, y0, t1=60.0):
sol = diffrax.diffeqsolve(
diffrax.ODETerm(rhs), diffrax.Kvaerno5(), t0=0.0, t1=t1, dt0=1e-4, y0=y0,
args=k, saveat=diffrax.SaveAt(t1=True),
stepsize_controller=diffrax.PIDController(rtol=1e-8, atol=1e-12),
max_steps=100_000, adjoint=diffrax.ForwardMode(),
)
return sol.ys[-1]
# 初期条件を 8 通り変えた「フローチューブ風」の実験
key = jax.random.PRNGKey(0)
base = jnp.array([0.05, 0.02, 0.0, 0.06, 0.0, 0.0])
factors = jnp.exp(jnp.log(2.0) * jax.random.normal(key, (8, N_SPEC)))
Y0 = jnp.where(base > 0, base * factors, 1e-12)
def h(theta):
"""観測演算子: k = k_ref * exp(theta) で解き、60 分後の全化学種の log 濃度を並べる"""
k = K_REF * jnp.exp(theta)
yT = jax.vmap(simulate, in_axes=(None, 0))(k, Y0)
return jnp.log(jnp.maximum(yT, 1e-30)).ravel()
theta0 = jnp.zeros(N_REAC)
jac_fn = jax.jit(jax.jacfwd(h))
jac_fn(theta0).block_until_ready() # JIT コンパイル
t = time.perf_counter()
J = jac_fn(theta0).block_until_ready()
t_ad = time.perf_counter() - t
h_jit = jax.jit(h)
h_jit(theta0).block_until_ready()
t = time.perf_counter()
eps = 1e-5
J_fd = jnp.stack([(h_jit(theta0.at[p].add(eps)) - h_jit(theta0.at[p].add(-eps))) / (2 * eps)
for p in range(N_REAC)], axis=1)
t_fd = time.perf_counter() - t
print(f"J shape = {J.shape}")
print(f"AD vs 中心差分の相対誤差 = {jnp.linalg.norm(J - J_fd) / jnp.linalg.norm(J):.1e}")
print(f"時間: AD {t_ad:.2f} s, 中心差分 {t_fd:.2f} s")
sigma = 0.05
F = J.T @ J / sigma**2
lam, V = jnp.linalg.eigh(F)
print("\nFIM の固有値（小さい順）と固有ベクトル:")
for m in range(N_REAC):
vec = " ".join(f"{v:+.2f}" for v in V[:, m])
print(f" lambda = {lam[m]:9.2e} v = [{vec}]")
sd = jnp.sqrt(jnp.diag(jnp.linalg.pinv(F)))
print("\nCramér–Rao 下限の標準偏差（ln k の単位）:")
for p in range(N_REAC):
tag = "識別可能" if sd[p] < 0.1 else "識別不能"
print(f" R{p + 1}: {sd[p]:.3g} {tag}")
実行結果です（Apple Silicon の CPU、時間は環境によって変わります）。
J shape = (48, 6)
AD vs 中心差分の相対誤差 = 1.0e-07
時間: AD 0.13 s, 中心差分 0.79 s
FIM の固有値（小さい順）と固有ベクトル:
lambda = 9.83e-08 v = [-0.00 -0.00 -0.00 +0.00 +0.71 +0.71]
lambda = 3.09e-03 v = [+0.58 +0.58 +0.58 -0.00 +0.00 +0.00]
lambda = 1.85e+03 v = [+0.34 +0.10 -0.44 +0.36 +0.52 -0.52]
lambda = 2.61e+03 v = [-0.23 -0.41 +0.64 +0.53 +0.20 -0.20]
lambda = 5.63e+03 v = [+0.67 -0.70 +0.03 -0.26 -0.05 +0.05]
lambda = 8.01e+03 v = [-0.23 -0.01 +0.24 -0.72 +0.43 -0.43]
Cramér–Rao 下限の標準偏差（ln k の単位）:
R1: 10.4 識別不能
R2: 10.4 識別不能
R3: 10.4 識別不能
R4: 0.0162 識別可能
R5: 2.25e+03 識別不能
R6: 2.26e+03 識別不能
小さいモデルでも、論文と同じ構造がはっきり出ています。
- 最小の固有値の固有ベクトル $(0, 0, 0, 0, 0.71, 0.71)$ は「R5 と R6 を同じ割合で大きくする」方向です。平衡定数 $k_5/k_6$ が変わらないので、観測される濃度はほとんど変わりません
- 2 番目の固有ベクトル $(0.58, 0.58, 0.58, 0, 0, 0)$ は、光定常サイクル R1〜R3 を同時に速くする方向です
- 遅い R4 だけが、観測から絶対値まで決まります
逆に言うと、たとえば $\theta_1 - \theta_3$（R1 と R3 の比）のように、これらの方向と直交する組み合わせは決まります。FIM を見れば、比で表し直してから推定すべきパラメータがフィッティングの前にわかる、というのが論文の主張です。
応用 2：NN を化学 ODE の中に入れて排出量を学習する
問題設定
2 つ目の応用は、排出インベントリ（どこからどれだけ汚染物質が出ているかの推定値）の補正です。POLLU を 0 次元のボックスモデルにし、排出と、背景濃度 $\mathbf{c}_{\mathrm{bg}}$ への希釈を加えます。
\frac{d\mathbf{y}}{dt} = S\,\boldsymbol{\omega}(t, \mathbf{y}) + \mathbf{E}(t) - \frac{\mathbf{y} - \mathbf{c}_{\mathrm{bg}}}{\tau}
$\tau = 12$ 時間は希釈の時定数です。排出は人為起源（anth）と植物起源（bio）の 2 部門 $s$ からなり、
\mathbf{E}(t) = \sum_s \mathbf{e}_s\, g_s\big(\mathbf{x}(t)\big)\, p_s(t)\, \exp\Big(f_s\big(\mathbf{x}(t)\big)\Big)
と書きます。$\mathbf{e}_s$ は化学種ごとの基準排出量、$p_s(t)$ は既知の日変化、$g_s$ はインベントリ側の大まかなスケーリング、$\mathbf{x} = (T, R, \log_{10} P, d_{\mathrm{anth}}, d_{\mathrm{bio}})$ は気温・日射・人口密度などの説明変数です。$f_s$ が学習したい 対数補正 で、インベントリどおりなら $f_s = 0$ です。真の補正 $f^{\mathrm{true}}_s$ は、線形モデルや加法モデルでは表せない滑らかな非線形関数として人工的に作られています。
$f_s$ は部門ごとの MLP（隠れ層 2 層・幅 32・softplus、出力を $f_{\max}\tanh(\cdot / f_{\max})$、$f_{\max} = 3$ で有界化）で表し、パラメータは 2 つの MLP で計 2434 個です。
オフラインの 2 段階学習
従来の方法では、まずシナリオ $n$ ごとに 4D-Var（4 次元変分法）で補正を推定します。制御変数 $\boldsymbol{\delta}_n$ は 1 時間ごとの節点での補正値で、
J_n(\boldsymbol{\delta}_n) = \frac{1}{2}\sum_{t,i} r_{t,i}^2 + \frac{1}{2}\sum_s \boldsymbol{\delta}_{n,s}^\top B^{-1} \boldsymbol{\delta}_{n,s},
\qquad
B_{jl} = \sigma_b^2 \exp\left(-\frac{|t_j - t_l|}{L}\right)
を最小化します。第 2 項は事前分布（インベントリ $\boldsymbol{\delta} = 0$）に近づける背景項で、$\sigma_b$ は事前の標準偏差、$L$ は時間方向の相関長です。残差 $r_{t,i}$ は、検出下限 $y_{\mathrm{DL}}$ を入れた対数濃度の差を $\sigma_{\mathrm{obs}}$ で割ったものです。
r_{t,i} = \frac{\ln \max(y_{t,i},\, y_{\mathrm{DL}}) - \ln \max(y^{\mathrm{obs}}_{t,i},\, y_{\mathrm{DL}})}{\sigma_{\mathrm{obs}}}
次に、推定した補正 $\hat{\boldsymbol{\delta}}$ を教師データとして NN を回帰で学習します。事後の標準偏差 $\sigma^{\mathrm{post}}$ が大きい（＝観測で決まっていない）点の重みを下げた、重み付き二乗誤差です。
L_{\mathrm{off}}(\boldsymbol{\theta}) = \frac{1}{N}\sum_{n,j,s} w_{n,j,s}\Big[f_s\big(\mathbf{x}_n(t_j); \boldsymbol{\theta}\big) - \hat{\delta}_{n,j,s}\Big]^2,
\qquad
w_{n,j,s} \propto \big(\sigma^{\mathrm{post}}_{n,j,s}\big)^{-2}
この方法の弱点は、各シナリオの推定が事前分布に引き寄せられるため、1 つのシナリオだけでは事前分布から切り分けられない弱い情報が、NN を学習する前に失われる ことです。
オンラインのファインチューニング
KPPax では化学 ODE が微分可能なので、NN を ODE の中に置いたまま、観測された濃度に直接合わせられます。
J(\boldsymbol{\theta}) = \frac{1}{2}\sum_{n \in \mathrm{train}} \sum_{t,i} r_{n,t,i}\big(\boldsymbol{\delta}_n(\boldsymbol{\theta})\big)^2,
\qquad
\boldsymbol{\delta}_n(\boldsymbol{\theta}) = \Big[f_s\big(\mathbf{x}_n(t_j); \boldsymbol{\theta}\big)\Big]_{j,s}
節点の値が、シナリオごとの自由な変数ではなく、全シナリオ共通の NN の出力になっている点がポイントです。共通の NN が「同じ気温帯を通るすべてのシナリオ」から弱い情報を少しずつ集められるので、個々の推定では失われていた情報を取り戻せます。
最適化には、1 次の Adam ではなく、ガウス・ニュートン法（レーベンバーグ・マーカートの減衰付き）を使います。連鎖律により、シナリオ $n$ の残差のパラメータについてのヤコビアンは $J_{\delta,n} D_n$ と分解できます。$J_{\delta,n}$ は 50 個の節点値についての残差のヤコビアン（4D-Var と同じもの）、$D_n = \partial \boldsymbol{\delta}_n / \partial \boldsymbol{\theta}$ は NN のヤコビアンで、こちらは ODE を解かずに計算できます。
H = \sum_n D_n^\top J_{\delta,n}^\top J_{\delta,n} D_n,
\qquad
\mathbf{g} = \sum_n D_n^\top J_{\delta,n}^\top \mathbf{r}_n,
\qquad
\Delta\boldsymbol{\theta} = -\big(H + \lambda\,\mathrm{diag}(H)\big)^{-1}\mathbf{g}
高価な ODE の感度 $J_{\delta,n}$ は 50 個の節点値についてフォワードモードで求めるだけなので、コストは NN のパラメータ数に比例して増えません。NN のパラメータ数が効くのは $2434 \times 2434$ の $H$ を解く部分だけです。15 回の反復が、150 シナリオで単一の NVIDIA GH200 GPU で約 2 分だったと報告されています。
論文によると、観測局データ・ノイズありの条件で Adam を 400 ステップ回しても、植物起源の補正はほとんど変わりませんでした。この方向はパラメータ空間で弱くしか制約されておらず、勾配が小さいためです。ガウス・ニュートン法は逆曲率でこの方向を拡大するので効きます。Adam でゼロからオンライン学習することもできましたが、オフライン学習＋ファインチューニングの 20 倍以上の計算時間がかかり、精度も劣りました。
結果
学習した補正 $f_s$ と真値 $f^{\mathrm{true}}_s$ の RMSE（学習範囲内）を、論文の表 1 から抜き出しました。
| 観測 | ノイズ | 人為起源 オフライン | 人為起源 ＋オンライン | 植物起源 オフライン | 植物起源 ＋オンライン |
|---|---|---|---|---|---|
| 全化学種 | なし | 0.063 | 0.056 | 0.162 | 0.041 |
| 全化学種 | あり | 0.063 | 0.043 | 0.145 | 0.028 |
| 観測局の 5 種 | なし | 0.063 | 0.053 | 0.287 | 0.107 |
| 観測局の 5 種 | あり | 0.055 | 0.036 | 0.309 | 0.128 |
植物起源の改善が特に大きく、全化学種を観測した場合、最も寒い気温帯（5〜12 ℃）の誤差は 0.26〜0.29 から 0.04〜0.06 に下がりました。気温が低い時間帯は植物起源の排出が小さく、1 シナリオの中では信号が弱いため、シナリオごとの推定ではほぼ事前分布のままになっていたものです。
限界と注意点
論文自身が挙げている限界です。
- 大規模メカニズムへの拡張：diffrax の陰解法は各ステップで右辺のヤコビアンを密行列のまま解くため、コストは化学種数の 3 乗で増えます。KPP は記号的な疎 LU 分解を使います。POLLU（20 種）では差は無視できますが、数千種の MCM（Master Chemical Mechanism）では支配的になるため、疎または行列フリーの線形ソルバーが最も重要な拡張とされています。KPP 生成の Fortran コードとの精度・速度の比較は別途報告予定です
- 対数空間の副作用：濃度を正の下限で切るため、切られた濃度についての勾配は 0 になります。初期濃度が 0 の化学種や、使い切られる化学種について微分する場合は注意が必要です
- FIM は局所的：1 点のヤコビアンに基づくので、曲がった谷に沿ったずれの大きさは過小評価します。プロファイル尤度や MCMC と組み合わせられるとしています
- 応用 2 は理想化された実験：真の補正は人工の関数で、化学メカニズムもモデルと真値で同じ、輸送なし、初期状態も既知です。実データでは初期状態の誤差が早朝の排出補正に混ざる可能性があり、初期状態との同時推定が必要とされています
- 2 次法のスケール：$H$ を NN の全パラメータについて作るので、2434 個なら可能でも、大きな NN では行列フリーや近似の 2 次法が必要です
- ノイズなし・全化学種観測の条件では、ファインチューニング中にパラメータが平らな方向へ大きく動いた（オフライン解からの距離が約 $2 \times 10^6$）ため、近接項や早期終了を勧めています
ハマりどころ：diffrax でフォワードモードを使うとき
おもちゃのコードを書いたとき、最初は adjoint を指定せずに jax.jacfwd をかけ、次のエラーになりました。
TypeError: can't apply forward-mode autodiff (jvp) to a custom_vjp function.
diffrax の diffeqsolve は、既定でリバースモード用の随伴（再帰的チェックポイント）を custom_vjp として定義しているため、そのままではフォワードモードで微分できません。adjoint=diffrax.ForwardMode() を渡すと、jax.jacfwd が使えるようになります。パラメータが少なく出力が多い FIM の計算では、論文と同じくフォワードモードが向いています。
まとめ
- KPPax は、KPP 言語で書かれた大気化学メカニズムから、微分可能・JIT 可能・ベクトル化可能な JAX のコードを生成するツールです
- 速度則を対数空間で書き、積を和（
segment_sum）に変えることで、疎な計算のまま自動微分できるようにしています - 硬い ODE の陰解法と随伴を diffrax に任せることで、積分法ごとの随伴を手で導く必要がなくなります
- 自動微分で得たヤコビアンから FIM を作ると、どの反応速度定数（の組み合わせ）がデータで決まるかを、推定の前に予測できます
- NN を化学 ODE の中に入れてガウス・ニュートン法で学習すると、シナリオごとの推定では失われる弱い情報を取り戻せ、植物起源の排出補正の誤差が 2.4〜5 分の 1 になりました
参考
- KPPax: preprocess differentiable kinetic ODE for machine learning in atmospheric chemistry（Peng, Boy, Liu, arXiv:2610.37655, 2026）
- Diffrax ドキュメント（Patrick Kidger）
- Diffrax: Adjoints（Patrick Kidger）
- jax.ops.segment_sum（JAX ドキュメント）
- On Neural Differential Equations（Kidger, arXiv:2202.02435, 2022）
- Optimistix: modular optimisation in JAX and Equinox（Rader, Lyons, Kidger, arXiv:2402.09983, 2024）
- KPP（The Kinetic PreProcessor）（GitHub）
- pykpp（Barron H. Henderson, GitHub）
動作環境：macOS 27.0.1（Apple Silicon）、Python 3.15.0、JAX 0.12.0.1、diffrax 0.7.2、equinox 0.13.8。
本記事は 2026年10月時点の情報に基づいています。