Probabilistic Retrofitting of Learned Simulators

Motivation

• Known: 딥러닝 기반 PDE surrogate 모델들은 대부분 MSE/MAE를 최소화하는 deterministic 방식으로 학습되며, 최근에는 diffusion이나 flow-based generative model을 활용한 확률적 프레임워크가 chaotic하고 불확실한 시스템을 더 잘 모델링한다는 것이 알려져 있다.
• Gap: Diffusion/flow 기반 확률적 모델은 성능이 우수하지만 처음부터 학습하려면 특수하고 자원 집약적인 학습 과정과 느린 반복적 추론이 필요해, 이미 상당한 자원이 투입된 고성능 deterministic backbone을 재활용할 수 없다는 문제가 있다.
• Why: PDE 모델링 분야가 대규모 파라미터와 데이터로 일반화 성능을 추구하는 foundation model 패러다임으로 전환되는 시점에서, 확률적 foundation model을 처음부터 학습하는 것은 경제적·환경적으로 지속 불가능하므로 기존 deterministic backbone을 재활용하는 방법은 시의적절하고 실용적인 대안이다.
• Approach: 사전학습된 deterministic 모델에 conditional layer normalisation 기반의 noise injection 메커니즘을 추가하고, CRPS(특히 unbiased 'fair' CRPS estimator)를 손실함수로 사용해 소량의 추가 학습만으로 확률적 모델로 retrofitting하는 방법을 제안한다.

Achievement

1. Architecture-agnostic 통합 레시피: 세 가지 서로 다른 아키텍처(pixel-space, latent-space 포함)와 다섯 개의 물리 시스템에 걸쳐 동일한 retrofitting 레시피가 효과적임을 입증했으며, deterministic backbone에 최소한의 코드 수정만 요구한다.
2. 효율성과 성능: single-system 사전학습 모델에서 rollout CRPS를 20-54% 감소, VRMSE를 최대 30% 개선했으며, 여러 시스템에 학습된 PDE foundation model을 대상 데이터셋에 retrofitting했을 때도 deterministic fine-tuning 대비 CRPS 최대 40%, VRMSE 최대 15% 개선을 달성했다.
3. 추론 시점 유연성: retrofitting된 모델이 test-time scaling 특성을 보여, 앙상블 멤버 수를 늘릴수록 예측 성능이 단조적으로 향상되어 재학습 없이 compute-accuracy trade-off를 조정할 수 있음을 보였다.
4. Ablation과 통찰: CRPS retrofitting의 학습 역학에 대한 광범위한 ablation을 수행하여 안정적인 확률적 적응을 위한 실용적 하이퍼파라미터 튜닝 권고사항을 제공했다.

How

• 기존 deterministic 모델 Mθ은 과거 상태 history ht로부터 다음 상태(또는 잔차)를 단일 point estimate로 예측하고 MSE/MAE로 최적화됨
• 제안 방법은 noise ϵ을 MLP로 projection하여 backbone을 modulate하는 stochastic branch를 도입해 확률적 모델 M̃θ로 변형
• 학습 목적함수로 strictly proper scoring rule인 CRPS, 구체적으로 유한 샘플 편향을 제거한 unbiased 'fair' CRPS estimator(식 1)를 사용해 state vector의 각 컴포넌트별로 독립적으로 적용 후 평균
• 추론 시 여러 개의 ϵ 벡터를 샘플링하여 단일 input history에 대해 sharp하고 다양한 예측들로 구성된 ensemble을 생성
• 단일 시스템에 대해 처음부터 학습된 모델과, 여러 시스템에 대해 학습된 PDE foundation model 두 가지 스케일의 사전학습 규모에서 검증

Originality

• CRPS 최적화가 사전학습에 사용된 MAE/MSE 손실과 밀접히 연관되어 있어, diffusion/flow 기반 모델의 denoising task로의 전환보다 정렬(misalignment)이 적어 사전학습된 feature가 더 강한 초기화를 제공한다는 가설을 세우고 이를 검증함
• weather forecasting 분야에서 성공적이었던 전략을 PDE 모델링의 다양한 아키텍처와 물리 시스템으로 확장 적용하여 그 전이 가능성(transferability)을 입증함
• 단일 시스템 모델뿐 아니라 multi-system PDE foundation model에도 동일한 retrofitting 방법론이 유효함을 보여, 기존 deterministic 자원을 재활용하는 실용적 대안을 최초로 체계적으로 검증함

Limitation & Further Study

• CRPS는 marginal distribution에 적용되는 스코어이기 때문에 변수 간 joint dependency나 공간적 상관관계를 완전히 포착하지 못할 가능성이 있어 추가적인 검증이 필요함
• conditional layer normalisation 기반 noise injection 메커니즘의 설계가 특정 backbone 구조(Transformer 계열)에 편향되어 있을 수 있어, 완전히 다른 유형의 아키텍처(e.g., GNN 기반)에 대한 일반화 검증이 더 필요함
• 다섯 개의 물리 시스템에 대한 검증에 그쳐, 더 다양하고 복잡한 실제 물리 시스템(e.g., 3차원 난류, 다중 스케일 결합 시스템)에서의 확장성 검증이 후속 연구로 요구됨
• ablation을 통한 하이퍼파라미터 권고사항이 제시되었으나, retrofitting 과정의 이론적 수렴 보장이나 학습 안정성에 대한 심층적 분석은 제한적임

Evaluation

총평: 사전학습된 deterministic PDE 모델을 CRPS 기반의 경량 retrofitting으로 확률적 모델로 전환하는 실용적이고 architecture-agnostic한 방법을 제안하며, 다양한 스케일과 아키텍처에서 강력한 실증 결과를 보여주는 시의적절하고 의미 있는 연구이다.