技術要約: Incremental Transformer Neural Processes (incTNP)

1. 問題提起

Neural Processes (NPs)、特にTransformer Neural Processes (TNPs) は、時空間予測からテーブルデータのモデリングに至るまで、メタ学習タスクにおいて強力な性能を示してきました。しかし、既存のTNPのバリアントは、データが逐次的かつ連続的に到着する（例：リアルタイムのセンサー読み取りやデータベースの更新）ストリーミング・シナリオにおいて、決定的なボトルネックに直面しています。

標準的なTNPは双方向の自己注意（self-attention）メカニズムに依存しています。新しい観測値が到着すると、これらのモデルは内部表現を更新するために、コンテキスト履歴全体を再エンコードする必要があります。これにより、更新あたりの計算コストは二次的 ($O(N^2)$) となり、長さ $N$ のストリームにおける累積コストは三次的 ($O(N^3)$) になります。ターゲットポイントが逐次的に予測され、コンテキストに追加される自己回帰（AR）展開モードでは、この非効率性がさらに悪化し、高頻度の更新を現実的なアプリケーションにとって極めて高価なものにします。

核心となる課題は、標準的なTNPの予測性能や確率的一貫性（ベイズ的整合性）を犠牲にすることなく、安価で増分的な更新（線形時間計算量）をサポートするNPアーキテクチャを設計することです。

2. 手法

著者らは、因果的マスキングを用いた大規模言語モデル（LLM）の進展から適応させた、Incremental Transformer Neural Process (incTNP) を提案しています。その手法は、以下の3つの柱に基づいています。

A. 因果的マスキングとKVキャッシュ

標準的なTNPが双方向アテンションを使用するのに対し、incTNPは**因果的マスキング（causal masking）**を自己注意メカニズムに採用しています。これにより、トークン $i$ がトークン $j \leq i$ のみに注目できる下三角のアテンションマスクが強制されます。

• メカニズム: 因果的マスキングをKey-Value (KV) キャッシュと組み合わせることで、モデルは過去のトークンのキー行列とバリュー行列を保持します。
• 利点: 新しいデータポイントが到着した際、モデルはキャッシュされた履歴に対して新しいトークンのみを処理します。これにより、コンテキスト更新の計算コストが $O(N^2)$ から線形時間 $O(N)$ へと削減されます。

B. 稠密な自己回帰トレーニング (incTNP-Seq)

データの効率を最大化し、因果的マスキングによる受容野の制限に伴う潜在的な性能低下を相殺するために、著者らは稠密な自己回帰トレーニング戦略（incTNP-Seqと呼称）を提案しています。

• アプローチ: データを互いに素なコンテキスト集合とターゲット集合に分割する代わりに、データセット全体を単一のシーケンス $D_{seq} = [(x_1, y_1), \dots, (x_N, y_N)]$ として扱います。
• マスキング: 因果的マスキングは、自己注意だけでなく、クロスアテンション・ブロックにも適用されます。これにより、$n$ 番目のターゲットに対する予測が、履歴 $(x_{1:n}, y_{1:n})$ のみに依存することが保証されます。
• 結果: これにより、すべての可能なプレフィックス・コンテキストを単一のフォワードパスで同時に計算できるようになり、すべてのコンテキストサイズに対して計算コストを償却し、より稠密な教師信号を提供することが可能になります。

C. 理論的枠組み: 暗黙的ベイズ性

論文では、因果的マスキングが標準的なNPに備わっている置換不変性（permutation invariance）の特性を損なうという懸念に対処しています。確率的一貫性に与える影響を定量化するために、著者らは「暗黙的ベイズ性（implicit Bayesianness）」（Mlodozeniec et al., 2024）という指標を適応させています。

• 彼らは、モデルの予測規則が条件付き交換可能な（ベイズ的な）分布からどの程度逸脱しているかを測定するために、**KLギャップ（KL Gap）**を定義しました。
• このフレームワークは、モデルの予測と真のデータ分布の間のカルバック・ライブラー（KL）ダイバージェンスを、「交換可能化された（exchangeabilified）」バージョンへのダイバージェンスと、ギャップ項へと分解します。ゼロのギャップは、完全な条件付き交換可能性を意味します。

3. 主な貢献

