Staged Depth Training: A Representation Curriculum for PINNs

Kejia Zhang; Youran Sun; Haizhao Yang

cs.LGAug 21, 2026

Staged Depth Training: A Representation Curriculum for PINNs

Authors: Kejia Zhang, Youran Sun, Haizhao Yang

Organizations: University of Maryland, College Park

Abstract

Representation quality is a central determinant of PINNs' performance, yet standard training leaves representations to emerge implicitly while fitting the final solution. We introduce representation curriculum, an ordered process in which representations are explicitly learned, transferred independently of their predictors, and progressively refined. We realize it with Staged Depth Training (SDT), which trains a shallow prefix under a temporary physics-informed head, discards the head, and freezes the learned prefix while adding depth, without equation-specific encodings or changes to the final architecture. Across the 20 default forward problems in PINNacle with three backbones, SDT improves 40 of 59 equal-budget problem--backbone cells by at least 5% and remains within that band in the rest, with a 32.8% geometric-mean error reduction on a PirateNet-style backbone. Mechanistic ablations suggest that the gain is not explained by optimizer restarts or shallow warm-starting alone. Representation visualizations and hyperparameter-basin analyses provide diagnostic evidence on representation geometry and local sensitivity to shared hyperparameters. On Poisson--Boltzmann 2D, SDT also more than doubles the fitted depth-scaling exponent for both backbones. These results support representation curriculum as a promising training strategy for improving PINNs while preserving the deployed architecture and inference cost.

Figures & tables

Figure 1: Overview of representation curriculum for PINNs. End-to-end training optimizes the whole hierarchy through the final output, whereas SDT learns a shallow representation under a temporary physics-informed head, freezes and transfers it while new layers train, and finally fine-tunes the full network. Every temporary head is discarded, so the deployed model matches the end-to-end baseline.

	MLP		ResNet		PirateNet-style
Problem	Vanilla	SDT (ours)	Vanilla	SDT (ours)	Vanilla	SDT (ours)
Burgers 1D	$1.426\text{\times}{10}^{-02}$	$1.320\text{\times}{10}^{-02}$	$1.289\text{\times}{10}^{-02}$	$1.313\text{\times}{10}^{-02}$	$1.360\text{\times}{10}^{-02}$	$1.253\text{\times}{10}^{-02}$
Burgers 2D	$5.236\text{\times}{10}^{-01}$	$5.068\text{\times}{10}^{-01}$	$4.915\text{\times}{10}^{-01}$	$5.059\text{\times}{10}^{-01}$	$4.795\text{\times}{10}^{-01}$	$4.868\text{\times}{10}^{-01}$
Poisson 2D (classic)	$6.261\text{\times}{10}^{-01}$	$5.968\text{\times}{10}^{-01}$	$6.722\text{\times}{10}^{-01}$ ‡	$5.975\text{\times}{10}^{-01}$	$6.775\text{\times}{10}^{-01}$	$3.876\text{\times}{10}^{-01}$
Poisson–Boltzmann 2D	$5.723\text{\times}{10}^{-01}$	$4.561\text{\times}{10}^{-01}$	$4.954\text{\times}{10}^{-01}$	$4.643\text{\times}{10}^{-01}$	$5.640\text{\times}{10}^{-01}$	$2.297\text{\times}{10}^{-01}$
Poisson 3D (complex geometry)	$2.872\text{\times}{10}^{-01}$	$2.165\text{\times}{10}^{-01}$	$4.806\text{\times}{10}^{-01}$	$2.527\text{\times}{10}^{-01}$	$4.550\text{\times}{10}^{-01}$	$2.393\text{\times}{10}^{-01}$

Table 1: Relative

L_{2}

error on all 20 default PINNacle forward problems; lower is better. Shading and boldface mark an improvement of at least 5% within a backbone among the 59 equal-budget comparisons, and none favors Vanilla by that margin. The symbol

\dagger

marks the four problems where both strategies stay above 0.9 L2RE for every backbone: retained for coverage, but not interpreted as comparative evidence. The symbol

\ddagger

marks the one run stopped early (5,000 of 20,000 updates) by numerical instability; that cell is excluded from the summary row and from the 5% counts.

Figure 2: Mechanistic ablation on P3D and NS2D: L2RE normalized by the matched Vanilla result (dashed line, lower is better), logarithmic axis. E2E-Restart has no consistent advantage over Vanilla, and SDT is lower in all four settings.

