Identifying Informative Environments for Cognition Parameter Inference via Bayesian Experimental Design

技術要約：ベイズ実験計画法による認知パラメータ推論のための情報量豊かな環境の特定

問題定式化
計算論的認知モデリングは、観測された逐次的な意思決定行動から、潜在的な認知メカニズム（例：計画ホライゾン、ワーキングメモリ容量）を推論することを目指している。ベイズ逆計画法は、この推論のための原理的な枠組みを提供するが、その有効性は実験環境に決定的に依存する。既存のアプローチは通常、環境を固定されたもの、あるいはランダムにサンプリングされるものとして扱っており、どの特定の環境がパラメータ回復に対して最も情報量の多いデータをもたらすかという問題に対処できていない。

本論文は、実験環境の選択を**ベイズ実験計画（BED）問題として定式化する。目標は、データ収集前に、潜在的な認知パラメータ $\theta$ に関する期待情報利得（EIG）**を最大化する環境 $\xi^*$ を特定することである。EIGは、環境 $\xi$ における行動軌跡 $D$ を観察した後の、$\theta$ に関する事後分布の不確実性の減少の期待値として定義される：
$$\xi^* = \arg \max_{\xi} \text{EIG}(\xi) = \arg \max_{\xi} \left( H(\theta) - \mathbb{E}_{p(D|\xi)}[H(\theta | D, \xi)] \right)$$
核心となる課題は、EIGの評価には多くのシミュレーションされたデータセットに対する繰り返しのベイズ事後推論が必要であり、そのため、網羅的な最適化は計算量的に極めて困難である点にある。

手法
著者らは、ベイズ逆計画法とスケーラブルな実験計画技術を組み合わせたパイプラインを提案している：

1. 生成的行動シミュレータ： 著者らは、潜在的な認知パラメータ $\theta$（計画ホライゾン $H$ およびワーキングメモリ $K$）を持つパラメータ化されたプランナーを用いてエージェントをモデル化している。エージェントは、$\theta$ に基づいて内部的な認知状態 $b_t$ を更新しながら、マルコフ決定過程（MDP）環境 $\xi$ と相互作用する。行動は、アクション値 $Q_\theta$ 上のソフトマックス・ポリシーを介してサンプリングされ、軌跡データ $D \sim p(D|\theta, \xi)$ を生成する。
2. 厳密なモンテカルロBED（MC-BED）： グラウンドトゥルースのベンチマークを確立するために、著者らはモンテカルロ・シミュレーションを用いてEIGを計算する。これには、パラメータと環境をサンプリングし、軌跡をシミュレートし、各シミュレートされたデータセットに対して厳密なベイズ事後推論を実行して、事後分布と事前分布の間のKLダイバージェンスを計算するプロセスが含まれる。
3. アモルタイズド・ベイズ実験計画（ABED）： MC-BEDの計算コストを克服するため、著者らはアモルタイズド（償却型）フレームワークを導入する：
   • ニューラル事後分布サロゲート： 厳密なベイズ事後分布を近似するニューラルネットワーク $q_\phi(\theta | D, \xi)$ を訓練する。訓練データは $\theta, \xi, D$ をサンプリングすることで生成され、厳密な事後分布が教師信号として機能する。
   • マルチタスク学習： サロゲートは、事後分布の蒸留（厳密な事後分布への適合）、教師あり分類（真の $\theta$ の予測）、および補助的な回帰ヘッド（厳密な事後分布と事前分布の間のKLダイバージェンス、すなわちデザイン・ユーティリティの予測）を組み合わせたマルチタスク目的関数を用いて訓練される。
   • 効率的なデザイン評価： 一度訓練されると、サロゲートは高コストな厳密推論に代わる。デザインの目的関数は、ネットワークの単一のフォワードパスによって近似され、これにより候補となる環境の迅速なランキングが可能になる。

主な貢献

• 定式化： 認知的なプランニング環境の設計を、固定されたタスク内の実験条件ではなく、環境自体をデザイン変数とするBED問題として定式化した。
• ベンチマーク： 認知パラメータ推論のための環境の情報量を評価するための、厳密なモンテカルロBEDベンチマークを確立した。
• アモルタイズド・フレームワーク： 厳密なMC-BEDベンチマークと密接に一致する環境ランキングを実現しつつ、繰り返しの推論をニューラルサロゲートに置き換えることで計算コストを大幅に削減する、アモルタイズドBEDフレームワークを導入した。
• 体系的な分析： 提案フレームワークをMouselab-MDPプロセス・トレーシング・パラダイム上で評価し、厳密推論、MC-BED、およびABEDを、ヒューリスティックおよび教師あり回帰ベースラインと比較した。

実験結果
実験は、深さ ($d$) と分岐係数 ($b$) が異なる木構造のMouselab-MDP環境に対して行われた。

• 回復可能性 vs 情報利得： 結果は、単一の環境が常に最適であるわけではないことを示している。有限サンプルでの事後分布の回復（例：結合推論における $d_5b_2$）を最大化する環境が、必ずしも期待情報利得（例：$d_6b_2$）を最大化する環境と一致するわけではない。
• アモルタイゼーションの性能： アモルタイズド・サロゲート（ABED）は、厳密なMC-BEDの環境ランキングを密接に再現した。事後分布の回復精度には変動があるものの、予測されたデザイン・ユーティリティ（KLダイバージェンス）は厳密な目的関数と一致しており、最も情報量の多い環境を正しく特定できる。
• ベースラインとの比較： ABEDは、構造的ヒューリスティック（例：「最大の木」）や、事後推論を行わずにデザインスコアを予測する教師あり回帰器よりも優れた性能を示した。ABEDは、統一されたフレームワーク内で認知パラメータ推論と環境評価を同時にサポートできる唯一の手法である。
• トレードオフ： 本研究は、期待情報利得、事後分布の回復可能性、および情報効率（クリックあたりのKL）の間のトレードオフを明らかにしている。複雑な戦略に対しては一般に深い環境が情報利得を最大化するが、単純な戦略に対しては浅い環境の方が高い効率や優れた回復可能性を提供する場合がある。

意義と主張
本論文は、ベイズ的パラメータ推論のための、情報量の多い認知実験を設計するための原理的なフレームワークを提供することを主張している。実験環境をデザイン変数として扱うことで、本研究は固定された評価設定を超え、潜在的な認知メカニズムに関する情報を最大化する環境を能動的に選択することへと進展させている。

著者らは、最適な環境は特定の潜在的な認知状態および推論の目的（回復か情報利得か）に依存するため、適応的な環境選択が必要であることを強調している。提案されたABBEDフレームワークは、進化する事後分布に基づいた逐次的かつ適応的な実験選択を可能にするために必要な構成要素（結合事後分布推定およびデザイン値予測）を提供し、より効率的な認知アセスメントへの道を開くものである。本研究は、実験環境自体が、認知モデルの推論プロセスの一部として最適化されるべきであることを主張している。