Towards Continuous-time Causal Foundation Models

本論文は、非結合観測による微細グリッド統合を通じて観測スケジュールに対する軌道法則不変性を保証する連続時間因果基盤モデルを提案し、多様な非線形ダイナミクスおよび不規則なデータシナリオにおいて、単純な統合手法を上回る性能を実証する。

要約

超知能のAIに、時間とともに変化するシステム（例えば、人体を巡る薬の動きや、突風に対する風洞の反応など）の未来を予測する方法を教えることを想像してみてください。

通常、AIモデルは時間を「スナップショット」で捉えます。1秒、2秒、3秒といった固定された瞬間をページとする、パラパラ漫画のようなイメージです。しかし、現実世界は時計の針が刻むのを待って動き出すわけではありません。それは連続して流れています。

この論文は、AIに単なるスナップショットではなく、その「流れ」を理解させる方法について述べています。以下に、簡単なアナロジーを用いて解説します。

1. 問題点：「停止と開始」の罠

著者たちは、ある一般的な過ちを指摘しています。連続的なプロセス（例えば、流れる川）をAIに教える際、特定の時刻に撮影されたスナップショットだけを見せると、AIは川そのものではなく、スナップショットの「スケジュール」を学習してしまうのです。

• アナロジー： 車の加速方法を学ぼうとしていると想像してください。
  • 悪い方法（離散的/単純）： 瞬きをするたびにスピードメーターを見るだけです。瞬きが遅ければ、ゆっくりとした加速に見え、瞬きが速ければ、速い加速に見えます。AIは「私がどのくらい速く瞬きするか」が速度を決定すると学習し、エンジンそのものを学習しません。
  • 結果： AIは混乱します。瞬きの新しいスケジュールを見せると、それは失敗します。なぜなら、AIは車の物理法則ではなく、あなたの瞬きのパターンを学習してしまったからです。

2. 解決策：「高速カメラ」

この論文は、これらのモデルを訓練する新しい方法として**「連続時間因果基盤モデル（Continuous-Time Causal Foundation Models）」**を提案しています。区切りごとに1枚のスナップショットを撮る代わりに、「高速カメラ」のアプローチを用います。

• アナロジー： 車を理解するために、エンジンが稼働している様子を超高速（1秒間に数千フレーム）で記録し、加速の完璧で滑らかなビデオを作成します。その後、AIにその滑らかなビデオを見せます。
• トリック： AIが「テスト」されるのは、1日1回患者をチェックする医師のように、遅いスナップショットだけであっても、AIはすでに高速訓練を通じて滑らかで連続的な物理法則を学習しています。それは「スナップショットの法則」ではなく、「川の法則」を知っているのです。

3. 訓練の3つのレベル

著者たちは、異なるモデルが時間をどのように処理するかを分類する「ティアリスト」を作成しました。

• ティア1（パラパラ漫画）： 古い方法です。AIは固定された時間ステップしか知りません。タイミングが変われば失敗します。
• ティア2（怠惰なカメラマン）： AIは連続的になろうとしますが、観測の間に1回だけ写真を撮ります。これはまだマシですが、時間間隔が変われば依然として混乱します。ぼやけた写真2枚だけに基づいて車の速度を推測するようなものです。
• ティア3（高速のプロ）： これがこの論文が達成したものです。AIは超微細なグリッド（数千の微小ステップ）上で物理をシミュレーションし、その後、AIに見せる必要がある特定の時刻だけを「提示」します。
  • 結果： AIはシステムの真の、不変の法則を学習します。観測が1秒ごと、1時間ごと、あるいはランダムな時刻に来ても、それは気にしません。

4. 実験：実際に機能するか？

チームは、2種類の「物理エンジン」を用いてこれをテストしました。

1. 線形： 単純な直線的な物理（バネなど）。
2. 非線形： 複雑でねじれた物理（カオス的な気象システムなど）。

彼らは「怠惰なカメラマン（ティア2）」と「高速のプロ（ティア3）」を対決させました。

• 発見： 高速のプロが毎回勝利しました。
• 驚き： 高速法で訓練されたAIは、観測間の「経過時間」をわざわざ教えられる必要さえありませんでした。それは自然に流れを理解したのです。しかし、怠惰な方法で訓練されたAIは、うまく機能させるために、時間間隔を明示的に教えられなければなりませんでした。

5. 実世界テスト（「ゼロショット」テスト）

著者たちは、これまで見たことのない実世界のデータ（ゼロショット）に、新しいAIを適用しようと試みました。

• 薬物動態学： 血液中の薬物濃度の予測（テオフィリンとワルファリン）。AIは合成データで訓練されたにもかかわらず、薬物の上昇と下降を驚くほどよく追跡できました。
• 物理システム： 風洞実験。AIは、ファン出力の急激な変化に対して風洞の速度がどのように反応するかを正常に予測しました。

