기술 요약: 시계열 PDE를 위한 파운데이션 모델의 증류(Distillation)

1. 문제 정의

편미분 방정식(PDE)은 유체 역학 등 시공간적 물리 과정을 모델링하는 데 필수적입니다. 신경망이 PDE를 위한 효율적인 대리 모델(surrogate models)로 부상하고 있지만, 다음과 같은 결정적인 트레이드오프에 직면해 있습니다:

• 파운데이션 모델 (Foundation Models, FMs): 다양한 물리 데이터셋(예: Poseidon, Walrus)으로 사전 학습된 대규모 모델은 제한된 데이터로도 다운스트림 작업에 대한 강력한 일반화 성능을 보여줍니다. 그러나 이들은 대개 특정 유도 편향(예: 등변성, equivariance)이 부족한 거대 트랜스포머 아키텍처에 의존하며, 계산 비용이 높고 추론 시간과 메모리 소모가 큽니다. 이는 제어 및 계획과 같은 실시간 응용 분야에서의 활용을 제한합니다.
• 소형 대리 모델 (Compact Surrogates): 푸리에 신경 연산자(Fourier Neural Operators)와 같은 작은 모델들은 효율적이지만, 높은 정확도를 달-성하기 위해 일반적으로 많은 양의 작업 특화 학습 데이터가 필요하며, 실험적 또는 부분적으로 관측된 시스템에서는 이러한 데이터를 확보하기 어려운 경우가 많습니다.

핵심 과제는, 타겟 도메인의 초기 조건 분포($p_0$)에 접근할 수 없거나(알려져 있지 않거나 샘플링 비용이 매우 큰 경우) 해당 분포를 직접 샘플링할 필요 없이, 거대하고 사전 학습된 PDE 파운데이션 모델의 일반화 능력을 작고 효율적인 학생(student) 모델로 전이하는 것입니다.

2. 방법론: 교사 롤아웃 확장 (Teacher Rollout Extension, TREX)

저자들은 거대 파운데이션 모델과 효율적인 다운스트림 대리 모델 사이의 간극을 메우기 위해 설계된 지식 증류 프레임워크인 **TREX(Teacher Rollout Extension)**를 제안합니다.

2.1 핵심 개념

TREX는 알려진 분포 $p_0$로부터 새로운 초기 조건을 샘플링해야 하는 표준 증류의 한계를 해결합니다. 많은 다운스트림 시나리오(예: 실험 데이터)에서 $p_0$는 알 수 없는 값입니다. 대신, TREX는 **교사 롤아웃(teacher rollouts)**을 통해 제한된 가용 정답(ground-truth) 궤적을 긴 합성 궤적으로 확장하여 보강합니다.

2.2 TREX 절차

1. 미세 조정 (Fine-tuning): 먼저 사전 학습된 파운데이션 모델(교사, $T$)을 사용 가능한 소수의 다운스트림 궤적 데이터셋에 대해 미세 조정합니다.
2. 롤아웃 생성 (Rollout Generation): 가용 가능한 소수의 정답 궤적의 초기 상태로부터 시작하여, 미세 조정된 교사를 자기회귀적(autoregressively)으로 롤아웃하여 긴 합성 궤적을 생성합니다.
   • 점유 측도 (Occupancy Measure): 이 과정은 단순히 짧은 관측된 궤적만을 샘플링하는 것이 아니라, 교사가 유도한 롤아웃 분포(자기회귀적 예측 중에 마주치는 상태들)로부터 샘플링합니다.
   • 노이즈 주입 (Noise Injection): 상태 공간의 커버리지를 넓히고 학생 모델에게 섭동(perturbed)된 상태로부터 시스템이 어떻게 진화하는지 가르치기 위해, 롤아웃 중에 주기적으로 가우시안 노이즈를 주입합니다. 그 후 교사는 이 섭동된 상태를 바탕으로 다음 상태를 예측합니다.
3. 학생 학습 (Student Training): 컴팩트한 학생 모델($S$)은 다음의 혼합 데이터로 학습됩니다:
   • 원래의 정답 궤적 전환(transitions).
   • 교사가 생성한 전환(노이즈가 없는 세그먼트와 노이즈가 포함된 세그먼트 모두 포함).
   • 학생 손실 함수는 정답 감독($L_{GT}$)과 교사 증류 손실($L_{TREX}$)을 결합합니다.

2.3 아키텍처 유도 편향 (Architectural Inductive Biases)

