Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

本論文は、合成的な時間的因果モデルを用いて事前学習された、ドメイン固有の学習を必要とせずに反事実的な縦断的アウトカムを予測する凍結された事前適合ネットワークであるCausalLongPFNを紹介しており、がん、HIV、MIMIC-IIIといったベンチマークにおいて競争力のある性能を示すとともに、従来のコホート固有のモデリングに代わる費用対効果の高い選択肢を提示している。

要約

あなたは、特定の治療計画に対して患者が今後数週間にわたってどのように反応するかを予測しようとしている医師だと想像してください。問題は、患者は複雑であるということです。彼らの健康状態は日々変化し、過去の治療が現在の状態に影響を与え、さらに食事やストレスといった他の要因が結果を妨害します。通常、これらの予測を行うには、遭遇する新しい患者グループごとに、全く新しい高度に専門化されたコンピュータモデルを構築しなければなりません。これは、新しい近隣に引っ越すたびに、その家を設計するために新しい建築家を雇うようなものです。それは時間がかかり、コストもかかり、多くのデータを必要とします。

この論文では、CAUSALLONGPFNと呼ばれる新しいツールを紹介しています。これを**「汎用的な医学的直感エンジン」**と考えてください。このエンジンは、実際の患者に出会う前に、すでにゲームのルールを学習済みです。

仕組みをシンプルな概念に分解すると、以下のようになります。

1. 「トレーニングキャンプ」（合成データによる事前学習）

実際の患者が到着するのを待つ代わりに、作成者たちは大規模な仮想の「トレーニングキャンプ」を構築しました。このキャンプでは、何百万もの異なる体型、疾患、治療反応を持つ、何百万もの架空の患者をシミュレートしました。彼らは、これらの架空の患者に複雑な挙動をプログラムしました。

• ゆっくりと回復する患者もいれば、急激に悪化する患者もいます。
• 治療がすぐに効く場合もあれば、効果が遅れて現れる場合もあります。
• 隠れた遺伝子に基づいて、異なる反応を示す患者もいます。

AIモデルはこのキャンプですべての時間を費やし、架空の患者の転帰（アウトカム）を予測することを学びました。モデルは単に答えを暗記したのではなく、治療、時間、そして生物学がいかに相互作用するかという「根底にある論理」を学んだのです。

2. 「凍結されたエキスパート」（新たな学習は不要）

ここが魔法のようなトリックです：モデルがトレーニングキャンプを終えた後、彼らはそれを**「凍結」**しました。モデルの脳をロックしたのです。モデルは新しいことを学んだり、内部設定を変更したりすることはできません。

新しい実患者のグループ（例えば、ある病院の癌患者）が到着したとき、モデルはゼロからスタートすることはありません。モデルは再学習する必要もありません。代わりに、モデルは**「カルテを読む超優秀なインターン」**のように振る舞います。

• サポート・トラジェトリー（支持軌跡）： 現在の病院の患者の例（「サポート」）をいくつかモデルに示します。これらは、インターンにその病院特有のスタイルを理解させるための、いくつかの症例ファイルを読ませるようなものです。
• クエリ（問い）： 「もしこの特定の患者に、治療Aを今後5日間行った場合、何が起こるか？」とモデルに尋ねます。
• 回答： モデルは、トレーニングキャンプで学んだことと、提示された症例ファイルから得たコンテキスト（文脈）を即座に組み合わせ、結果を予測します。これは、「勾配降下法」（通常の再学習プロセスにおけるテクニカルな用語）を一切行うことなく行われます。

3. 「タイムトラベル・シミュレーター」

このモデルは**縦断的（longitudinal）**なデータを扱うように設計されており、つまり「時間」を理解しています。単に明日の結果を推測するのではなく、未来をステップ・バイ・ステップでシミュレートします。

• 1日目を予測します。
• その予測を2日目の出発点として使用します。
• このプロセスを繰り返し、5日目、6日目、あるいは7日目に何が起こるかをシミュレートします。

