Walrus: A Cross-domain Foundation Model for Continuum Dynamics

Motivation

• Known: 기존 연구들은 PDE 기반 physics-informed 모델, hybrid 모델, 그리고 MPP, Poseidon(Herde et al., 2024) 등 pretraining 기반 foundation model을 통해 데이터 기반 물리 시뮬레이션 방향을 제시해왔으나 대체로 2D 또는 고정 해상도의 상대적으로 동질적인 데이터셋에 국한되어 있었다.
• Gap: 물리 시뮬레이션 분야에서는 language/vision 분야와 달리 데이터 이질성(해상도, 차원, 물리 필드 다양성)과 장기 예측의 불안정성이 다양한 물리 동역학으로부터의 학습을 저해하며, 현대 하드웨어에서 효율적인 2D-3D 혼합 학습을 지원하는 방법이 부재하다는 문제가 있었다.
• Why: 진정한 cross-domain physical foundation model을 구축하려면 다양한 물리 도메인과 차원, 해상도를 아우르는 안정적이고 효율적인 pretraining 방법론이 필수적이며, 이는 과학computing 전반에서 시뮬레이션 대체 및 전이 학습 활용성을 크게 높일 수 있다.
• Approach: harmonic-analysis 기반의 patch jittering 안정화 기법, load-balanced 분산 2D-3D 학습 전략, compute-adaptive tokenization(CSM)을 결합하여 transformer 기반 foundation model인 Walrus를 설계하고 19개의 이질적 물리 시나리오에서 pretraining하였다.

Achievement

Figure 7. Pretraining VRMSE from HALFWALRUS and naively trained models. From pretraining alone, it appears that the less

1. 장기 예측 안정성 개선: patch jittering 기법이 84%의 pretraining 시나리오에서 long-horizon error를 감소시킴을 확인함.
2. 훈련 처리량 향상: topology-aware sampling을 통해 training throughput을 262% 향상시킴.
3. 2D-3D 통합 학습: 2D 데이터를 3D 공간에 무작위로 임베딩된 평면으로 취급하는 2D-to-3D augmentation을 통해 서로 다른 차원의 데이터를 단일 모델에서 효율적으로 결합 학습함.
4. 성능 우위 입증: 단기 및 장기 예측 지평 모두에서, downstream task와 pretraining 데이터 전반에서 기존 foundation model 대비 우수한 성능을 보임.
5. 다양성 우선 사전학습의 효과 검증: ablation을 통해 안정성, 처리량, 전이 성능 측면에서 기존 접근법 대비 각 기여 요소들의 가치를 확인함.

How

Figure 2. All raw data is projected into 3D by appending singleton

• 아키텍처: 공간-시간 factorized transformer 구조 채택, 공간 축에는 axial RoPE를 사용한 parallelized attention, 시간 축에는 T5-style relative position encoding을 사용한 causal attention 적용, 두 블록 모두 QK normalization 사용.
• Compute-Adaptive Compression: Convolutional Stride Modulation(CSM)을 encoder/decoder에 적용해 해상도에 따라 downsampling/upsampling 수준을 동적으로 조절, 고정된 tokens-per-axis를 유지하며 학습 효율 극대화.
• Shared Encoder-Decoder: 동일 차원(d)의 모든 물리 시스템이 첫 projection만 다른 encoder/decoder block을 공유하도록 하여 일반화 가능한 특징 학습을 유도, Hierarchical MLP(hMLP) 및 transposed hMLP로 경량 인코딩/디코딩 수행.
• 정규화: RMSGroupNorm(GroupNorm을 그룹별 RMSNorm으로 구현)을 표준 정규화로 사용하고, 입력과 출력(Δu)의 분포 차이를 고려한 비대칭 정규화 적용.
• Patch Jittering: harmonic-analysis에서 파생된 안정화 기법으로 high-frequency artifact 누적을 줄여 장기 예측 안정성 개선.
• 2D-to-3D Augmentation: 원시 데이터를 singleton 차원을 추가해 3D로 project하여 2D/3D 데이터를 통합 처리.
• Topology-aware Sampling: 분포의 topology와 sampling 전략을 연계해 load balancing 및 분산 학습 처리량 개선.

Originality

• 기존 physical foundation model들이 대부분 2D 혹은 단일 해상도에 국한된 것과 달리, 19개의 이질적 물리 도메인(astrophysics부터 classical fluids까지)과 2D/3D 데이터를 동시에 포괄하는 최초 시도 중 하나.
• harmonic-analysis 이론에 기반한 patch jittering이라는 새로운 안정화 기법을 제안하여, 기존의 노이즈 주입이나 augmentation 기법과 차별화됨.
• CSM 기반 adaptive-compute tokenization을 통해 고정 압축 encoder의 한계를 극복하고 다양한 해상도를 native하게 처리.
• topology-aware sampling을 통한 load-balanced 분산 학습이라는 시스템 수준의 기여로 순수 모델링을 넘어선 실질적 확장성 문제 해결.

Limitation & Further Study

• 발췌된 부분만으로는 19개 시나리오의 세부 구성, 실제 downstream task의 범위와 다양성, 정량적 성능 비교치(SOTA 대비 개선폭)의 구체적 수치가 충분히 드러나지 않음.
• patch jittering이 84% 시나리오에서 효과가 있다는 것은 나머지 16% 시나리오에서는 효과가 제한적이거나 부정적일 수 있음을 시사하며 이에 대한 원인 분석이 필요.
• 현재는 fluid-like continuum dynamics에 초점이 맞춰져 있어 고체역학, 복잡한 경계조건, 매우 긴 시간 지평에서의 안정성 등 더 넓은 물리 시스템으로의 일반화는 후속 연구 과제로 남음.
• 1.3B 파라미터 규모의 모델 학습에 필요한 대규모 컴퓨팅 자원 요구가 커뮤니티의 재현 및 추가 실험에 진입장벽이 될 수 있음.

Evaluation

총평: 물리 시뮬레이션을 위한 cross-domain foundation model 구축에 있어 데이터 이질성과 장기 안정성이라는 핵심 난제를 이론에 기반한 실용적 기법들로 해결한 인상적인 시스템 및 모델링 연구이며, 코드와 가중치 공개를 통해 커뮤니티 기여도가 클 것으로 기대된다.