Identifying Informative Environments for Cognition Parameter Inference via Bayesian Experimental Design

기술 요약: 베이지안 실험 설계(Bayesian Experimental Design)를 통한 인지 파라미터 추론을 위한 정보 집약적 환경 식별

문제 정의
계산 인지 모델링은 관찰된 순차적 의사결정 행동으로부터 잠재적인 인지 메커니즘(예: 계획 수평선, 작업 기억 용량)을 추론하는 것을 목표로 한다. 베이지안 역계획(Bayesian inverse planning)은 이러한 추론을 위한 원칙적인 프레임워크를 제공하지만, 그 효용성은 실험 환경에 결정적으로 의존한다. 기존의 접근 방식은 일반적으로 환경을 고정되어 있거나 무작위로 샘플링되는 것으로 취급하며, 어떤 특정 환경이 파라미터 회복(parameter recovery)에 가장 유익한 데이터를 제공하는지에 대해서는 다루지 못한다.

본 논문은 실험 환경의 선택을 베이지안 실험 설계(Bayesian Experimental Design, BED) 문제로 정식화한다. 목표는 데이터를 수집하기 전, 잠재적 인지 파라미터 $\theta$에 대한 **기대 정보 이득(Expected Information Gain, 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. 생성적 행동 시뮬레이터 (Generative Behavioral Simulator): 저자들은 잠재적 인지 파라미터 $\theta$(계획 수평선 $H$ 및 작업 기억 $K$)를 가진 매개변수화된 플래너를 사용하여 에이전트를 모델링한다. 에이전트는 $\theta$에 기반하여 내부 인지 상태 $b_t$를 업데이트하며 마르코프 결정 과정(MDP) 환경 $\xi$와 상호작용한다. 행동은 액션 가치 $Q_\theta$에 대한 소프트맥스 정책(softmax policy)을 통해 샘플링되며, 이를 통해 궤적 데이터 $D \sim p(D|\theta, \xi)$가 생성된다.
2. 정밀 몬테카를로 BED (Exact Monte Carlo BED, MC-BED): 그라운드 트루스 벤치마크를 구축하기 위해, 저자들은 몬테카를로 시뮬레이션을 사용하여 EIG를 계산한다. 이는 파라미터와 환경을 샘플링하고, 궤적을 시뮬레이션하며, 모든 시뮬레이션 데이터셋에 대해 정밀한 베이지안 사후 추론을 수행하여 사후 분포와 사전 분포 사이의 KL 발산을 계산하는 과정을 포함한다.
3. 아모티즈드 베이지안 실험 설계 (Amortized Bayesian Experimental Design, ABED): MC-BED의 계산 비용을 극복하기 위해, 저자들은 아모티즈드(amortized) 프레임워크를 도입한다:
   • 신경 사후 분포 대리 모델 (Neural Posterior Surrogate): 정밀한 베이지안 사후 분포를 근사하는 신경망 $q_\phi(\theta | D, \xi)$를 학습시킨다. 학습 데이터는 $\theta, \xi, D$를 샘플링하여 생성되며, 정밀한 사후 분포가 감독 신호 역할을 한다.
   • 멀티태스크 학습 (Multi-Task Learning): 이 대리 모델은 사후 분포 증류(정밀 사후 분포와 일치시키기), 지도 학습 기반 분류(그라운드 트루스 $\theta$ 예측), 그리고 정밀 사후 분포와 사전 분포 간의 KL 발산(설계 유용성)을 예측하는 보조 회귀 헤드를 결합한 멀티태스크 목적 함수로 학습된다.
   • 효율적 설계 평가 (Efficient Design Evaluation): 일단 학습되면, 대리 모델은 비용이 많이 드는 정밀 추론을 대체한다. 설계 목적 함수는 네트워크의 단일 순전파(forward pass)를 통해 근사되며, 이를 통해 후보 환경들을 빠르게 순위화할 수 있다.

주요 기여

• 정식화: 인지 계획 환경 설계를 고정된 태스크 내의 실험 조건이 아닌, 환경 자체가 설계 변수인 BED 문제로 정식화하였다.
• 벤치마킹: 인지 파라미터 추론을 위한 환경의 정보 집약성을 평가하기 위해 정밀한 몬테카를로 BED 벤치마크를 구축하였다.
• 아모티즈드 프레임워크: 반복적인 추론을 신경 대리 모델로 대체함으로써, 정밀한 MC-BED 벤치마크와 밀접하게 일치하는 환경 순위를 달성하면서도 계산 비용을 크게 줄이는 아모티즈드 BED 프레임워크를 도입하였다.
• 체계적 분석: 본 프레임워크를 Mouselab-MDP 프로세스 트레이싱 패러다임에 적용하여, 정밀 추론, MC-BED, ABED를 휴리스틱 및 지도 학습 기반 베이스라인과 체계적으로 비교하였다.

실험 결과
실험은 다양한 깊이($d$)와 분기 계수($b$)를 가진 트리 구조의 Mouselab-MDP 환경에서 수행되었다.

• 회복 가능성 vs. 정보 이득 (Recoverability vs. Information Gain): 결과에 따르면 단일 환경이 균일하게 최적은 아니다. 유한 샘플 사후 회복 가능성을 최대화하는 환경(예: 결합 추론을 위한 $d_5b_2$)이 기대 정보 이득을 최대화하는 환경(예: $d_6b_2$)과 항상 일치하는 것은 아니다.
• 아모티즈드 성능 (Amortization Performance): 아모티즈드 대리 모델(ABED)은 정밀한 MC-BED의 환경 순위를 밀접하게 재현한다. 사후 회복 정확도는 다양할 수 있으나, 예측된 설계 유용성(KL 발산)은 정밀한 목적 함수와 일치하여 가장 정보 집약적인 환경을 올바르게 식별할 수 있게 한다.
• 베이스라인과의 비교: ABED는 구조적 휴리스틱(예: "가장 큰 트리") 및 사후 추론 없이 설계 점수를 예측하는 지도 학습 회귀 모델보다 우수한 성능을 보인다. ABED는 인지 파라미터 추론과 환경 평가를 단일 프레임워크 내에서 동시에 지원하는 유일한 방법이다.
• 트레이드오프 (Trade-offs): 연구는 기대 정보 이득, 사후 회복 가능성, 그리고 정보 효율성(클릭당 KL) 사이의 트레이드오프를 밝혀낸다. 복잡한 전략을 위해서는 일반적으로 더 깊은 환경이 정보 이득을 최대화하는 반면, 단순한 전략을 위해서는 더 얕은 환경이 더 높은 효율성이나 더 나은 회복력을 제공할 수 있다.

의의 및 주장
본 논문은 베이지안 파라미터 추론을 위한 정보 집약적 인지 실험을 설계하기 위한 원칙적인 프레임워크를 제공한다고 주장한다. 실험 환경을 설계 변수로 취급함으로써, 본 연구는 고정된 평가 설정을 넘어 잠재적 인지 메커니즘에 대한 정보를 최대화하는 환경을 능동적으로 선택하는 단계로 나아간다.

저자들은 **적응형 환경 선택(adaptive environment selection)**의 필요성을 강조한다. 최적의 환경은 특정 잠재적 인지 상태와 추론 목적(회복 vs. 정보 이득)에 따라 달라지므로, 고정된 환경은 차선책이 될 수밖에 없다. 제안된 ABED 프레임워크는 공동 사후 분포(joint posterior)에 기반하여 진화하는 사후 분포에 따라 순차적이고 적응적인 실험 선택을 가능하게 하는 필수 구성 요소(공동 사후 추정 및 설계 가치 예측)를 제공하며, 이는 확장 가능하고 참여자 효율적인 인지 평가를 향한 경로를 제시한다. 본 연구는 실험 환경 자체가 인지 모델 추론 과정의 일부로서 최적화되어야 한다고 주장한다.