1. アーキテクチャ: 因果的マスキングとKVキャッシュを利用して、線形時間のコンテキスト更新を可能にし、因子化モードおよび自己回帰モードの両方において、桁違いの高速化を実現するTNPのバリアントであるincTNPを提案。
2. トレーニング戦略: 稠密な自己回帰トレーニング戦略 (incTNP-Seq) が、標準的なTNPと同等またはそれを上回る予測性能を、等しいトレーニング予算の下で達成できることを実証し、因果的マスキングによる制限を効果的に緩和することを示した。
3. 一貫性の分析: 新しい暗黙的ベイズ性指標を用いた経験的検証。研究の結果、因果的構造にもかかわらず、incTNPは標準的な置換不変TNPと同等のベイズ的一貫性を保持する予測規則を維持していることが示された。

4. 実験結果

著者らは、1Dガウス過程回帰、テーブルデータ回帰（UCIベンチマーク）、および定常温度予測（HadISDデータセット）を含む、合成および実世界のタスクを用いてincTNPを評価しました。

• 予測性能:

  • 静的コンテキスト設定において、incTNP-Seqは、合成および実世界のタスクの両方で標準的なTNP-Dのベースラインと同等、あるいはそれを上回る性能を達成しました。
  • 自己回帰 (AR) モードにおいて、incTNP-SeqはTNP-Dを大幅に上回りつつ、計算コストを削減しました。例えば、温度予測タスクにおいて、精度を損なうことなくTNP-Dに対して3.5倍の高速化を実現しました。
  • 本モデルは強力な**長さ汎化（length generalization）**を示し、最大1,024ポイントのコンテキストで訓練されているにもかかわらず、30,000ポイントを超えるコンテキスト・ストリームに対しても良好に動作しました。

• 計算効率:

  • ストリーミング更新: incTNPは、ステップあたりの更新コストを $O(N^2)$ から $O(N)$ に削減しました。
  • 累積コスト: 長さ $N$ のストリームにおいて、累積コストはincTNPでは二次的 ($O(N^2)$) にスケールしますが、標準的なTNPでは三次的 ($O(N^3)$) にスケールします。
  • AR推論: AR展開において、incTNPは複雑度を $O(N_t \cdot N_s^2)$ から $O(N_t \cdot N_s)$ に低減し、長大なストリームに対する高忠実度なAR推論を可能にしました。

• 確率的一貫性:

  • KLギャップ指標を用いた解析により、incTNP-Seqは非因果的なTNP-Dと同等のKLギャップを示すことがわかりました。これは、ストリーミング設定における不一致の主な要因が、incTNPの特定の因果的マスキングメカニズムではなく、多くのNPバリアントに共通する条件付き一貫性の欠如にあることを示唆しています。
  • incTNPは、予測精度とベイズ的一貫性の両面において、スケーラブルなストリーミングガウス過程ベースライン（WISKIなど）を上回りました。

5. 意義と主張

本論文は、incTNPが、ストリーミング・ニューラルプロセスにおける計算のスケーラビリティと確率的コヒーレンスの間の溝を埋めることに成功したと主張しています。

• スケーラビリティ: 因果的マスキングとKVキャッシュを活用することで、incTNPは、長い履歴ストリームを伴う複雑なタスクにおいて、高頻度のリアルタイム更新を計算可能にし、これまでTNPのバリアントには欠けていた能力を実現しました。
• パフォーマンス: 因果的マスキングが性能を低下させるという直感に反して、稠密なトレーニング戦略により、incTNP-Seqは標準的なTNPと同等、あるいはそれ以上の性能を（Sim-to-Realの転送シナリオにおいても）達成できます。
• 一貫性: 本研究は、因果的マスキングが必ずしも「合理的」なベイズ更新を損なうわけではないという仮説に挑戦しています。経験的な結果は、inc-TNPが置換不変なベースラインと同様の暗黙的ベイズ性を維持していることを示しており、データ順序が時間によって固定されるストリーミング推論におけるその使用を正当化しています。

著者らは、incTNPが、歴史的情報を破棄することなく、大規模で進化し続けるデータセットに対して構造化データ（テーブルデータや時系列データなど）の基盤モデルをスケールさせるための実用的な経路を提供すると結論付けています。