UniPINN: A Unified PINN Framework for Multi-task Learning of Diverse Navier-Stokes Equations

Physics-Informed Neural Networks (PINNs) have shown promise in solving incompressible Navier-Stokes equations, yet existing approaches are predominantly designed for single-flow settings. When extended to multi-flow scenarios, these methods face three key challenges: (1) difficulty in simultaneously capturing both shared physical principles and flow-specific characteristics, (2) susceptibility to inter-task negative transfer that degrades prediction accuracy, and (3) unstable training dynamics caused by disparate loss magnitudes across heterogeneous flow regimes. To address these limitations, we propose UniPINN, a unified multi-flow PINN framework that integrates three complementary components: a shared-specialized architecture that disentangles universal physical laws from flow-specific features, a cross-flow attention mechanism that selectively reinforces relevant patterns while suppressing task-irrelevant interference, and a dynamic weight allocation strategy that adaptively balances loss contributions to stabilize multi-objective optimization. Extensive experiments on three canonical flows demonstrate that UniPINN effectively unifies multi-flow learning, achieving superior prediction accuracy and balanced performance across heterogeneous regimes while successfully mitigating negative transfer. The source code of this paper will be released on https://github.com/Event-AHU/OpenFusion

💡 Research Summary

Physics‑informed neural networks (PINNs) have become a powerful tool for solving forward and inverse problems in fluid dynamics by embedding the governing Navier‑Stokes equations directly into the loss function. However, most existing PINN research focuses on a single flow configuration, which limits their applicability to real‑world scenarios where multiple flow regimes—different Reynolds numbers, viscosities, boundary conditions, and geometries—must be handled simultaneously. Training separate networks for each flow leads to redundant parameters, high computational cost, and missed opportunities for knowledge transfer across flows that share the same underlying physics. Moreover, naïve multi‑task integration of heterogeneous flows suffers from three major challenges: (1) difficulty in jointly learning universal physical laws and flow‑specific features; (2) negative transfer, where interference from unrelated tasks degrades performance; and (3) unstable optimization caused by loss components that differ by several orders of magnitude, resulting in gradient pathology.

The paper introduces UniPINN, a unified multi‑task PINN framework designed to overcome these obstacles. UniPINN comprises three complementary components. First, a shared‑specialized architecture separates a common backbone that learns universal Navier‑Stokes operators from task‑specific heads that encode individual flow parameters (viscosity, density, boundary conditions) and fine‑grained features such as boundary‑layer dynamics or vortex shedding. This design preserves parameter efficiency while maintaining the fidelity of each flow. Second, a cross‑flow attention module combines self‑attention (to highlight salient regions within a single flow) with cross‑attention (to identify and exchange similar topological patterns across different flows). By selectively amplifying shared physical patterns and suppressing incompatible features, the attention mechanism mitigates negative transfer and encourages constructive knowledge sharing. Third, a dynamic weight allocation (DWA) strategy monitors the residual distributions of each flow’s PDE, boundary, and initial condition losses in real time and automatically adjusts their weighting in the total loss. This adaptive balancing prevents any single task from dominating the gradient direction, thereby stabilizing training across heterogeneous regimes.

Extensive experiments were conducted on three canonical incompressible flows: lid‑driven cavity, pipe (Poiseuille) flow, and Couette flow. These cases span a range of Reynolds numbers and boundary conditions, providing a rigorous testbed for multi‑flow learning. UniPINN consistently outperformed conventional single‑task PINNs, achieving an average reduction of more than 30 % in L2 error across all metrics. In the high‑Reynolds cavity case, UniPINN accurately captured vortex core locations and strengths, which single‑task baselines struggled to resolve. Ablation studies demonstrated that the cross‑flow attention module alone reduced negative transfer by over 70 %, while the DWA component eliminated loss‑scale‑induced instability, leading to smoother convergence curves and higher final accuracy. The shared‑specialized backbone contributed to faster pre‑training and reduced overall parameter count compared with training independent networks for each flow.

The authors also position UniPINN relative to existing multi‑task learning techniques such as soft parameter sharing and uncertainty‑based weighting. They argue that those methods, originally devised for purely data‑driven tasks, do not account for the unique characteristics of physics‑constrained optimization—namely, the presence of PDE residuals with vastly different magnitudes and the need to preserve strict physical consistency. UniPINN’s integration of physics‑aware attention and dynamic loss balancing directly addresses these gaps.

Finally, the paper outlines future directions, including scaling the framework to three‑dimensional turbulent flows, fluid‑structure interaction problems, and real‑time inference for control applications. The authors commit to releasing the source code and datasets on GitHub, facilitating reproducibility and further research. In summary, UniPINN represents a significant step toward unified, efficient, and robust multi‑flow PINN modeling by harmonizing shared physical knowledge, selective cross‑flow feature exchange, and adaptive optimization.