CRUMB: Efficient Prior Fitted Network Inference via Distributionally Matched Context Batching

技術要約：CRUMB – 分布一致型コンテキスト・バッチングによる効率的な事前適合ネットワーク推論

問題提起
TabPFNやTabICLに代表される事前適合ネットワーク（PFN）は、インコンテキスト学習が可能な表形式の基盤モデルの一種である。これらはラベル付きの訓練セットをコンテキストとして受け取り、単一のフォワードパスでテストクエリの予測を生成する。これは、小規模から中規模のデータセットにおいて、従来の勾配ブースティング手法を凌駕することが多い。しかし、決定的なボトルネックがそのスケーラビリティを制限している。トランスフォーマーベースのPFNにおける自己注意（self-attention）メカニズムは、訓練サンプル数（$N$）に対して二次関数的（$O(N^2)$）にスケールする。このため、$N > 50,000$となる大規模なデータセットに対する推論は計算量的に極めて困難である。

既存の解決策は、精度と効率性の間のトレードオフに直面している。一様サンプリング（Uniform subsampling）は、訓練データをランダムに破棄することで速度を向上させるが、情報を犠牲にする。$k$-近傍法（$k$NN）のようなクエリ依存の手法は、テストポイントに対してコンテキストを最適化することで精度を高めるが、バッチ処理を妨げ、$T$個のテストポイントに対して$T$回の個別のフォワードパスを必要とする。Mixture of In-Context Prompters（MICP）のような他のアプローチは、訓練データをクラスタリングすることでクエリをバッチ化しようと試みるが、特に共変量シフト（covariate drift）が発生した場合、訓練コンテキストをテストデータの特定の分布に適合させることに失敗する場合がある。

手法
著者らは、モデルの再学習を必要としない、アーキテクチャに依存しない3段階の推論ラッパーであるCRUMB（Clustered Retrieval Using Minimised-MMD Batching）を提案する。この手法は、以下のステージを通じて、コンテキストの品質とバッチング効率の間の緊張関係を解消する：

1. テスト側クラスタリング（Test-Side Clustering）: テストクエリ（$D_{test}$）を、標準化された入力特徴量を用いた$k$-meansにより、$K$個のクラスター（$C_1, \dots, C_K$）に分割する。これにより、共通の訓練コンテキストを共有できる類似したテストポイントをグループ化することが可能になる。
2. 分布一致型コンテキスト選択（Distributionally Matched Context Selection）: 各テストクラスター$C_k$に対し、全訓練プール（$D_{train}$）から小さな訓練サブセット$S_k$（サイズ $n \ll N$）を選択する。幾何学的な近接度に基づく手法（例：重心最近傍法）とは異なり、CRUMBは、テストクラスターの経験分布と候補となる訓練サブセットとの間の**最大平均不一致（MMD）**を貪欲に最小化することによって$S_k$を選択する。
   • これは、テストクラスターへの近接性（関連性）を報酬とし、選択された訓練ポイント間の冗長性（多様性）を罰則とする、カーネルハーディング（kernel herding）を用いた貪欲アルゴリズムによって達成される。
   • このプロセスを加速させるため、著者らはガウスRBFカーネルを近似するためにランダムフーリエ特徴量（RFF）を採用し、バッチ化された貪欲選択を用いている。
3. バッチ推論（Batched Inference）: PFNモデルは$K$回の独立したフォワードパスを実行する。各パスでは、モデルは選択された訓練サブセット$S_k$をコンテキストとして受け取り、クラスター$C_k$全体をクエリとして受け取る。これにより、総フォワードパス数は（$k$NNにおける）$T$から$K$へと削減され、GPUの効率的な利用が可能になると同時に、分布の整合性が維持される。

主な貢献

