Expanding the Chaos: Neural Operator for Stochastic (Partial) Differential Equations

Motivation

• Known: Neural operator(NO)는 FNO, DeepONet, GINO 등으로 함수공간 간 매핑을 one-shot으로 학습할 수 있으며, Neural SPDE 등 기존 연구들은 관측된 stochastic process로부터 구성한 feature basis 위에서 결정론적 함수를 NO로 파라미터화하는 방식으로 SDE/SPDE에 stochasticity를 반영해왔다.
• Gap: 기존 NO 기반 SDE/SPDE 학습 연구들은 stochastic forcing을 NO에 통합하는 원리적 근거가 부족했고, 고전적인 Wiener–chaos expansion 이론과 neural operator 설계를 명시적으로 연결하여 chaos coefficient가 만족하는 결정론적 propagator 시스템을 유도하고 이를 다양한 SDE/SPDE 다운스트림 과제에 폭넓게 적용한 사례가 없었다.
• Why: SDE와 SPDE는 diffusion generative model, 금융 모델링, 유체역학, 통계역학 등 자연과학 및 머신러닝 전반의 핵심 도구이므로, 이들의 solution operator를 빠르고 원리적으로 학습할 수 있는 방법은 시뮬레이션 가속과 새로운 학습 관점 제공 측면에서 중요하다.
• Approach: driving Brownian(또는 Q-Brownian) noise path를 countable Gaussian coordinate로 재구성하고 이를 orthonormal Wick–Hermite basis로 사영하여 얻은 chaos coefficient를 결정론적 ODE/PDE 시스템으로 특성화한 뒤, 이 propagator를 neural operator backbone으로 파라미터화하여 SDE/SPDE solution을 재구성한다.

Achievement

Figure 3. Comparison between ground truth and F-SPDENO solutions of the 2D

1. WCE 기반 NO 프레임워크 구축: SPDE(semi-linear)와 SDE(다차원)에 대해 Wiener–chaos expansion으로 solution을 chaos coefficient 형태로 표현하고, 이 coefficient들이 만족하는 결정론적 ODE/PDE 시스템을 Theorem 1, 2로 증명하여 NO 파라미터화의 이론적 근거를 마련했다.
2. 광범위한 실험적 검증: 표준 SPDE 벤치마크, SDE 기반 diffusion 모델의 one-step image sampling(CIFAR10), 위상학적 graph interpolation, 금융 시계열 extrapolation, parameter estimation, manifold SDE 홍수 예측 등 다양한 도메인에서 경쟁력 있는 정확도를 달성했다.
3. law-matching 설정으로의 확장: 학습 solution과 짝지어진 Brownian trajectory가 관측되지 않는 상황에서도, 규정된 noise prior로부터 auxiliary Brownian path를 샘플링하여 프레임워크가 여전히 효과적임을 보였다.

How

Figure 2. General working flow of F-SPDENO: Discretized Q-Brownian motion as the initial input, Wick features computed f

• Gaussian random variable ξij를 orthonormal basis {ej}와 Brownian increment의 적분으로 정의하고, 이를 통해 1차원 및 다차원 Brownian motion, Q-Brownian motion(Karhunen–Loève expansion 활용)을 truncation order n으로 재구성(Lemma 1, 2).
• 재구성된 noise coordinate로부터 orthonormal Wick–Hermite feature(Wick polynomial)를 구성하여 유한차원 noise representation η을 얻음.
• solution의 chaos coefficient가 만족하는 결정론적 propagator 시스템(SDE의 경우 ODE 시스템, semilinear SPDE의 경우 PDE 시스템)을 도출(Theorem 1, 2).
• 이 propagator를 FNO, DeepONet 등 다양한 NO backbone으로 파라미터화하여 (χ0, η) → û(·) 매핑을 학습, 이산 space-time grid 상에서 MSE 기반 empirical L2 손실로 최적화.
• SPDE에는 F-SPDENO, 이미지 diffusion에는 U-SDENO, 그래프 보간에는 G-SDENO 등 태스크별 변형 아키텍처 적용.

Originality

• Wiener–chaos expansion이라는 고전적 확률론 도구를 neural operator 설계와 명시적으로 연결하여, 임의적인 feature 구성이 아닌 이론적으로 정당화된 orthonormal Wick–Hermite basis를 사용한 점.
• chaos coefficient가 만족하는 결정론적 ODE/PDE propagator 시스템을 SDE와 semilinear SPDE 모두에 대해 정리(Theorem)로 제시하여 원리적 근거를 마련한 점.
• diffusion 이미지 생성, 그래프 보간, 금융 예측, manifold SDE 등 이질적인 도메인 전반에 하나의 통합 프레임워크를 적용하고, Brownian path가 관측되지 않는 law-matching 설정까지 다룬 점.

Limitation & Further Study

• truncation order n에 따른 근사 오차와 계산 비용 간의 trade-off에 대한 체계적 분석이 제한적일 수 있으며, 고차원 noise space나 매우 긴 시간축에서의 확장성(scalability)이 충분히 검증되지 않았을 가능성이 있다.
• law-matching 설정에서 auxiliary Brownian path 샘플링이 실제 noise prior와 얼마나 잘 일치하는지에 대한 이론적 보장이 부족해 보이며, 이 부분에 대한 후속 연구가 필요하다.
• 실험이 다양한 도메인을 다루지만 각 도메인에서의 baseline 비교가 제한적이거나 task-specific NO 변형(U-SDENO, G-SDENO, F-SPDENO)의 설계 선택에 대한 ablation이 충분히 제시되지 않았을 수 있다.

Evaluation

총평: 고전적 Wiener-chaos expansion 이론을 neural operator 설계에 원리적으로 접목하여 SDE/SPDE solution operator 학습에 새로운 이론적 기반을 제공하고, 다양한 도메인에서 폭넓게 검증한 의미 있는 연구이다. 다만 확장성과 law-matching 설정에서의 이론적 보장에 대한 추가 검증이 필요하다.