Stage Advantage (SA) in BO-SA-PINNs

• Stage Advantage (SA) is a self-adaptive mechanism within BO-SA-PINNs that dynamically rebalances loss terms and redistributes collocation points to target high-error regions in PDE solutions.
• It employs an EMA-based adaptation of loss weights and residual-driven adaptive refinement (RAR-D) to incrementally improve sample efficiency and reduce L2 error significantly.
• Empirical results on benchmarks like Helmholtz and Maxwell demonstrate that integrating SA can lower L2 errors by up to 60% compared to static or uniform sampling approaches.

Stage Advantage (SA) refers to the self-adaptive Stage 2 mechanism within the BO-SA-PINNs (Bayesian Optimization–Self-Adaptive Physics-Informed Neural Networks) multi-stage framework for training PINNs to solve partial differential equations (PDEs). The SA stage dynamically optimizes both the relative weights of PINN loss terms and the spatial distribution of collocation points, leveraging observed loss statistics and residual-driven adaptive sampling. This enables targeted improvement in regions of high modeling error and addresses the limitations of static hyperparameter choices and uniform domain sampling (Zhang et al., 14 Apr 2025).

1. Role and Motivation for the Stage Advantage Mechanism

The Stage Advantage (SA) mechanism is the dedicated Stage 2 component of the BO-SA-PINNs architecture, situated between the initial Bayesian Optimization (BO)-guided pre-training and the final L-BFGS-based refinement. Its dual objectives are: (i) to dynamically re-balance the weights of the loss terms corresponding to PDE residual, boundary, initial, and data constraints, focusing optimization on the "hard" parts of the solution; and (ii) to adaptively re-distribute collocation points in the computational domain, intensifying sampling in regions with high PDE residual while maintaining global domain coverage. This approach addresses two common failure modes in PINN frameworks: (a) suboptimal static loss weights starving some constraints of gradient flow, and (b) uniform sampling failing to resolve sharp fronts, boundary layers, or localized features (Zhang et al., 14 Apr 2025).

2. Algorithmic Structure of Stage Advantage

After Bayesian Optimization and a 500-iteration ADAM warm-start, the Stage Advantage mechanism executes an additional 4,000–5,000 ADAM steps. Each iteration involves the following pipeline:

1. Observation of instantaneous loss values for the PDE residual ($\mathcal{L}_R^{(t)}$), boundary ($\mathcal{L}_B^{(t)}$), and initial condition ($\mathcal{L}_I^{(t)}$).
2. Exponential Moving Average (EMA)-based update of the respective loss term weights ($\omega_R$, $\omega_B$, $\omega_I$).
3. Periodic execution of Residual-based Adaptive Refinement with Distribution (RAR-D) to augment the collocation set with new points in error-prone regions.
4. Application of updated weights and sample set in subsequent iterations.

This continues until a convergence criterion or a fixed iteration budget is reached, after which the adapted network proceeds to Stage 3 (L-BFGS training) (Zhang et al., 14 Apr 2025).

3. EMA-Based Loss-Weight Adaptation

The Stage Advantage mechanism employs EMA on both raw losses and loss weights, parameterized by $\beta$ and $\gamma$. For each loss component ($m^{(t)}_R$, $m^{(t)}_B$, $m^{(t)}_I$), EMA updates are applied as:

$$m_X^{(t)} = \beta m_X^{(t-1)} + (1-\beta) \mathcal{L}_X^{(t)}, \quad X \in \{R, B, I\}$$

Provisional loss weights are then computed via normalization:

$\omega'_X^{(t)} = \frac{m_X^{(t)}}{m_R^{(t)} + m_B^{(t)} + m_I^{(t)}}$

The final loss weights are further updated using a second EMA:

$\omega_X^{(t)} = \gamma \omega_X^{(t-1)} + (1-\gamma) \omega'_X^{(t)}$