• MMDベースのコンテキスト検索: 論文では、訓練サブセットを選択するためのメカニズムとして、貪欲なMMD最小化（カーネルハーディング）を導入している。著者らは、この分布一致アプローチが、同一のクラスタリングフレームワーク内において、重心最近傍法やVoronoi一様検索戦略よりも優れた性能を示すことを実証している。
• テスト側クラスタリングによるバッチコンテキスト選択: テストクエリをクラスタリングし、各クラスターに対して分布的に一致した単一の訓練コンテキストを割り当てることで、CRUMBはクエリの関連性を損なうことなく効率的なバッチングを可能にする。
• 共変量シフトへの堅牢性: 本手法は共変量シフトを自然に扱うことができる。訓練コンテキストが（固定された訓練分布ではなく）現在のテストバッチの分布に一致するように選択されるため、CRUMBはデータのシフトに適応する。これに対し、固定された訓練クラスターに依存するMICPのような手法は、テスト分布が訓練分布から離れると大幅に性能が低下する。

実験結果
著者らは、3つのPFNアーキテクチャ（TabPFNv2, TabICLv1, TabICLv2）を用いて、TabArenaベンチマーク（51の多様な表形式データセット）でCRUMBを評価した。

• ベースラインとの比較: TabArenaベンチマークにおいて、CRUMBは同一のコンテキスト予算において、一様サンプリングおよびMICP技術を大幅に上回る性能を示した。CRUMBは、性能面では$k$NNと統計的に近いレベルにあるが、これを$T$回のパスではなく、固定された$K$回のフォワードパス（例：$K=20$）で達成しており、大幅な推論時間の短縮を実現している（報告された実験では、35秒対361秒）。
• 大規模データセット: 最大120,000サンプル（例：GiveMeSomeCredit, Higgs）を含むデータセットを用いた実験では、早期停止（MMDの改善がプラトーに達した時点でコンテキスト選択を終了する）を用いたCRUMBが、特にTabPFNv2およびTabICLv1において、MICPに対してより頻繁に勝利した。
• 共変量シフト: テスト分布が第1主成分に沿ってシフトする制御されたドリフト実験において、CRUMBのMICPに対する優位性は、ドリフトの強度とともに単調に増加した。ゼロドリフト時、CRUMBは+4.9%リードしており、最大ドリフト（完全にアウトオブサンプル）時には、その差は+17.1%まで拡大した。これは、MICPの固定された訓練クラスターが不整合になる一方で、CRUMBが訓練コンテキストをシフトしたテスト分布に再整列させる能力を持っていることに起因する。

意義と主張
本論文は、CRUMBが、アーキテクチャの変更やモデルの微調整を必要とせずに、大規模な表形式データセットに対するPFN推論を実用的なものにすると主張している。テスト側のクラスタリングと、分布一致（MMD）に基づくコンテキスト検索を組み合わせることで、CRUMBは、クエリごとの検索の精度と、一様バッチングの効率性の間の「スイートスポット」を実現している。

著者らは、核心的な革新はテスト適応型設計にあると強調している。すなわち、幾何学的な近接性だけでなく、分布的な整合性（MMD）に基づいてテストポイントをクラスタリングし、コンテキストを選択するという点である。この設計は、既存のコンテキスト選択手法がしばしば失敗する共変量シフトの領域において、自然な堅牢性メカニメントを提供する。結論として、CRUMBは基礎となるPFNの事前学習による性能上限を継承するものの、PFNを大規模なデータセットに適用することを妨げていた計算上の障壁を効果的に取り除くものであるとしている。

限界
著者らは、CRUMBがクラスタリングを行うためにテストクエリの全セットを事前に必要とすることを認めており、これは（セントロイドをキャッシュすることで実質的にMICP型のルーティングに帰着させるような修正を加えない限り）真のオンライン・ストリーミング設定には適さないことを示唆している。加えて、ユークリッド$k$-meansおよび等方性ガウスカーネルへの依存は、高次元または混合型の特徴空間において性能を制限する可能性がある。最後に、純粋な推論ラッパーであるため、CRUMBはPFNの事前学習による事前分布と特定のデータセットとの間の不一致を補償することはできない。