Physics-Informed Neural Networks for coupled stiff transport systems

Abstract: Purpose: Physics-Informed Neural Networks (PINNs) struggle with stiff, regime-changing transport equations due to instability, loss imbalance, and violations of physical consistency. This paper investigates these failures through the Marshak wave equations - a canonical benchmark from radiative transport - where initial and boundary conditions differ by up to 12 orders of magnitude, and proposes targeted modifications to the standard PINN framework to overcome them. Design/methodology/approach: Three modifications are introduced: (1) a ScaledSigmoid final activation enforcing physical bounds and positivity of the unknowns; (2) a logarithmic MSE loss replacing the standard quadratic loss for initial and boundary conditions, enabling training across extreme scale disparities; and (3) explicit enforcement of global conservation laws derived from the governing equations as an additional physics loss term. Monte Carlo sampling with exponential time weighting is used throughout. Findings: The proposed framework successfully recovers the Marshak wave dynamics - including the hot, cold, and wave-front regions - in agreement with a reference Implicit Monte Carlo solution, with run times under 30 minutes. Ablation studies confirm that each ingredient is essential: linear activation, absence of the logarithmic loss, or removal of the PDE term each independently cause the method to fail qualitatively. Originality/value: This work identifies and resolves three concrete failure modes of standard PINNs on stiff hyperbolic systems with nonlinear coupling. The combination of bounded activations, scale-aware loss functions, and conservation law enforcement constitutes a novel and practically validated framework, with applicability to radiative transport and other coupled stiff PDE systems in engineering.

• The paper introduces a novel PINN architecture integrating scale-aware loss functions to robustly solve stiff, coupled transport PDEs.
• It employs positivity-preserving bounded activations and logarithmic MSE loss to enforce physical constraints across 12 orders of magnitude.
• Experimental results demonstrate accurate replication of thermal dynamics and wavefront propagation, outperforming standard PINN methods.

Physics-Informed Neural Networks for Coupled Stiff Transport Systems

Introduction

The paper "Physics-Informed Neural Networks for coupled stiff transport systems" (2606.21236) targets a central limitation of the canonical PINN framework: its inability to solve stiff, regime-changing hyperbolic PDEs under highly disparate initial and boundary conditions. Using the Marshak wave equation as a representative case, the authors identify three distinct sources of failure: magnitude imbalance, loss of physical consistency, and instability in enforcing boundary/initial constraints. To address these, the proposal combines three architectural and methodological advances: positivity-preserving bounded activations, logarithmic scaling of loss terms, and explicit conservation law enforcement. The result is a robust PINN capable of solving tightly coupled radiative transport systems across 12 orders of magnitude scale disparity, with strong numerical evidence for both accuracy and stability.

Figure 1: Illustration of a neural network generating solution functions with automatic differentiation support for inputs and parameters.

Problem Setup and Challenges

PINNs are built to approximate solutions to PDEs by embedding physics into the loss function. However, when applied to stiff systems such as the Marshak wave equations—where initial/boundary values differ by up to $10^{12}$—standard PINNs break due to:

• Severe imbalance in loss terms, skewing optimization toward large-scale values.
• Inability to enforce positivity and boundedness, especially for nonlinear couplings involving divisions (e.g., $\frac{1}{T^3}$).
• Poor satisfaction of conservation laws intrinsic to hyperbolic transport phenomena.

The Marshak system, coupling intensity $u(t,x,\omega)$ and temperature $T(t,x)$, exemplifies this stiffness. It admits regions of sudden transition (wavefront), necessitating a network architecture capable of robustly enforcing both physical constraints and global conservation.

Network Architecture and Activation Strategies

The architecture utilizes a single neural network with a polynomial embedding of inputs ($t$, $x$), followed by separate branches for $u$ and $T$, each with tailored output activations.

Figure 2: Actual architecture of the neural network, featuring shared trunk and branched outputs for the unknowns.

The core innovation is the ScaledSigmoid activation function:

$$\text{ScaledSigmoid}(x) = v_{min} + (v_{max} - v_{min}) \frac{1}{1 + e^{-x}}$$

This ensures the output strictly obeys physical bounds ($u\in[u_\text{init},u_\text{leftbcu}]$, $\frac{1}{T^3}$0), prevents negative values, and curtails instability from near-zero denominators in nonlinear terms.