Figure 3: Third-layer ResNet representations on Heat 2D (multiscale): the same 2,400 points under identical t-SNE settings, colored by the percentile of $u(\boldsymbol{z})$ . SDT forms a coherent band whereas Vanilla is locally mixed; $V_{20}$ is 33.1% lower (Appendix E.4 ).

Figure 4: Hyperparameter response surfaces for MLP on Heat 2D (complex geometry): circles mark the evaluated grids, contours are quadratic fits to

\log_{10}\operatorname{L2RE}

. At

\operatorname{L2RE}\leq 0.05

the basin area is 0.673 for Vanilla and 1.203 for SDT.

Figure 5: Depth scaling on Poisson–Boltzmann 2D (log–log axes), three layers added per level. The methods coincide at depth three by construction; beyond it SDT attains lower L2RE at every matched depth.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Staged quantity	Carried across stages	Depth changes	Same final model
Greedy layerwise ( 2 )	network depth	layer weights	yes	yes
Progressive stacking ( 8 )	Transformer depth	copied blocks	yes	yes
Staged Training ( 23 )	width and depth	loss-preserving growth map	yes	yes
RAPTR ( 20 )	active subnetwork	full weights	no	yes
LayerLock ( 5 )	which layers are trainable	full weights	no	yes
PirateNets ( 28 )	effective block depth	full weights	no	yes

Appendix

Table A.1: Comparison with the closest staged or progressive methods. “Depth changes” distinguishes methods that alter the number of layers during training from those that hold a fully instantiated network and vary its loss, its trainable subset, or its inputs. “Same final model” records whether training returns a network architecturally identical to an end-to-end baseline of the same target size.

Problem	$d_{\mathrm{in}}$	$d_{\mathrm{out}}$	Type	$N_{r}$	$N_{b}$	$N_{i}$	$N_{\mathrm{test}}$	L2RE reference
Burgers1D	2	1	TimePDE	8192	2048	2048	8192	Reference data
Burgers2D	3	2	TimePDE	32768	8192	8192	32768	Reference data
Poisson2D-Classic	2	1	PDE	8192	2048	–	8192	Reference data
PoissonBoltzmann2D	2	1	PDE	8192	2048	–	8192	Reference data
Poisson3D-ComplexGeometry	3	1	PDE	32768	8192	–	32768	Reference data
Poisson2D-ManyArea	2	1	PDE	8192	2048	–	8192	Reference data

Appendix

Table B.1: Per-problem sampling and evaluation protocol inherited from PINNacle. The columns

N_{r}

N_{b}

N_{i}

, and

N_{\mathrm{test}}

give nominal interior, boundary, initial, and DeepXDE residual-test counts. The DeepXDE residual-test set is distinct from the reference set used to compute the reported L2RE.

$d$	$o$	Problems	MLP/ResNet	PirateNet-style
2	1	Burgers1D; P2D-Classic; PB2D; P2D-ManyArea; Wave1D; KS	50900	61006
3	1	P3D-Complex; four Heat2D; two Wave2D	51000	61106
3	2	Burgers2D; Gray–Scott	51100	61206
2	3	NS2D-LidDriven; NS2D-BackStep	51100	61206
3	3	NS2D-LongTime	51200	61306
5	1	PoissonND	51200	61306

Appendix

Table C.1: Exact trainable parameter counts of the final depth-six, width-100 models. Vanilla and staged training have the same count in every row.

		Final training physics loss			Staged L2RE endpoints			L2RE ratio
Problem	Backbone	Vanilla	Staged	ratio	S1 ( $D{=}3$ )	final	final/S1	final	med. last 5
Heat2D-Multiscale	MLP	$5.127\text{\times}{10}^{-4}$	$1.029\text{\times}{10}^{-4}$	$0.201$	$0.368$	$1.924\text{\times}{10}^{-2}$	$0.052$	$0.410$	$0.365$
Heat2D-Multiscale	ResNet	$5.014\text{\times}{10}^{-4}$	$6.194\text{\times}{10}^{-5}$	$0.124$	$1.256$	$1.717\text{\times}{10}^{-2}$	$0.014$	$0.376$	$0.401$
Heat2D-Multiscale	PirateNet-style	$7.227\text{\times}{10}^{-4}$	$8.710\text{\times}{10}^{-5}$	$0.121$	$0.082$	$1.890\text{\times}{10}^{-2}$	$0.231$	$0.215$	$0.385$
NS2D-LidDriven	MLP	$3.677\text{\times}{10}^{-4}$	$1.490\text{\times}{10}^{-5}$	$0.041$	$0.731$	$3.406\text{\times}{10}^{-2}$	$0.047$	$0.374$	$0.483$
NS2D-LidDriven	ResNet	$4.679\text{\times}{10}^{-4}$	$1.062\text{\times}{10}^{-4}$	$0.227$	$0.735$	$3.944\text{\times}{10}^{-2}$	$0.054$	$0.760$	$0.806$

