Scalable Simulation-Based Model Inference with Test-Time Complexity Control

Motivation

• Known: Bayesian model selection은 marginal likelihood(evidence)를 통해 Occam's razor 방식의 model complexity penalty를 부여하는 원칙적인 프레임워크이며, simulation-based inference(SBI)는 likelihood-free 상황에서 amortized posterior 추정을 가능하게 한다.
• Gap: 기존 방법들은 각 candidate model마다 별도로 MCMC 등 Bayesian inference를 수행해야 해 모델 수가 많아지면 계산이 불가능해지며, amortized model-selection 방법들도 학습 시점에 고정된 model prior 또는 complexity penalty를 하드코딩하여 데이터를 보기 전에 parsimony 가정을 미리 결정해야 하는 한계가 있다.
• Why: 과학적 발견 과정에서 연구자들은 흔히 수많은 plausible model family 중에서 데이터에 가장 부합하는 모델을 선택해야 하는데, 이 갭을 해결하지 못하면 combinatorial하게 증가하는 model space에서 신뢰할 수 있는 model selection과 사후 재조정이 불가능해 실제 과학 응용(dMRI, symbolic regression 등)에 확장하기 어렵다.
• Approach: PRISM은 transformer 기반 encoder-decoder를 사용해 model complexity를 조절하는 hyperparameter λ에 조건화된 hierarchical model prior p(M|λ) 하에서 model posterior와 parameter posterior를 동시에 근사 학습하며, 이를 통해 재학습 없이 test-time에 λ를 조정해 parsimony trade-off를 탐색할 수 있다.

Achievement

Figure 4. PRISM on symbolic regression task.

1. 대규모 model family로의 확장성: synthetic symbolic regression task에서 수십억(billions) 개에 달하는 combinatorial model instantiation을 갖는 model family에 대해 PRISM이 효과적으로 amortize됨을 보였다.
2. Test-time complexity 조절: 별도의 재학습 없이 λ 값을 조정함으로써 symbolic regression에서 항의 개수(term 수) 등 모델 복잡도를 유연하게 탐색하고 선택할 수 있음을 입증했다.
3. 실제 과학 응용에서의 model selection: diffusion MRI(dMRI) 데이터에 대한 biophysical multi-compartment model(예: Ball-and-Sticks)들 사이에서 synthetic 및 in-vivo neuroimaging 데이터 모두에 대해 model selection을 수행하고, 기존 SBI pipeline(Manzano-Patron et al., 2025)보다 fiber orientation estimation 등에서 우수한 성능을 보였다.

How

Figure 3. Architecture overview. PRISM is based on two parallel

• 두 개의 병렬 구조(discrete model 추론용, continuous parameter 추론용)를 가진 transformer 기반 encoder-decoder 아키텍처(Fig. 3)를 설계.
• 학습 단계에서 model family M으로부터 hierarchical prior p(M|λ)에 따라 샘플링하고, loss L = L(M, p(M|x,λ)) + L(θ_M, p(θ|M,λ,x))를 최적화하여 model posterior와 조건부 parameter posterior를 함께 근사.
• Schröder & Macke (2024)의 이전 model-inference 접근법에서 사용된 Gaussian mixture density network 기반 parameter posterior 및 비활성 차원에 대한 analytic marginalization 필요성 등 제약적 설계를 제거.
• 추론 단계에서는 λ를 조정하여 원하는 parsimony 수준을 설정한 후, combinatorial model space 내에서 모델을 선택·탐색하고 해당 모델의 파라미터 posterior를 추론하여 downstream 분석에 활용.
• Symbolic regression 벤치마크와 dMRI 데이터셋(synthetic 및 in-vivo)에서 실험을 수행해 PRISM의 확장성과 실용성을 검증.

Originality

• 기존 amortized model-selection 방법들이 model prior/complexity penalty를 학습 시점에 고정하는 것과 달리, λ에 조건화된 네트워크를 통해 test-time에 모델 복잡도를 자유롭게 조절 가능하게 한 최초의 시도.
• discrete model structure와 continuous parameter를 하나의 encoder-decoder에서 joint posterior로 동시에 추론하는 통합 프레임워크 제안.
• Gaussian mixture density network 기반의 이전 접근법(Schröder & Macke, 2024)이 가진 analytic marginalization 등 제약을 제거하여 더 유연하고 확장 가능한 구조로 재설계.
• 수십억 개 model instantiation 규모의 combinatorial model space에 대한 amortized inference를 실증적으로 시연한 점.

Limitation & Further Study

• 논문 발췌 내용상 λ의 구체적 파라미터화 방식이나 hierarchical prior p(M|λ)의 설계가 다양한 도메인에 얼마나 일반화되는지에 대한 상세한 이론적 분석이 제한적으로 보인다.
• dMRI 응용에서 고려한 multi-compartment model 종류가 제한적일 수 있으며, 더 복잡하거나 이질적인 model family(서로 다른 파라미터 차원, 시뮬레이터 구조를 가진 경우)로의 확장 가능성에 대한 추가 검증이 필요하다.
• Marginal likelihood 기반 evidence의 근사 정확도, 특히 model posterior 추정의 calibration에 대한 심층적 검증(다양한 misspecification 시나리오 포함)이 후속 연구로 요구된다.
• 대규모 combinatorial model space에서 학습 비용(simulation budget, 학습 시간) 및 스케일링 법칙에 대한 정량적 분석이 추가로 필요해 보인다.

Evaluation

총평: PRISM은 model selection과 parameter inference를 하나의 amortized framework로 통합하면서 test-time에 model complexity를 조절할 수 있게 한 실용적이고 확장 가능한 기여로, symbolic regression과 dMRI라는 서로 다른 도메인에서의 검증을 통해 그 실효성을 잘 보여준다.