일반적으로 거대하고 유연한 트랜스포머인 교사와 달리, 학생 아키텍처는 작업 특화된 유도 편향을 포함하도록 선택됩니다. 저자들은 **텐서화된 푸리에 신경 연산자(Tensorized Fourier Neural Operator, TFNO)**를 활용하는데, 이는 교사에게는 명시적으로 강제되지 않았던 평행 이동 등변성(translation equivariance)을 강제하고 고주파 노이즈를 억제합니다.

3. 주요 기여

1. 최초의 체계적 연구: 본 논문은 일반적인 사전 학습된 PDE 파운데이션 모델을 컴팩트한 다운스트림 대리 모델로 증류하는 것에 대한 최초의 체계적인 조사를 제시합니다.
2. TREX 프레임워크: 노이즈가 포함된 교사 롤아웃을 통해 감독 데이터를 생성함으로써, 명시적인 초기 조건 샘플러 없이도 학습이 가능한 새로운 증류 방법을 도입했습니다.
3. 유도 편향 통합: 일반적인 파운데이션 모델을 적절한 유도 편향(예: 등변성)을 가진 작은 아키텍처로 증류함으로써, 교사의 예측 정확도를 유지하면서도 물리적 특성을 준수할 수 있음을 입증했습니다.

4. 실험 결과

저자들은 다양한 유체 역학 벤치마크(Navier-Stokes 및 Compressible Euler 방정식)에 대해 Poseidon과 Walrus를 교사로 사용하여 TREX를 평가했습니다.

• 정확도: TREX로 학습된 학생 모델은 저데이터 환경(4~32개의 궤적)에서 거대 파운데이션 모델 교사의 정확도와 일치하거나 이를 능가했습니다. 일부 사례에서는 정답 데이터 없이 오직 교사가 생성한 데이터만으로 학습했을 때 학생 모델이 교사보다 더 나은 성능을 보이기도 했습니다.
• 효율성:
  • 파라미터 수: 학생 모델은 파라미터 수를 수 차례 규모로 줄였습니다 (예: Poseidon-L의 약 628M에서 TFNO의 약 174K로, 약 3600배 감소).
  • 추론 속도: 추론 시간이 한 자릿수 이상 단축되었습니다 (예: 스텝당 약 270ms에서 22ms로 감소).
  • 메모리: 순전파(forward pass)를 위한 GPU 메모리 소비량이 약 3.2배 감소했습니다.
• 안정성 및 등변성:
  • 학생 모델은 교사 모델에 비해 오차 증가 속도가 느리며, 더 안정적인 장기 롤아웃을 보여주었습니다.
  • TFNO 학생 모델은 비등변성(non-equivariant)인 교사에 비해 현저히 낮은 평행 이동 등변성 오차를 보였으며, 자기회귀적 롤아웃 중 공간적 이동(spatial shifts) 하에서도 일관성을 유지했습니다.
• 절제 실험 (Ablations):
  • 본 방법론은 롤아웃 길이와 노이즈 하이퍼파라미터에 대해 견고합니다.
  • 롤아웃 중 노이즈 주입은 특히 높은 진폭에서 유익한 것으로 나타났습니다.
  • 학생 아키텍처(UNet, 표준 FNO)가 정답 데이터만으로는 교사의 성능을 능가할 수 없는 경우에도 증류가 효과적임을 확인했습니다.

5. 의의 및 주장

본 논문은 TREX가 파운데이션 모델의 강력한 일반화 능력과 실용적인 대리 모델의 효율성 요구 사이의 트레이드오프를 효과적으로 해결한다고 주장합니다. 제한된 다운스트림 데이터를 교사 롤아웃을 통해 확장함으로써, TREX는 컴팩트한 모델이 방대한 데이터셋이나 복잡한 아키텍처 없이도 장기 역학 및 물리적 대칭성을 학습할 수 있게 합니다.

저자들은 이 연구를 계산 비용과 추론 속도가 중요한 실제 과학 및 공학 응용 분야(예: 제조 제어, 에너지 시스템)에서 파운데이션 모델 기반의 대리 모델을 배포 가능하게 만드는 단계로 규정합니다. 또한, 이 방법이 미세 조정된 교사의 품질에 의존하기는 하지만, 교사가 특정 유도 편향을 결여하고 있더라도 학생 모델이 등변성과 같은 물리적 제약 조건을 강제할 수 있도록 지식을 성공적으로 전이한다는 점을 강조합니다.