Generative Neural Operators through Diffusion Last Layer

Motivation

• Known: Neural operator(DeepONet, FNO 등)는 discretization invariant한 함수 공간 간 매핑을 학습하는 강력한 프레임워크로 자리잡았으며, 최근에는 Bayesian neural operator나 함수 공간 생성모델 등 확률적 operator learning 연구가 활발히 진행되고 있다. 또한 diffusion model과 flow matching은 고차원 필드에 대한 조건부 생성 모델링에 효과적이며, latent 공간에서의 생성모델링을 통해 확장성을 개선하는 연구도 존재한다.
• Gap: 기존 표준 neural operator는 결정론적 점 추정만 제공하여 예측 불확실성을 포착하지 못하며, 기존의 확률적 확장(무한차원 함수공간 diffusion, pixel space diffusion, 일반적 latent diffusion 등)은 계산 비용이 크거나 operator learning 고유의 discretization invariance와 구조적 이점을 온전히 유지하지 못한다는 한계가 있다.
• Why: 과학 응용에서 확률적 forcing, 미해결된 subgrid 효과, 불확실한 계수, 카오스적 민감성 등으로 인해 target field가 본질적으로 확률적인 경우가 많아 신뢰할 수 있는 surrogate 모델링을 위해서는 uncertainty quantification이 필수적이다. DLL은 기존 neural operator backbone에 경량 모듈만 추가하여 이러한 확률적 모델링과 UQ를 가능하게 하므로 실용적 가치가 크다.
• Approach: Neural operator encoder가 생성하는 입력 의존적 저랭크 기저 함수 위에서 계수 벡터를 정의하고, 이 계수 공간에서 conditional diffusion(또는 flow matching) model을 학습하여 backbone의 discretization invariance를 유지한 채 조건부 생성모델로 확장한다.

Achievement

Figure 3. Stochastic Burgers’ equation. Columns compare the ground truth and different surrogate models. For each method

1. 분포 충실도(distributional fidelity): 무작위 forcing이 있는 stochastic PDE benchmark(예: stochastic Burgers, stochastic Darcy flow)에서 DLL이 pixel space 및 conventional latent diffusion baseline과 경쟁력 있는 성능을 보인다.
2. rollout 안정성 개선: 결정론적 long horizon rollout task(예: KS equation)에서 DLL이 backbone 대비 rollout 안정성을 향상시킨다.
3. 유용한 불확실성 추정: autoregressive rollout에서 발생하는 compounding error 하에서도 DLL이 유의미한 predictive uncertainty 추정치를 제공한다.
4. 경량 모듈화: DLL은 기존 neural operator backbone에 부착 가능한 lightweight probabilistic head로, 구조적 이점을 보존하면서 결정론적 operator를 조건부 생성기로 전환한다.

How

Figure 2. Training and inference pipeline for DLL. (1) Operator encoder. A NO backbone produces basis functions Φ(a), an

• Neural operator(NO) 백본의 encoder를 이용해 입력 a에 의존하는 저랭크 기저 함수(Karhunen–Loève expansion에서 영감을 받은 형태)를 생성.
• 해당 기저에 대응하는 계수 벡터 x ∈ R^{dx}를 target field의 유한차원 표현으로 사용.
• 계수 공간에서 조건부 velocity field v_φ(x_t, t, c)를 flow matching/velocity prediction 방식으로 학습(식 (3), (4)의 forward noising process와 conditional velocity matching objective 사용).
• 학습된 velocity field로 probability flow ODE를 역방향으로 적분하여 p(u|a)로부터 샘플링.
• 결정론적 문제와 확률적 문제(Problem 2.1, 2.2)를 conditional diffusion model이라는 통합된 확률적 프레임워크로 다룸.
• 2-Wasserstein distance를 이용해 학습된 분포와 목표 분포 간의 거리를 비교, 이론적으로 velocity 오차와 endpoint 분포 오차 간의 관계를 분석(Proposition 2.3 등).
• Stochastic Burgers, stochastic Darcy flow, KS equation 등 benchmark에서 pixel space diffusion, conventional latent diffusion과 비교 실험 수행.

Originality

• 결정론적 neural operator backbone에 diffusion 기반 확률적 output head를 부착하는 모듈형 설계(diffusion last layer)를 제안, 기존 backbone 구조를 재사용 가능하게 함.
• Karhunen–Loève expansion에서 영감을 받은 입력 의존적 저랭크 기저를 활용해 target field를 압축하고, 이 계수 공간에서 conditional diffusion을 학습하는 접근은 무한차원 함수공간 diffusion이나 일반적 latent diffusion과 차별화됨.
• 결정론적 문제와 확률적 문제를 conditional diffusion model이라는 단일 확률적 프레임워크로 통합하여 다루는 관점 제시.
• 2-Wasserstein distance 및 velocity based transport stability 관점에서 diffusion 학습을 이론적으로 정당화.

Limitation & Further Study

• 발췌된 abstract 및 본문 앞부분만으로는 계수 공간의 차원(dx) 선택이나 저랭크 기저의 수 등 하이퍼파라미터가 성능과 계산 비용에 미치는 영향에 대한 구체적 분석이 부족해 보인다.
• 저랭크 근사가 고차원 또는 복잡한 다중 스케일 물리 시스템에서 충분한 표현력을 가질 수 있는지에 대한 검증이 제한적일 수 있다.
• 실험이 상대적으로 저차원 PDE benchmark(Burgers, Darcy flow, KS equation)에 집중되어 있어, 실제 대규모 과학 응용(예: 기상 예보, 3차원 난류)으로의 확장성 검증이 후속 연구로 필요하다.
• Bayesian neural operator 등 epistemic uncertainty를 다루는 기존 방법과의 정량적 비교가 충분히 제시되었는지 추가 확인이 필요하다.

Evaluation

총평: Neural operator의 구조적 이점과 diffusion model의 확률적 표현력을 저랭크 계수 공간에서 효율적으로 결합한 실용적이고 독창적인 접근으로, uncertainty aware neural operator 연구에 의미 있는 기여를 한다.