UniFluids: Unified Neural Operator Learning with Conditional Flow-matching

2026-03-19

AI Summary

This work proposes UniFluids, a unified framework for operator learning of partial differential equations (PDEs) that addresses the challenge of jointly handling PDE systems across varying dimensions and physical variables. By introducing conditional flow matching into unified operator learning for the first time, UniFluids integrates a diffusion-based Transformer architecture with a four-dimensional spatiotemporal representation, enabling joint training and parallel generation of solution operators for diverse PDEs. The framework employs an x-prediction strategy to enhance solution accuracy and demonstrates high predictive performance, strong scalability, and robust generalization across dimensions and physical scenarios. Extensive experiments on 1D, 2D, and 3D benchmark PDEs validate its effectiveness and versatility.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Learning with Manifolds
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

Abstract

Partial differential equation (PDE) simulation holds extensive significance in scientific research. Currently, the integration of deep neural networks to learn solution operators of PDEs has introduced great potential. In this paper, we present UniFluids, a conditional flow-matching framework that harnesses the scalability of diffusion Transformer to unify learning of solution operators across diverse PDEs with varying dimensionality and physical variables. Unlike the autoregressive PDE foundation models, UniFluids adopts flow-matching to achieve parallel sequence generation, making it the first such approach for unified operator learning. Specifically, the introduction of a unified four-dimensional spatiotemporal representation for the heterogeneous PDE datasets enables joint training and conditional encoding. Furthermore, we find the effective dimension of the PDE dataset is much lower than its patch dimension. We thus employ $x$-prediction in the flow-matching operator learning, which is verified to significantly improve prediction accuracy. We conduct a large-scale evaluation of UniFluids on several PDE datasets covering spatial dimensions 1D, 2D and 3D. Experimental results show that UniFluids achieves strong prediction accuracy and demonstrates good scalability and cross-scenario generalization capability. The code will be released later.

Problem

• PDE
• operator learning
• unified framework
• conditional generation
• spatiotemporal representation

Innovation

• conditional flow-matching
• unified neural operator
• diffusion Transformer
• spatiotemporal representation
• x-prediction