Other activations (tanh, ReLU variants, SafeSoftplus) were empirically inferior, failing to balance expressivity, positivity, and boundedness required by the physics.

Loss Functions: Logarithmic MSE and Conservation Enforcement

Traditional quadratic error loss is ill-suited for the scale disparity: it overwhelmingly weights large-magnitude conditions, marginalizing accuracy in low-magnitude regions. The authors introduce a logarithmic MSE term:

$\frac{1}{T^3}$1

for initial and boundary loss components, which normalizes errors across scales, allowing simultaneous enforcement of disparate constraints.

Additionally, conservation laws derived from the governing equations are integrated as explicit physics loss terms:

$\frac{1}{T^3}$2

computed for representative time slices via Monte Carlo sampling. This guarantees global consistency, essential for stiff hyperbolic systems, and provides leverage where local PINN residuals may be insufficient.

Numerical Methodology and Implementation

Monte Carlo sampling with exponentially weighted time points (concentrating at $\frac{1}{T^3}$3), polynomial embedding, and batch sizes of $\frac{1}{T^3}$4 underpin the stochastic minimization. Parameter choices ensure the maximum principle is reflected in activation bounds.

Experimental Results and Analysis

The PINN framework is benchmarked against Implicit Monte Carlo reference solutions for the Marshak wave. Critical metrics include the emergence of distinct cold, hot, and wavefront regions in $\frac{1}{T^3}$5 profiles, adequacy of initial/boundary enforcement, and propagation speed.

Figure 3: Reference solution for the Marshak wave problem obtained via Implicit Monte Carlo, showing wavefront progression.

The PINN solution, obtained with the modified architecture and loss, exhibits:

• Accurate reproduction of thermal regions and wavefront propagation.
• Strong agreement in both magnitude and spatial structure compared to reference results.
• Effective handling of initial/boundary mismatches across 12 orders of magnitude.

Figure 4: Execution result after $\frac{1}{T^3}$6 iterations with clear separation between hot, wavefront, and cold regions.

Convergence analysis shows solution stabilization after several thousand iterations, with persistent accuracy in critical regions.

Figure 5: Convergence evolution over iterations, illustrating the establishment and persistence of correct dynamics.

Initial and boundary constraints are tightly enforced:

Figure 6: Initial and left boundary conditions for the PINN solution, demonstrating precise agreement with target values.

Ablation Studies

Ablation studies validate the indispensability of each ingredient:

• Switching to linear activation or omitting the PDE loss results in qualitative failure—incorrect profiles and unstable dynamics.
• Removing the logarithmic scaling of the loss seriously impairs enforcement of low-scale boundary/initial constraints, with persistent mismatches even after prolonged training.

Figure 7: Profiles with linear activation or without PDE loss—demonstrating breakdown in physical correctness.

Figure 8: Boundary conditions with MSE loss in place of logarithmic—revealing major mismatches across several orders of magnitude.

Implications and Future Directions

The architectural and loss function modifications enable PINNs to robustly solve nonlinearly coupled stiff hyperbolic PDEs, opening practical applications in radiative transport, plasma physics, and stiff multiphysics simulation. The bounded activation and scale-aware loss mechanisms are generally extensible for other physically constrained ML tasks where parameter regimes span multiple orders of magnitude.

The explicit enforcement of conservation laws via loss terms could be generalized to other domains (e.g., electromagnetics, coupled reaction-diffusion), potentially synergizing with operator-learning frameworks and hybrid solver-PINN approaches.

Improved robustness may permit scalable surrogate modeling, uncertainty quantification, and inverse problem solutions where classical numerical methods are prohibitive or infeasible.

Conclusion

This work establishes a validated PINN methodology for stiff, regime-changing coupled transport systems, resolving critical sources of failure observed in canonical PINNs. By integrating positivity-preserving bounded activations, logarithmic scale loss functions, and explicit conservation enforcement, it provides both theoretical structure and empirical evidence for stable, accurate solutions under extreme scale disparity. The implications extend PINNs' operational envelope toward highly nonlinear, multiphysics scenarios, anticipating further advances in architecture design and physics-augmented learning for scientific computing.