Preserving Conservation Laws in the Time-Evolving Natural Gradient Method via Relaxation and Projection Techniques

1. Problem Statement

Neural networks, particularly Physics-Informed Neural Networks (PINNs), have shown promise in solving Partial Differential Equations (PDEs). However, standard global training approaches often lack temporal causality, leading to instability in time-dependent problems. While local training frameworks like the Neural Galerkin scheme and the Time-Evolving Natural Gradient (TENG) method address causality by updating parameters sequentially over time, they face a critical limitation: the violation of physical conservation laws.

Specifically, when applying TENG to conservative systems (e.g., those conserving mass, energy, or momentum), two fundamental issues arise:

1. Target Solution Drift: The target solution ($u_{target}$), typically generated by explicit Runge-Kutta methods, often fails to preserve quadratic or nonlinear invariants (like energy), causing the training data itself to deviate from physical reality.
2. Parameter Manifold Drift: The parameter updates in TENG occur in a tangent space approximation that does not inherently enforce conservation constraints. Consequently, the network parameters $\theta(t)$ drift away from the theoretical invariant manifold ($M_I$), leading to non-physical long-term behavior and energy dissipation.

2. Methodology: The RP-TENG Framework

The authors propose RP-TENG (Relaxation-Projection TENG), a dual-strategy enhancement to the TENG method designed to strictly preserve conservation laws while maintaining high accuracy.

A. Relaxation Algorithm (Structure-Preserving Targets)

To address the first issue, the authors integrate a Relaxation Runge-Kutta approach to generate the target solution $u_{target}$.

• Mechanism: Instead of using a fixed time step $\Delta t$, a relaxation parameter $\gamma$ (or $r$) is introduced to scale the time step adaptively.
• Process: The target is computed as $u_{target} = \hat{u}_{\theta_n} + \gamma \Delta t f(\hat{u}_{\theta_n})$. The scalar $\gamma$ is solved via a nonlinear equation (using Newton's method) such that the invariant $I(u_{target})$ exactly equals the invariant of the previous state $I(\hat{u}_{\theta_n})$.
• Benefit: This ensures the training target strictly adheres to conservation laws (both linear and nonlinear) without sacrificing the formal order of accuracy of the underlying time discretization.

B. Projection Algorithm (Manifold Constrained Updates)

To address the second issue, a Projection Technique is applied to the network parameters after the TENG update step.

• Mechanism: Even with a conservative target, the mapping from physical space to parameter space may cause drift. The authors project the updated parameters $\theta_{temp}$ back onto the conservation manifold $M_I$.
• Process: This is formulated as a constrained optimization problem:
  $$\theta^* = \arg \min_{\eta \in \Theta} \frac{1}{2} \|\eta - \theta_{temp}\|^2_2 \quad \text{subject to} \quad \hat{I}_j(\hat{u}(\eta, \cdot)) = c_j$$
  The solution is found using Lagrange multipliers, iteratively correcting the parameters to satisfy the invariant constraints.
• Benefit: This forces the network trajectory to remain strictly on the invariant manifold, preventing long-term accumulation of conservation errors.

3. Key Contributions

1. Dual-Strategy Enhancement: The paper introduces a novel combination of relaxation (for the target) and projection (for the parameters) to the TENG framework, creating a fully structure-preserving solver.
2. Theoretical Rigor: The method addresses the specific failure modes of standard TENG regarding quadratic and nonlinear invariants, which are notoriously difficult for explicit integrators to preserve.
3. Algorithmic Implementation: The authors provide a complete algorithm (Algorithm 1) detailing the integration of these techniques, including the handling of the mass matrix singularity and the iterative solution for Lagrange multipliers.
4. Non-Intrusive Correction: The relaxation parameter $\gamma$ is shown to remain extremely close to 1 (deviations $\sim 10^{-12}$), indicating that the physical constraints are enforced with minimal perturbation to the system's intrinsic dynamics.

4. Numerical Results

The authors validated RP-TENG on three benchmark PDEs using Python/JAX with double-precision arithmetic:

• Inviscid Burgers Equation (Mass Conservation):
  • Standard TENG showed significant mass drift ($10^{-9}$ to $10^{-5}$).
  • RP-TENG reduced mass error to machine precision ($\sim 10^{-14}$), effectively eliminating drift.
• Korteweg-de Vries (KdV) Equation (Quadratic/Energy Invariant):
  • Standard TENG failed to preserve the quadratic energy invariant, leading to significant Hamiltonian drift.
  • RP-TENG suppressed energy error to machine epsilon ($\sim 10^{-15}$).
  • The method also improved the relative $L_2$ error growth, preventing phase drift and artificial dissipation common in unconstrained models.
• Acoustic Wave Equation (Hamiltonian Preservation):
  • RP-TENG achieved a relative error of $3.44 \times 10^{-3}$ compared to $3.55 \times 10^{-3}$ for the baseline, demonstrating that enforcing conservation laws also enhances predictive accuracy.
  • Long-term simulations ($t \in [0, 5]$) confirmed strict Hamiltonian conservation.

5. Significance and Conclusion

This work bridges a critical gap between geometric numerical integration and neural network solvers.

• Physical Fidelity: It demonstrates that for neural PDE solvers to be viable for long-term simulations of conservative systems, they must explicitly enforce invariants, not just minimize loss functions.
• Stability: By keeping the solution on the invariant manifold, the method significantly improves the stability and reliability of neural solvers, preventing the "blow-up" or unphysical dissipation often seen in unconstrained neural approaches.
• Future Impact: The framework provides a robust foundation for extending structure-preserving neural methods to more complex dissipative systems (e.g., phase-field models) and stochastic differential equations.

In summary, RP-TENG offers a mathematically rigorous and computationally efficient pathway to train neural networks that respect the fundamental laws of physics, ensuring that the learned dynamics are not only accurate but also physically consistent over long time horizons.