これは、離陸だけでなく、たとえ飛行中に天候が変化したとしても、パイロットの決定に基づいて飛行経路全体をシミュレートするフライトシミュレーターのようなものです。

4. なぜこれが重要なのか（結果）

著者たちは、この「凍結されたエキスパート」を、従来の方法（病院ごとに新しいモデルを構築する方法）と比較検証しました。

• テスト： 彼らは、癌、HIV、血液希釈剤（ワルファリン）、および実際のICUの記録のデータを使用しました。
• 結果： 凍結されたモデルは、特定の疾患に合わせて特別に訓練されたモデルと同等、あるいはそれ以上の性能を発揮しました。
• 大きな勝利： このモデルは、シミュレーションにおいて「もし〜だったら」というシナリオを試すことが倫理的に不可能な、実際のICUデータにおいて特に優れた成果を上げました（なぜなら、実際の患者に対して異なる治療を試すことはできないためです）。モデルは、観察されたデータのみに基づいて、次に何が起こる可能性が高いかを予測することができました。

まとめ

この論文は、新しい医学的データセットごとにカスタムモデルを構築する必要は必ずしもない、と主張しています。代わりに、幅広い「もしも」のシナリオ（合成データ）を用いて一つの大規模なモデルを訓練し、それを凍結された、すぐに使えるツールとして利用できるのです。

これは、あらゆる種類の料理を仮想のキッチンで練習してきたマスターシェフがいるようなものです。新しい食材（新しい患者グループ）を持ってきたとき、彼らは料理の仕方を学び直す必要はありません。ただ、あなたの食材を見て、膨大な事前学習による直感を用いて、即座に料理の準備を始めることができるのです。

重要な注意点： 著者は、このツールは予測と研究のためのものであり、臨床における最終的な生死に関わる決定を下すためのものではない、と慎重に述べています。これは医師が潜在的な結果を理解するのを助けるものですが、他のあらゆる因果モデルと同様に、標準的な医学的仮定に基づいています。これは強力な研究ツールであり、医師の判断に取って代わるものではありません。

---

実装とモデルの詳細について
コードと事前学習済みモデルの重みは、以下で公開されています。

• Code on GitHub: https://github.com/Amirhossein-Zare/causal-long-pfn
• Pretrained model weights on Hugging Face: https://huggingface.co/Amirhossein-Zare/causal-long-pfn

技術概要

技術要約：因果的縦断的事前分布適合ネットワーク（Causal Longitudinal Prior-Fitted Networks）

問題の定式化

本論文は、縦断的な観測データにおける将来の治療シーケンス下での潜在的アウトカムを予測するという課題に取り組んでいる。核心となるタスクは、時刻 $t$ までの観測された履歴 $H_t$ に基づく、計画された将来の治療シーケンス $\bar{a}$ に対する履歴条件付き潜在的アウトカム $E[Y(\bar{a})_{t+\tau} | H_t]$ の推定である。

この問題は、主に以下の3つの要因によって複雑化している：

1. 時変共変量（Time-varying confounding）: 各ステップでの治療割り当ては、それ自体が先行する治療の結果である共変量に依存する。
2. 異質な患者ダイナミクス: 個々の軌跡は、複雑で非線形な状態進化と潜在的な不均一性を示す。
3. データの制限: 観測コホートは、新しいドメインやシミュレータごとに信頼できるディープ・シーケンス・モデルをゼロから訓練するには、規模が小さすぎる場合が多い。

既存の縦断的因果推論推定器（例：RMSN, CRN, G-Net, Causal Transformer）は、通常、新しいコホートごとにハイパーパラメータ選択や傾向スコア・モデリングを含む、個別の教師あり学習の実行を必要とする。このパイプラインはコストが高く、ドメインごとに繰り返しの訓練が必要な場合には非実用的である。

