A Robust Foundation Model for Conservation Laws: Injecting Context into Flux Neural Operators via Recurrent Vision Transformers

TL;DR

This paper introduces a novel neural operator architecture that integrates recurrent Vision Transformers with Flux Neural Operators, enhancing robustness and generalization in solving conservation laws without explicit PDE knowledge.

Contribution

The authors develop a hypernetwork-based model that encodes solution dynamics with ViT and generates context-conditioned neural operators, improving conservation law solutions.

Findings

Model preserves robustness and generalization of Flux NO.

Achieves reliable solutions across diverse conservative systems.

Outperforms standard neural operators in long-term predictions.

Abstract

We propose an architecture that augments the Flux Neural Operator (Flux NO), which combines the classical finite volume method (FVM) with neural operators, with ViT-based context injection. Our model is formulated as a hypernetwork: it extracts solution dynamics over a finite temporal window, encodes them with a recurrent Vision Transformer, and generates the parameters of a context-conditioned neural operator. This enables the model to infer and solve conservation laws without explicit access to the governing equation or PDE coefficients. Experimentally, we show that the proposed method preserves the robustness, generalization ability, and long-time prediction advantages of Flux NO over standard neural operators, while delivering reliable numerical solutions across a broad range of conservative systems, including previously unseen fluxes. Our code is available at https://github.com/xx257xx/CONTEXT_FLUX_NO.