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

TL;DR

UniPINN is proposed, 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.

Abstract

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

Figures (4)

• Figure 1: Single-task vs. unified multi-flow learning. (a) Conventional single-task PINNs train independent networks for each flow type, incurring parameter redundancy and neglecting the shared Navier-Stokes structure across flows. (b) UniPINN unifies multiple flow types under a shared backbone with task-specific heads and cross-flow attention, enabling joint learning and knowledge transfer while preserving flow-specific fidelity.
• Figure 2: UniPINN architecture. It consists of three components: (1) Task-shared backbone network: extracts common physical features from spatiotemporal coordinates and task embeddings for all flow types. (2) Task-specific feature extraction: the task-aware multi-task attention selectively reinforces the shared representation; the task-specific dedicated layers decode the enhanced features into stream function and pressure, with velocity obtained via automatic differentiation. (3) Dynamic weight allocation: adaptively reweights equation, boundary, and data loss terms across flows to ensure stable joint optimization.
• Figure 3: Predicted flow field visualization (streamline plot). The model accurately captures characteristic physical structures: (a) Lid-driven cavity primary vortex; (b) Pipe flow parabolic profile; (c) Couette flow linear profile.
• Figure 4: Training convergence analysis (loss curves) during multi-flow PINN training. (a) Unweighted Total Loss; (b) Equation Loss (Navier-Stokes residuals); (c) Boundary Condition Loss; (d) Data Loss. All loss components and the total loss show stable convergence under the DWA strategy.