手法：CAUSALLONGPFN

著者らは、広範な時間的構造因果モデル（TSCM）の事前分布に対して、縦断的因果予測を償却（amortize）するように設計された、事前分布適合型インコンテキスト予測器である CAUSALLONGPFN (Causal Longitudinal Prior-Fitted Networks) を提案する。

1. 広範な事前分布を用いた合成データによる事前学習

本モデルは、広範なTSCMの事前分布からサンプリングされた合成エピソードを用いて、完全に事前学習される。この事前分布は、単一のシミュレータを複製するのではなく、広範なクラスの縦断的因果ダイナミクスをカバーするように設計されている。合成事前分布の主な特徴は以下の通りである：

• 因果的時系列グラフ: 疎な同時およびラグのある依存関係と、非循環的な瞬時グラフ。
• 非線形メカニズム: 状態座標は、多様な基本非線形関数（identity, tanh, sinusoidal, ReLUなど）と様々なノイズ分布を用いた、疎な非線形自己回帰更新に従う。
• ダイナミカル・モチーフ: 遅延効果や制御フィードバックといった定性的なメカニズムを捉えるため、アクション・メモリ、飽和、恒常性、フィードバック制御、平滑化出力チャネルなどの構造化されたモチーフが重ね合わされている。
• 交絡を伴う行動ポリシー: 治療は、潜在的な個体不均一性（$Z_i$）に影響を受ける状態依存的な確率的ポリシーからサンプリングされ、時変的な治療と共変量のフィードバックを生み出す。
• アウトカムモデル: スカラーのアウトカムは、直接的かつ累積的な治療効果を持つ自己回帰的リードアウトを介して生成される。

2. アーキテクチャ

CAUSALLONGPFNは、デュアルエンコーダ・アーキテクチャを採用している：

• 因果的履歴エンコーダ（Causal History Encoder）: 軌跡レベルの因果的Transformer（マスク付き自己注意機構を使用）であり、縦断的シーケンスを履歴表現へと写像する。これにより、時刻 $r$ における表現は、時刻 $r$ までに利用可能な情報のみに依存することを保証する。
• PFNコンテキストエンコーダ（PFN Context Encoder）: インコンテキスト適応を行うTransformerベースのコンテキストエンコーダである。これは、サポート・トラジェトリ（順序のない集合として扱われる）とクエリ・トークンを自己注意機構によって結合して処理する。サポート・トラジェトリの順序に対する位置エンコーディングは割り当てられない。
• ガウス混合予測ヘッド（Gaussian-Mixture Prediction Head）: 最終的なクエリ表現は、正規化された次ステップのアウトカムに対する5成分のガウス混合分布をパラメータ化し、点予測と不確実性の推定の両方を提供する。

3. インコンテキスト予測とロールアウト

テスト時、モデルは**凍結（frozen）**された状態で動作する。モデルは以下を受け取る：

• サポート・トラジェトリ: 新しいドメイン/タスクからの事例。
• クエリ履歴: 時刻 $t_{obs}$ までに観測されたもの。
• 提案された将来の治療シーケンス: 計画された介入。

モデルは、勾配更新、傾向スコア・モデリング、または敵対的バランシングを行うことなく、予測分布を返す。多ステップ予測（$\tau > 1$）の場合、モデルは自己回帰的なプラグイン・ロールアウトを実行する。すなわち、1ステップ先のアウトカム分布を予測し、その混合平均をクエリ履歴に挿入し、指定された治療シーケンスの下でプロセスを繰り返す。

主な貢献

