Learning to Think in Physics: Breaking Shortcut Learning in Scientific Diffusion via Representation Alignment

Motivation

• Known: 기존의 physics-informed diffusion model들은 PINN, FNO/PINO 등과 함께 최종 출력에만 PDE residual 제약을 부과하여 물리적 일관성을 강제해왔으며, 이는 boundary condition 분포가 고정된 경우 어느 정도 효과를 보여왔다.
• Gap: 출력 단에만 물리 제약을 부과하는 방식은 네트워크 내부가 물리적 추론을 우회하고 학습 데이터의 spurious correlation, discretization artifact를 암기하는 shortcut learning에 취약하며, boundary condition이 shift된 분포 외(OOD) 상황에서 취약성을 드러낸다는 근본적 한계가 존재한다.
• Why: 중간 표현을 물리 상태와 정렬시키는 inductive bias를 제공하면 모델이 통계적 상관관계 암기가 아닌 실제 물리 법칙을 내재화하게 되어, scientific diffusion model의 신뢰성과 일반화 성능을 근본적으로 향상시킬 수 있기 때문에 AI for Science 분야에서 중요한 문제다.
• Approach: diffusion backbone의 선택된 중간 레이어에 경량 1×1 projection head를 부착해 hidden activation을 물리량으로 decode하고, 여기에 first-principles PDE residual loss를 적용하여 학습 중에만 감독 신호로 사용하고 추론 시에는 제거하는 Physics-Informed Representation Alignment(REPA-P) 방식을 제안한다.

Achievement

1. Zero-overhead 물리 정렬: projection head를 추론 시 제거함으로써 추가 연산 비용 없이 중간 표현에 물리적 제약을 부과하는 방법을 제시했다.
2. 다중 PDE benchmark에서의 일관된 성능 향상: Darcy flow, topology optimization, electrostatic potential, turbulent channel flow 등 4개 과제에서 수렴 속도 최대 2배 가속, physics residual 최대 66.4% 감소, OOD robustness 최대 49.3% 개선을 달성했다.
3. 아키텍처 독립성 입증: U-Net과 Diffusion Transformer(DiT) 두 backbone 모두에서 일관된 성능 향상을 보여 architecture-agnostic함을 실증했다.
4. 효율적인 레이어 선택 통찰: ablation을 통해 소수의 핵심 중간 레이어 감독만으로도 대부분의 성능 이득을 얻을 수 있으며, 이는 output-level physics loss와 상호 보완적임을 규명했다.

How

• 표준 diffusion model(DDPM 계열)을 기반으로 물리 필드 x0를 예측하는 denoising objective(Ldata)를 사용
• Darcy flow의 질량 보존식과 같은 governing PDE와 boundary condition을 discretized residual operator R(x0)로 정의
• virtual-observable framework를 채택하여 residual을 Gaussian 변수로 모델링하고, diffusion posterior variance에 비례하는 분산(σ²(t)=β̃t)을 부여
• 선택된 중간 레이어의 hidden activation에 1×1 projection head를 부착해 물리 상태 변수로 decode
• decode된 물리 상태에 PDE residual loss를 직접 계산하여 gradient를 backbone 파라미터로 전파, 이를 통해 shortcut path를 차단하고 물리 법칙을 내재화하도록 유도
• 이 physics-informed representation alignment loss를 기존 output-level physics loss 및 denoising loss와 결합하여 학습
• 학습 완료 후에는 projection head를 제거하여 순수 backbone만으로 추론 수행

Originality

• 외부 teacher network 없이 first-principles PDE residual 자체를 "teacher" 신호로 활용하는 teacher-free representation alignment 개념을 제안
• 기존의 output-only physics loss나 generic deep supervision(보조 loss만 추가)과 달리, 중간 표현을 물리 상태로 직접 decode하고 PDE gradient를 짧은 경로로 역전파시켜 shortcut learning을 credit assignment 관점에서 근본적으로 차단
• projection head를 추론 시 완전히 제거해 zero inference overhead를 달성하는 실용적 설계
• U-Net과 DiT 등 이질적인 backbone에 동일하게 적용 가능한 architecture-agnostic 설계

Limitation & Further Study

• 어떤 중간 레이어를 선택해 supervision을 적용할지에 대한 원칙적 기준이 명확히 제시되지 않고 ablation에 의존하는 경향이 있어, 새로운 아키텍처나 task에 일반화할 때 추가 탐색이 필요할 수 있음
• 검증된 4개 PDE task는 비교적 규칙적인 격자 기반 문제로, 복잡한 3D 비정형 메쉬나 다중 물리(multi-physics) 결합 시스템에 대한 확장성은 추가 검증이 필요함
• virtual-observable framework의 Gaussian residual 가정이 강한 비선형 PDE나 불연속 해를 가지는 문제에서도 유효한지에 대한 이론적 분석이 부족함
• 향후 연구로 레이어 선택을 자동화하는 방법론이나, 더 복잡한 실세계 물리 시스템(예: multi-physics coupling, irregular geometry)으로의 확장이 필요함

Evaluation

총평: 중간 표현에 물리 법칙을 직접 정렬시켜 shortcut learning을 억제한다는 아이디어는 간결하면서도 효과적이며, zero-overhead 설계와 architecture-agnostic 특성이 실용성을 높인다. 다양한 PDE benchmark에서의 실증적 개선이 인상적이나, 레이어 선택 원칙과 더 복잡한 물리 시스템으로의 확장 가능성에 대한 추가 검토가 필요하다.