結論

この論文は、AIのためのより優れた「タイムマシン」を構築します。データポイント間の間隔を単に暗記するのではなく、物事が変化する滑らかで連続的な法則（高速シミュレーションを用いて）をAIに学習させることで、データが奇妙で不規則な時刻に到着する場合でも、AIは未来を予測する能力が大幅に向上します。

この論文が主張していないこと：

• これはまだ医師やエンジニアを代替する準備ができているとは主張していません。
• すべての種類の時系列問題を解決するとは主張していません。
• 実世界でのテストは「予備的」であり、重要な状況で使用される前にはさらに多くの作業が必要であると認めています。

これは基礎的な一歩です。AIに時間を一連の踏み石としてではなく、流れる川として見るように教えることで、AIは宇宙の法則をはるかに良く学習できることを証明したのです。

技術概要

技術的概要：連続時間因果基盤モデルへの道

1. 問題定義

Prior-Data Fitted Networks（PFN）は、合成構造因果モデル（SCM）上でトランスフォーマーを事前学習させることで、因果推論を表形式データおよび離散時間時系列に拡張することに成功しました。しかし、既存の時間的因果事前分布は離散整数グリッド上で動作します。これらを連続時間へ単純に拡張しようとする試み、すなわちメカニズムを確率微分方程式（SDE）として書き直し、各観測ギャップごとに一度積分するだけでは、真の連続性は達成されません。

核心的な問題は、SDE を観測間隔でのみステップさせる（単純な積分）場合、軌道の結合法則が特定の観測スケジュールに依存してしまうことです。その結果、事前分布は実質的に「SDE 衣装をまとった離散時間マルコフモデル」のままとなり、データ生成プロセスがいつ観測されるかによって不変であるべきという要件を満たせません。この制限は、薬物動態学（臨床的に選択されたサンプリング時間）、可変遅延イベントを伴う物理系、欠損データを有する電子医療記録など、不規則でスケジュールが不均一なデータを持つ領域において決定的に重要です。

2. 手法

2.1. 連続時間因果事前分布の定義

本論文は、連続時間因果事前分布に対する厳密な基準を確立します：サンプリングされた軌道の結合法則は、観測スケジュールに対して不変でなければならないという基準です。観測スケジュールは、基礎となる時間的 SCM（TSCM）の一部ではなく、純粋な測定として扱われます。

この基準に基づき、著者は 3 段階の分類体系を提案します：

• ティア（A）離散：整数グリッド上でのみ定義される標準的な離散時間 SCM。
• ティア（B）単純連続：観測ギャップごとに 1 回積分される SDE（観測グリッド上のオイラー・マルヤマ法）。軌道の法則はギャップサイズ $\Delta_i$ に依存して変化し、連続性の基準を満たしません。
• ティア（C）微細グリッド連続：SDE を微細グリッド（$\Delta_{fine} \ll \min \Delta_{obs}$）上で積分し、その後観測スケジュールにサブサンプリングします。$\Delta_{fine} \to 0$ において、これは真の SDE 法則に収束し、有限ステップにおいて連続性の基準を近似して満たします。

2.2. 連続時間事前分布の構築

提案された構築は、以下のコンポーネントを持つランダム有向非巡回グラフ（DAG）上でティア（C）を実現します：

• グラフサンプリング：変数はランダム DAG または標準的な構造（例：バックドア、フロントドア、道具変数）からサンプリングされます。隠れた交絡因子を含めることも可能です。
• メカニズムファミリー：
  • 線形ドリフト：ドリフトが親変数の線形結合であるオルンシュタイン・ウレンベック（OU）過程。
  • 非線形ドリフト：線形な親の和を置き換え、軌道の安定性を確保するために有界化された tanh 活性化関数を持つ小型の多層パーセプトロン（MLP）。
• レジームスイッチング：軌道の一部は、構造的ブレイク（例：薬理学における吸収相対して排泄相）をモデル化する、粘着性マルコフ遷移行列を持つ連続時間レジームスイッチング TSCM に従います。
• 介入：事前分布は、特定のウィンドウにおけるハード（値の設定）、ソフト（ドリフトのシフト）、時間変化型介入をサポートします。反事実的推論は、同じウィーナーノイズを再利用することで生成されます。
• シミュレーション：軌道は、各微細ステップでリサンプリングされたブラウン増分を用いてオイラー・マルヤマ法で SDE を微細グリッド上で積分し、その後不規則な観測スケジュールにサブサンプリングすることで生成されます。

2.3. アーキテクチャ：$\Delta t$ 認識型 PFN エンコーダ