1. 縦断的因果予測のための事前分布適合モデル: CAUSALLONGPFNは、計画された縦断的治療シーケンス下での履歴条件付き潜在的アウトカム予測を行う、初のPFNスタイルのモデルである。これは、テスト時の適応を必要としない凍結されたモデルとして機能する。
2. 縦断的因果タスクに対する合成事前分布: 著者らは、治療と共変量のフィードバック、潜在的な不均一性、非線形ダイナミクス、遅延/累積効果、およびレジーム変化を備えた多様なタスクを生成するTSCM事前分布を設計した。
3. 縦断的インコンテキスト推論のためのアーキテクチャ: 因果的Transformer履歴エンコーダとPFNコンテキストエンコーダ、およびガウス混合ヘッドを組み合わせた新しいデュアルエンコーダ。
4. 自己回帰的カウンターファクチュアル・ロールアウト: 学習された1ステップ予測器を、決定論的なプラグイン・ロールアウトを介して多ステップ予測へと拡張。
5. ゼロショット評価: 分岐可能な反事実的ベンチマークおよび事実的な実世界データにおいて、単一の凍結モデルをドメイン特化型学習済みベースライン（MSM, RMSN, G-Net, CRN, Causal Transformer, G-Transformer）と比較した包括的な評価。

結果

モデルは、がん腫瘍成長、ワルファリンPK/PD、HIV治療ダイナミクス、およびMIMIC-III ICU軌跡の4つのベンチマークで評価された。

• ドメイン均衡性能: CAUSALLONGPFNは、ドメイン均衡1ステップ正規化RMSE（0.222）において最高値を達成し、MSMやRMSNを僅差で上回った。5ステップ予測では、RMSNとG-Netに次いで全体で3位となったが、MSM、CRN、およびTransformerベースのベースラインを上回った。
• ドメイン別性能:
  • MIMIC-III (Factual): モデルは1ステップおよび5ステップ予測の両方で第1位となり、ドメイン固有の学習なしに実世界の臨床軌跡への強力な転移性能を示すことを実証した。
  • 反事実的ベンチマーク (Cancer, HIV, Warfarin): モデルは競争力を維持し、1ステップのタスクで2位または3位となった。しかし、より長期の予測（例：がんの5ステップ予測）においては、特定の適合のために十分なターゲットドメインデータが存在する場合に優位性を持つ、専門化されたドメイン学習済み回帰モデル（RMSN, CRN）の方が低い誤差を達成しており、これを示唆している。
• 不確実性の較正（Calibration）: ガウス混合ヘッドは有用な分布情報を提供した。較正の精度はドメインによって異なり、ワルファリンは最も良好な較正を示し、MIMIC-IIIは高い不均一性により広い信頼区間を示した。

意義と主張

本論文は、以下の条件下において、広範な合成的因果事前学習が、繰り返しのドメイン特化型学習に代わる有用な凍結された代替手段を提供できると主張している：

• 特化したモデルの再学習がコスト面または実用面で困難である場合。
• 新しいコホートへの迅速な適応が求められる場合。
• 反事実的な教師信号が入手できない場合（MIMIC-IIIのような実世界の事実的予測タスクのように）。

著者らは、CAUSALLONGPFNが、観測データを解釈するために必要な標準的な因果仮定（一貫性、正値性、逐次的交換可能性）を取り除くものではないことを強調している。むしろ、推定問題を償却するものである。結果は、十分に広範な合成事前分布が、治療応答タスク間で再利用可能な構造を捉えることができ、モデルを強力な汎用インコンテキスト予測器にすることを示唆している。ただし、論文内では、特に長期の反事実的予測において、十分なターゲットドメインのデータと検証信号がある場合には、ドメイン特化型の学習が依然として価値を持つことも謙虚に述べている。

本研究は、単独の臨床意思決定システムではなく、因果的シーケンスモデリングおよび仮説生成のための研究ツールとして位置付けられており、因果仮定や事前分布のサポートが不適切な場合に予測を過信することに対して警告を発している。

コードとモデルの利用可能性

実装コードはGitHub（https://github.com/Amirhossein-Zare/causal-long-pfn）で、事前学習済みモデルの重みはHugging Face（https://huggingface.co/Amirhossein-Zare/causal-long-pfn）で公開されています。