Appendix

Table E.1: Diagnostics for the eleven matched pairs. The physics-loss columns use the unit-weighted training loss of Eq. ( 1 ) and are independent of the reference solution; ratios below one favor SDT. “S1” is the staged model at the first stage boundary, so “final/S1” measures how much the depth added after stage one contributes. The last two columns compare the reported final-update ratio with a ratio formed from the median of the last five logged evaluations.

Physics-Informed ML

Physics-informed machine learning combines data-driven models with physical laws to solve, predict, and infer scientific systems.

See all 843 papers

From Simple to Complex: Curriculum-Guided Physics-Informed Neural Networks via Gaussian Mixture Models
May 19, 2026Jianan Yang, Yiran Wang, Shuai Li +3PDE Solving Curriculum Learning
Xi’an Jiaotong University, Xi’an, China · College of International Business, Zhejiang Yuexiu University, Shaoxing, China · Shaoxing Key Laboratory of Intelligent Monitoring and Prevention of Smart Society, Shaoxing, China +1
Curriculum Learning of Physics-Informed Neural Networks based on Spatial Correlation
May 14, 2026Xujia Chen, Xinyue Hu, Letian Chen +2PDE Solving Curriculum Learning
Department of Automation, Tsinghua University, Beijing 100084, P.R. China
Lightweight Geometric Adaptation for Training Physics-Informed Neural Networks
Apr 16, 2026Kang An, Chenhao Si, Shiqian Ma +1Neural Network Optimization Second-Order Optimization
School of Data Science The Chinese University of Hong Kong, Shenzhen Shenzhen, China

May 19, 2026J/K move · Enter open · S save

From Simple to Complex: Curriculum-Guided Physics-Informed Neural Networks via Gaussian Mixture Models

Jianan Yang, Yiran Wang, Shuai Li, Fujun Cao+2

Xi’an Jiaotong University, Xi’an, China · College of International Business, Zhejiang Yuexiu University, Shaoxing, China · Shaoxing Key Laboratory of Intelligent Monitoring and Prevention of Smart Society, Shaoxing, China+1

PDE Solving Curriculum Learning Physics-Informed ML

Open paper AskPDF arXiv

Physics-informed neural networks (PINNs) offer a mesh-free framework for solving partial differential equations (PDEs), yet training often suffers from gradient pathologies, spectral bias, and poor convergence, especially for problems with strong nonlinearity, sharp gradients, or multiscale features. We propose the Curriculum-Guided Gaussian Mixture Physics-Informed Neural Network (CGMPINN), which integrates Gaussian mixture modeling with dynamic curriculum learning. Specifically, a GMM is periodically fitted to the PDE residual distribution to quantify spatially varying learning difficulty. A smooth curriculum schedule progressively shifts training focus from easy to harder regions, while precision-based variance modulation suppresses unreliable clusters during early optimization. This dual curriculum is governed by a shared curriculum parameter and can be combined with self-adaptive loss balancing. We further establish theoretical guarantees, including sublinear convergence of the gradient norm for the induced time-varying loss, uniform equivalence between the curriculum-weighted and standard PDE losses, and a generalization bound with an explicit weighting-induced bias characterization. Experiments on six benchmark PDEs spanning elliptic, parabolic, hyperbolic, advection-dominated, and nonlinear reaction-diffusion types show that CGMPINN consistently achieves the lowest relative

L_2

and maximum absolute errors among all compared methods, reducing relative

L_2

error by up to 97.8% over the standard PINN at comparable cost. Our code is publicly available at https://github.com/Mathematics-Yang/CGMPINN.

Abstract

Physics-Informed ML

Explore similar work

From Simple to Complex: Curriculum-Guided Physics-Informed Neural Networks via Gaussian Mixture Models

Curriculum Learning of Physics-Informed Neural Networks based on Spatial Correlation

Lightweight Geometric Adaptation for Training Physics-Informed Neural Networks