このモデルは、介入前のウィンドウ上で動作する因果トランスフォーマーエンコーダを使用します。

• 時間埋め込み：学習された整数位置埋め込みの代わりに、モデルは連続時間のフーリエ埋め込みを使用します：$\phi(t) = W_\phi [\sin(2\pi f_k t), \cos(2\pi f_k t)]$。
• ギャップ埋め込み：観測間ギャップ（$\Delta t_i$）は、$\log(1+\Delta t_i)$ 変換の後に同じファミリーを用いて埋め込まれます。
• 推論：モデルは、観測データ、タイムスタンプ、介入仕様、およびクエリ時刻を入力として受け取り、介入下での結果の分布を予測します。

3. 主な貢献

1. 連続性基準：観測スケジュールに対する軌道法則の不変性を要求する形式的な定義であり、3 段階の分類体系を通じて運用化されました。
2. ティア（C）構築：微細グリッド積分、ランダム DAG、OU/MLP ドリフト、および不規則スケジュールを用いた連続時間事前分布の実用的な実現。
3. 実証的検証：エンコーダ×積分器の制御された $2 \times 2$ 除去研究により、微細グリッド積分が単純積分よりも優れており、特に評価グリッドが細かくなるほどその差が顕著であることを実証しました。

4. 実験結果

4.1. 除去研究

著者は、2 つの事前分布（線形 OU および非線形ニューラルドリフト）と 2 つの積分器（単純対微細）、および 2 つのエンコーダ（位置情報のみ対時間認識型）を用いて PFN を訓練しました。

• 積分器の性能：微細グリッド積分は、両方の事前分布および評価離散化にわたる 8 つの実験セルすべてにおいて、単純積分を上回りました。性能差（$\Delta$）は、評価グリッドが細かくなるにつれて単調に増大しました（例：ニューラル事前分布において、評価サブステップが細かくなるにつれて、差は +0.0048 から +0.0088 へ増加）。これは、微細グリッド訓練がモデルを真の SDE 極限に整合させるのに対し、単純訓練は離散化バイアスを導入することを確認しました。
• エンコーダの性能：時間認識型エンコーダ（ギャップのフーリエ埋め込み）の利点は、積分器に依存していました。
  • 単純積分の場合、時間認識型エンコーダは位置情報のみのエンコーダを大幅に上回り、スケジュール依存のダイナミクスを補償しました。
  • 微細積分の場合、エンコーダの選択は実証的に不変（差なし）であり、データ生成プロセスが十分にスケジュール不変となり、明示的なギャップ特徴の必要性がなくなったことを示唆しています。

4.2. ゼロショット転移（予備的）

本論文は、微調整なしで 3 つの実世界データセットに対する予備的なゼロショット転移結果を報告します：

• 薬物動態学（テオフィリンおよびワルファリン）：モデルは、投与量駆動の軌道を追跡するワルファリン血漿濃度において、強い相関（$r \approx 0.88$）を達成しました。テオフィリンの性能は中程度でした（線形モデルで $r \approx 0.53$）。著者は、濃度データの狭いクラスタリングにより単純ベースラインに対する RMSE の改善は小さいものの、ピアソン相関が動的な追跡を確認したと指摘しています。
• 物理系（因果チャンバー）：風洞インパルス装置において、混合メカニズム PFN は RPM ダイナミクスで $r = 0.95$ のピアソン相関を達成し、線形モデル（$r = 0.39$）を大幅に上回りました。これは、モデルが非線形で飽和する指数関数的ダイナミクスを成功裡に捉えたことを示唆しています。

5. 意義と主張

本論文は、離散モデルに対する「SDE 衣装」を超えて、因果基盤モデルに対する厳密な連続性基準を提供すると主張しています。主な意義は、評価グリッドが細かくなるにつれて性能差が増大するという証拠により、この基準を実現するために微細グリッド積分が必要であることを実証した点にあります。

著者は、実世界応用に関する主張において慎重です：

• ゼロショット転移結果は「予備的」かつ「補強的」と記述されており、まだ分野固有のベースライン（例：PK における NONMEM）と競争できる段階ではありません。
• 因果チャンバーでの成功は、構造的に適さない「白色ノイズ」ベンチマークから、明示的な二値介入と実ダイナミクスを持つデータセットへの切り替えを必要としました。
• 論文は、マルチシード反復の必要性、現在のニューラルドリフトが時間相関ノイズ（マルコフノイズのみ）を捉えられないという限界、および実データ転移の予備的性質を含む限界を認識しています。

この研究は、不規則な観測スケジュールを持つ SDE 駆動 TSCM のファミリー全体にわたる因果推論をトランスフォーマーが償却することを可能にする構築を提供し、真の連続時間因果推論への基礎的な一歩として位置づけられています。