Fitting Large Nonlinear Mixed Effects Models Using Variational Expectation Maximization

Mohamed Tarek と Pedro Afonso による論文「Variational Expectation Maximization を用いた大規模非線形混合効果モデルの適合」の詳細な技術的サマリーを以下に示す。

1. 問題定義

非線形混合効果（NLME）モデルは、薬物動態学、定量的システム薬理学（QSP）、および臨床試験における階層的および縦断データの分析の標準となっている。これらのモデルは、被験者固有のランダム効果（$\eta$）を積分することで周辺尤度を最大化し、集団パラメータ（$\theta$）を推定する。

核心的な課題：

• 計算上の非現実性： モデルが複雑化するにつれて（例えば、ニューラルネットワーク、高次元の常微分方程式（ODE）、または数千のパラメータを組み込む場合）、周辺尤度を計算するために必要な積分は、計算量的に実行不可能となる。
• 既存手法の限界： 一次条件推定（FOCE）や確率的近似 EM（SAEM）などの従来のアルゴリズムは、数値積分またはマルコフ連鎖モンテカルロ（MCMC）サンプリングに依存している。これらの手法はパラメータ数やランダム効果の数に対してスケーラビリティが低く、数千のパラメータを持つモデル（例：DeepNLME モデル）では収束に失敗したり、過度な時間を要したりすることが多い。

2. 手法：変分期待最大化（VEM）

著者は、スケーラブルな代替手段として**変分期待最大化（VEM）**の採用を提案する。MCEM や SAEM におけるランダム効果の事後分布からのサンプリングの代わりに、VEM はパラメトリック分布（変分族）を用いて事後分布を近似する。

主要な技術的構成要素：

• 変分推論（VI）フレームワーク：
  • 目標は、周辺尤度を直接最大化するのではなく、**エビデンス下限（ELBO）**を最大化することである。
  • ELBO は以下のように定義される：$\text{ELBO} = \mathbb{E}_{q}[\log p(y, \eta|\theta)] - \text{KL}(q(\eta) || p(\eta|y, \theta))$。
  • ELBO の最大化は、変分分布 $q(\eta)$ と真の事後分布との間の Kullback-Leibler（KL）ダイバージェンスを最小化することに等価である。
• 再パラメータ化トリック：
  • 勾配ベースの最適化を可能にするため、著者は再パラメータ化トリックを使用する。ランダム効果は、固定されたベース分布の決定論的変換として生成される（例：$\eta = T(\xi; \kappa)$ ここで $\xi \sim \mathcal{N}(0, I)$）。
  • これにより、サンプリングプロセスを通じて勾配が伝播し、効率的なオプティマイザの使用が可能となる。
• 同時最適化：
  • 従来の EM が E ステップと M ステップを厳密に交互に行うのとは異なり、Pumas 実装では、勾配ベースの手法を用いて集団パラメータ（$\theta$）と変分パラメータ（$\kappa$）の同時最適化を行う。
• 自動微分（AD）：
  • 出力に対してパラメータ数が非常に多いモデル（DeepNLME で一般的）に対しては、逆モード ADが利用される。これは、数千のパラメータに対する ELBO の勾配を効率的に計算するために不可欠である。
  • 実装は、小規模モデル向けのフォワードモードと、大規模モデル向けのリバースモードの両方の AD をサポートしている。
• 最適化戦略：
  • 決定論的 VEM： 事前サンプリングされたベース変数と、安定した高速な収束を実現するL-BFGSオプティマイザ（準ニュートン法）を使用する。
  • 確率的 VEM： 反復ごとのリサンプリングと、大規模集団のミニバッチ処理に適したAdamオプティマイザを使用する。
• 制約の処理： 制約付きランダム効果（例：対数正規分布）の場合、ガウス変分族で近似する前にドメイン変換（例：対数変換）を適用する。

3. 主要な貢献

1. NLME 用のスケーラブルなアルゴリズム： 本論文は、NLME モデルに特化した VEM 変種を導入し、従来の手法が失敗する領域である15,000 超の集団パラメータと16 以上のランダム効果を持つモデルの適合を可能にする。
2. Pumas への実装： このアルゴリズムは、Julia の高性能計算機能と自動微分を活用し、統計ソフトウェア Pumas に実装されている。
3. 理論的関連性： 著者は、VEM、変分オートエンコーダ（VAE）、確率的 PCA（PPCA）、およびラプラス近似との関係を明確にした。また、ラプラス法は本質的に VEM 目的関数の 2 次テイラー近似であることを示した。
4. 実証的検証： 本論文は、小規模（ワルファリン）および大規模（DeepNLME）モデルの両方において、VEM と FOCE を比較する厳密なベンチマークを提供している。

4. 実験結果

著者は Pumas ソフトウェアを使用して 2 つのモデルで実験を行った。

A. ワルファリンモデル（小規模検証）

• 設定： 31 人の被験者と 464 件の観測値を持つ標準的な PK-PD モデル。
• 比較： VEM 対 FOCE。
• 結果：
  • 精度： VEM は FOCE とほぼ同一のパラメータ推定値と対数尤度を生み出した（対数尤度：-994.5 対 -995.3）。
  • 速度： VEM は著しく高速であった（16.6 秒対 44.4 秒）。
  • 結論： VEM は、標準的な薬物動態モデルに対する正確かつ効率的な代替手段である。

B. DeepNLME Friberg モデル（大規模スケーラビリティ）

• 設定： 15,410 の集団パラメータと 16 のランダム効果を持つニューラルネットワーク（MLP）を組み込んだ人工的に大規模なモデル。
• 比較： VEM（5 反復）対 FOCE（1 反復を試行）。
• 結果：
  • VEM： リバースモード AD を使用して、約 3092 秒（1 反復あたり約 618 秒）で 5 反復を完了した。
  • FOCE： 24 時間以内に 1 反復も完了できなかった。
  • 結論： VEM は、FOCE が計算的に非現実的となる複雑で高次元のモデルに成功裏にスケーリングする。

5. 意義と含意

• DeepNLME と QSP の実現： 数千のパラメータを持つモデルを適合できる能力は、DeepNLME（メカニズムモデルへのニューラルネットワークの統合）や複雑な**定量的システム薬理学（QSP）**モデルへの扉を開く。これらは以前は NLME として適合させるには計算コストが高すぎた。
• 最適化パラダイムの転換： 本論文は、大規模問題に対して MCMC ベースのサンプリング（SAEM/MCEM）から、現代の深層学習インフラ（リバースモード AD、GPU/TPU 互換性）を活用した勾配ベースの変分推論への転換を提唱している。
• 実用的な限界： 著者は、VEM がハイパーパラメータ（サンプル数、変分族の選択）に敏感であり、FOCE が推定中にこれらを行うのとは異なり、最適化後に経験的ベイズ推定量（EBE）と対数尤度の近似を計算するために別々のステップが必要であることを認めている。
• 将来の方向性： この手法は、従来の手法が計算コストの「壁」に直面する大規模な結合バイオマーカーモデルや複雑な生物学的システムの適合における標準的なツールとなる準備ができている。

要約すると、本論文は、逆モード自動微分と組み合わせた変分期待最大化が、薬物動態学における次世代の複雑で高次元な非線形混合効果モデルの適合のための、堅牢でスケーラブルかつ正確なフレームワークを提供することを示している。