技術要約：事前学習済み信念表現を用いた効率的な適応型データ収集（POLAR）

問題提起

適応型データ収集は、ベイズ実験計画法（BED）、ベイズ最適化（BO）、能動学習（AL）といった逐次学習問題を含みます。これらの設定において、学習者は有限の予算の下で、ダウンストリームの効用（例：情報利得の最大化、グローバルな最適解の特定、または予測精度の向上）を最大化するように、クエリ（設計）を逐次的に選択しなければなりません。

獲得ポリシーを学習するための既存のアプローチは、一般に以下の2つの問題のあるカテゴリに分類されます。

1. 事後分布に基づく手法： 観測されたデータに対して確率モデルを適合させ、事後分布を計算することに依存し、その事後分布を用いて獲得関数（例：期待改善量、期待情報利得）を評価します。これらの手法は、事後分布近似の質によってボトルネックが生じます。モデルの誤設定や近似誤差は、獲得ルールを歪ませる可能性があり、これらの誤差は逐次的なラウンドを通じて累積します。
2. 直接的なポリシー学習手法： 生の観測履歴から決定へと直接マッピングを行います（例：強化学習による）。これらは明示的な事後推論を回避しますが、生のデータから決定への複雑なマッピングをゼロから学習する必要があります。これは、データセットを表現する方法と、その表現に基づいて行動する方法を同時に学習する必要があるため、サンプル効率の低下や学習の困難さを招くことがよくあります。

手法：POLAR

著者らは、表現学習とポリシー学習を分離するフレームワークである POLAR（Policy Learning with Belief Representations）を導入しています。核心となる洞察は、最適な逐次的決定は、観測履歴が十分な 信念状態（belief state） を通じてのみ依存するという点です。POLARは、ゼロからこの状態を学習したり、明示的な事後分布近似に頼ったりするのではなく、事前学習済みの予測基盤モデル（foundation models） を信念状態エンコーダとして活用します。

アーキテクメント

1. 信念エンコーダ（バックボーン）： POLARは、多様な合成タスクで事前学習されたTabPFNやTabICLなどの表形式基盤モデル（Tabular Foundation Model: TFM）を利用します。履歴 $D_{t-1}$ と候補となる設計 $x$ が与えられると、TFMは候補に条件付けられた隠れ表現 $r_\psi(x, D_{t-1})$ を生成します。

   • TFMが出力する最終的な予測分布とは異なり、この隠れ表現は 信念表現 として扱われます。これは、予測がデータによってどのように変化するか（予測不確実性、局所的な構造、コンテキスト依存の振る舞いなど）を、特定の事後分布近似にコミットすることなくエンコードします。
   • 本フレームワークは、有限の候補集合 $C_t$ から選択を行う プールベースの設定 で動作します。バックボーンは履歴を一度処理し、プール内のすべての候補に対する表現を生成します。

2. ポリシーヘッド： 軽量で学習可能なポリシーヘッド $g_\eta$ は、信念表現 $r_\psi$ を各候補に対するスカラーのロジットへとマッピングします。これらのロジットに対するソフトマックス関数によって、確率的ポリシー $\pi_\eta$ が定義されます。デプロイ時には、最も高いスコアを持つ候補が選択されます。

3. 学習目的関数：

   • ポリシー損失： ポリシーは、シミュレーションされたロールアウトを通じてポリシー勾配を用いて学習されます。目的関数は、タスク固有の効用（例：BEDにおける期待情報利得、BOにおける終端の最良値）を最大化します。著者らは、分散を抑えるために、完全な長期ホライゾンの価値関数ではなく、ローカルなクレジット割り当て（1ステップの増分）を使用しています。
   • バックボーンの適応（オプション）： 特定のダウンストリームタスクの分布にバックボーンを適合させるために、補助的な教師あり予測損失（$L_{pred}$）がバックボーンのパラメータに適用されます。
   • 分離（Decoupling）： 重要な設計上の選択として、ポリシー勾配はバックボーンには伝播されません。バックボーンは補助的な予測損失によってのみ更新され、ポリシーヘッドはポリシー勾配によって更新されます。これにより、高分散な強化学習の信号が、事前学習された表現を不安定化させることを防ぎます。

理論的基礎