All weights are clamped to a predefined interval (for example, [0.01, 0.25]) to prevent degeneration. The hyperparameters are set to $\beta = \gamma = 0.999$ (Zhang et al., 14 Apr 2025).

4. Residual-Based Adaptive Refinement with Distribution (RAR-D)

The RAR-D algorithm systematically augments the domain point set to resolve high-error regions. Each cycle:

1. Samples a large candidate pool (e.g., 1,000 points) uniformly within the problem domain.
2. Evaluates the squared residual $r(x_j) = |\mathcal{N}^*(x_j; \theta) - f(x_j)|^2$ at each $x_j$.
3. Normalizes the residuals and forms a discrete probability distribution over the candidate set.
4. Draws a fixed number of new collocation points (e.g., 50) i.i.d. from this probability law.
5. Appends selected points to the current training set.

This RAR-D procedure can repeat in tandem with ADAM updates, for example every 50 steps, and typically for 100 cycles during Stage 2 (Zhang et al., 14 Apr 2025).

5. Quantitative Impact and Ablation Studies

Ablation experiments demonstrate the Stage Advantage’s contribution to accuracy and sample efficiency:

Model Variant	2D Helmholtz L₂ Error	2D Maxwell (ε_c=1.5) L₂ Real/Imag
BO-SA-PINN (full SA)	$3.21 \times 10^{-4}$	$(1.74 \times 10^{-2}, 1.17 \times 10^{-2})$
No EMA	$7.34 \times 10^{-4}$	$(1.85 \times 10^{-2}, 1.24 \times 10^{-2})$
No SA (no EMA, no RAR-D)	$7.75 \times 10^{-4}$	$(2.03 \times 10^{-2}, 1.58 \times 10^{-2})$

The results show that EMA re-weighting alone reduces error by 5–15%, while full SA (EMA + RAR-D) provides an additional 10–20% improvement over BO-only pre-training. In all cases, inclusion of the Stage 2 SA mechanism lowers the ultimate $L_2$ error by 30–60% for a fixed total iteration budget, as illustrated by the reduction from $L_2 \approx 7.75 \times 10^{-4}$ (no SA) to $L_2 = 3.21 \times 10^{-4}$ (with SA) in the Helmholtz case (Zhang et al., 14 Apr 2025).

6. Pipeline Flow and Practical Considerations

The practical pipeline for Stage Advantage involves initialization with the pre-trained model, ADAM-based iterations with dynamic weights and adaptive sampling, and return of both the self-adapted model and the updated sampling set at the conclusion of the stage. The key algorithmic flow:

• Input: $\mathcal{N}^*$ (pre-trained), sampling sets, initial weights from BO.
• For $t = 1$ to $T_\mathrm{SA}$:
  1. Perform ADAM step with current weights and samples.
  2. Update losses and compute EMA.
  3. Update and clamp weights.
  4. Every fixed $k$ steps: run RAR-D, expand sampling set.

• Output: Self-adapted model, updated training points, latest weights.

This methodology ensures the PINN is continually refocused onto underfit solution regions and poorly enforced constraints, while maintaining computational efficiency by reusing the existing ADAM training budget (Zhang et al., 14 Apr 2025).

7. Significance and Limitations

The Stage Advantage mechanism, by jointly leveraging EMA-driven loss re-weighting and RAR-D adaptive sampling, constitutes a systematic approach for mitigating the central bottlenecks in PINN training—namely, gradient starvation and poor domain coverage. These improvements translate into lower iteration counts and order-of-magnitude reductions in $L_2$ error for canonical PDE benchmarks, including Helmholtz and Maxwell equations. A plausible implication is that SA mechanisms could generalize to other operator-learning and scientific machine learning frameworks facing analogous challenges. However, explicit limitations or trade-offs are not discussed, leaving open questions of scalability, computational overhead, and effectiveness for highly irregular or discontinuous solutions (Zhang et al., 14 Apr 2025).