디자인과 추론을 동시에 학습하는 JADAI 프레임워크

본 논문은 파라미터 추정 문제에서 실험 설계 변수를 능동적으로 최적화함으로써 정보 획득을 극대화하고자 하는 목표를 갖는다. 이를 위해 저자들은 JADAI(Jointly Amortizing Adaptive Design and Bayesian Inference)라는 새로운 프레임워크를 제안한다. JADAI는 세 개의 주요 구성 요소—설계 정책 네트워크, 히스토리 네트워크, 그리고 확산 기반 포스터리어 추정기—를 하나의 통합 학습 목표 아래 결합한다. 설계 정책은 현재까지 관측된 데이터와 설계 이력을 히스토리 네트워크에 전달받아, 다음 실험에서 사용할 설계 변수 ξ를 선택한다. 히스토리 네트워크는 시뮬레이션으로부터 생성된 (θ, x) 쌍을 시계열 형태로 인코딩함으로써, 정책과 추정기가 과거 정보를 충분히 활용하도록 돕는다. 가장 핵심적인 기술적 기여는 포스터리어 추정기에 확산 모델을 도입한 것이다. 기존 시뮬레이션 기반 추론(SBI)에서는 흐름 모델, 변분 오토인코더, 혹은 최근의 일관성 모델 등을 사용했지만, 이러한 방법들은 고차원·다중모달 분포를 정확히 표현하는 데 한계가 있었다. 확산 모델은 단계별 노이즈 주입과 역전파 과정을 통해 복잡한 분포를 점진적으로 학습하므로, JADAI는 10차원 이상 파라미터와 다중 피크를 갖는 사후 분포에서도 높은 정확도를 달성한다. JADAI의 학습 목표는 “증분 포스터리어 오류 감소”를 정량화한 일반화 손실 L이다. 구체적으로, t번째 실험 후의 사후 pₜ(θ|x₁:ₜ, ξ₁:ₜ)와 t‑1번째 실험 후의 사후 pₜ₋₁(θ|x₁:ₜ₋₁, ξ₁:ₜ₋₁) 사이의 KL 발산 혹은 MMD 차이를 계산하고, 이를 전체 실험 시퀀스에 걸쳐 누적한다. 이 손실은 기대 정보 이득(EIG)과 동등한 의미를 가지며, 설계 변수 ξ가 실제로 사후 불확실성을 얼마나 감소시키는지를 직접적인 학습 신호로 제공한다. 설계 변수는 연속적인 파라미터이므로, 정책 네트워크는 재파라미터화 기법이나 스코어 함수 추정(Score‑Function Estimator)을 이용해 미분 가능하게 만든다. 따라서 설계와 추정이 서로의 그래디언트를 공유하며 공동 최적화가 가능해진다. 실험에서는 표준 적응 설계 벤치마크(예: 1D Gaussian, 2D 로지스틱 회귀, 고차원 물리 시뮬레이터)와 다중모달 시뮬레이터(예: 파동 전파, 전염병 확산 모델)를 대상으로 기존 정책 기반 BAD(Deep Bayesian Adaptive Design)와 비동시 SBI(Sequential Neural Posterior Estimation)와 비교하였다. 결과는 JADAI가 동일한 샘플 수에서 더 높은 최종 포스터리어 정확도와 낮은 평균 사후 오류를 기록했으며, 특히 다중모달 상황에서 기존 방법이 수렴하지 못하거나 편향된 추정을 보이는 반면, JADAI는 안정적으로 다중 피크를 복원했다. 또한, 설계 정책이 학습 과정에서 점진적으로 정보가 풍부한 설계로 수렴하는 모습을 시각화했으며, 이는 손실 함수가 설계와 추정 사이의 상호작용을 효과적으로 캡처함을 증명한다. 한계점으로는 확산 모델 학습에 필요한 대규모 시뮬레이션 데이터와 GPU 메모리 요구량이 크다는 점, 그리고 현재 구현이 연속형 설계 변수에 최적화돼 있어 이산형 설계(예: 실험 장비의 스위치 선택)에는 직접 적용하기 어려울 수 있다는 점을 언급한다. 향후 연구 방향으로는 메타‑학습을 통한 시뮬레이터 독립성 강화, 이산 설계 공간에 대한 정책 네트워크 확장, 그리고 강화학습 기반 탐색 전략과의 결합을 제시한다. 전반적으로 JADAI는 설계와 추정을 별도로 다루던 기존 패러다임을 깨고, 두 과정을 공동으로 최적화함으로써 샘플 효율성과 추정 정확도 모두를 크게 향상시킬 수 있음을 실증적으로 보여준다.