本論文は、統一的な決定論的視点を提供し、ベイズ最適なポリシー $\pi^*$ はタスクに関連する事後分布 $p(w|D_{t-1})$ の関数であることを示しています。事後分布は履歴の関数であるため、デコーダを介して事後分布の復元を可能にするあらゆる表現 $e_t$ は、有効な信念表現となります。POLARは、事前学習されたTFMが、最適な行動を決定するのに十分なほど豊かでありながら、ポリシーヘッドへのマッピングを容易にするほど構造化された表現を提供しており、生の履歴（複雑すぎる）や明示的な事後分布（近似誤差に弱い）の両極端を回避できると仮定しています。

主な貢献

1. 統一的視点： 著者らは、適応型データ収集を、信念表現の上でのポリシー学習として定式化し、生の履歴ベースのアプローチと事後分布ベースのアプローチの間の溝を埋めました。
2. 分離されたアーキテクチャ： 事前学習済みの基盤モデルを固定（またはファインチューニング）された信念エンコーダとして使用し、別途ポリシーヘッドを訓練するという、シンプルなアーキテクチャを提案しました。これにより、表現学習と意思決定の分離を実現しています。
3. サンプル効率： 実証結果により、POLARは多様なタスク（BED、BO、AL）において、既存の償還化（amortised）手法と比較して大幅に少ない学習サンプル（最大100倍少ない）で、最先端の性能を達成することが示されました。

実験結果

著者らは、主に以下の3つのベンチマークでPOLARを評価しています。

1. ベイズ実験計画法 (BED):

   • タスク： 位置特定（2Dおよび5D）および定数弾力性（CES）。
   • 結果： POLARは、既存の最先端の償還化手法（DAD、RL-BOED、ALINE）を大幅に上回りました。5Dの位置特定タスクにおいて、POLARはRL-BOEDおよびDADよりも 10,000倍少ない学習サンプル でピーク性能を上回り、ALINEに対しては 100倍少ないサンプル で上回りました。
   • アブレーション解析：
     • 固定 vs. ファインチューニングされたバックボーン： 補助的な予測損失を用いてバックボーンをファインチューニングすることは、固定する場合よりも高い性能をもたらしましたが、固定バージョンでも競争力のある性能を示しました。
     • 結合 vs. 分離された勾配： ポリシー勾配をバックボーンの更新に許容する（結合された）場合、性能が著しく低下し、分離の必要性が確認されました。
     • 事前学習済み vs. スクラッチ： アーキテクチャをゼロから学習すると、収束が大幅に遅くなり、性能も低くなったため、事前学習済み基盤モデルの価値が浮き彫りになりました。

2. ベイズ最適化 (ハイパーパラメータ最適化):

   • ベンチマーク： HPO-B (Hyperparameter Optimisation Benchmark)。
   • 結果： POLARは、6つの探索空間において、最小の平均後悔（regret）と最高の平均ランクを達成し、古典的なGP、Meta-GP、およびNAPのようなエンドツーエンドの償還化手法を上回りました。事前学習されたTFMバックボーンを使用することで、ゼロから表現を学習するよりも強力な性能が得られることが示されました。

3. 分子最適化:

   • ベンチマーク： DOCKSTRING（高次元の分子ドッキング）。
   • 結果： 512次元の離散空間において、POLARはTanimotoカーネルを用いたGPやランダムフォレスト、および手設計された獲得関数を用いたTabICLのベースラインと比較して、最小の後悔を達成しました。これは、信念表現から直接獲得ルールを学習することが、固定された獲得戦略よりも優れていることを裏付けています。

意義と主張

本論文は、POLARが償還化されたデータ収集の スケーラビリティと効率性 における重要な一歩であることを主張しています。表現と決定を共同で学習するパラダイムから、事前学習済みの予測モデル を用いた信念表現に基づく行動へとシフトすることで、POLARは直接的なポリシー学習に伴うサンプル効率の低さと学習の不安定性、および事後分布ベースの手法に伴う近似誤差を軽減します。

著者らは、このアプローチは新しい推定器や学習アルゴリズムを設計することではなく、事前学習済みの予測モデルを信念エンコーダとして再利用する ことであると強調しています。これにより、小さなポリシーヘッドを標準的なポリシー勾配を用いて訓練することが可能となり、劇的に低い計算コストで優れた性能を実現しています。このフレームワークは、使用するタスク固有の効用が異なるだけで、BED、BO、ALに適用可能な統一されたソリューションとして提示されています。