Bifurcations from Weight Space: JEPAWG on Parametric Dynamical Systems
Authors: Tobias Göbel, Miranda C. N. Cheng
Organizations: Institute for Mathematics, Academia Sinica, Taipei, Taiwan · Department of Physics, National Taiwan University, Taipei, Taiwan · Institute of Physics, University of Amsterdam, Amsterdam, Netherlands
Trained network weights are increasingly treated as data, but there is rarely ground truth about what a population of models should encode, which hampers systematic validation of weight-space methods. This makes parametric dynamical systems a natural test-bed: the dynamics change qualitatively at submanifolds of parameter space that are, in well-studied families, known analytically or numerically. Fine-tuning a forecasting network at each point of a parameter grid yields a model zoo indexed by a physical parameter, against which weight-space diagnostics can be checked. We build such model zoos for four dynamical systems and analyse them with JEPAWG, a recently proposed joint-embedding hypernetwork built around a learned parameter--weight representation. It was first introduced for neural samplers for lattice field theories, where its latent geometry located the phase transition and estimated the 2D Ising critical exponent. Applied to dynamical systems, JEPAWG is competitive on forecasting with a hypernetwork trained directly for that task, despite being supervised on the parameter--weight correspondence rather than on forecasts. Beyond forecasting, the norm of the parameter-derivative of the decoded weights locates Hopf points and curves, tracks a period-doubling cascade closely enough to give a qualitative estimate of the Feigenbaum constant, and detects a global bifurcation, which has no eigenvalue signature and is therefore invisible to linearisation. As such, JEPAWG provides interpretable insight into the dynamics that is not available from other hypernetworks. Together, these results indicate that JEPAWG is a workable instrument for reading physical structure off a population of trained networks.
Figures & tables
Figure 1: JEPAWG pipeline. (1) A conditional forecaster is trained on trajectories from the parametric family and fine-tuned at each of M parameter values, yielding the weight grid {(pi,θi)}i=1M (Section 4.1 ). (2) Encoders, predictor and decoder are trained in two stages (Eqs. 4 , 5 ), giving the inference path θ^(p)=D(P(Ep(p))) . (3) The decoded weights are evaluated as a hypernetwork and probed via G(p)=∥∇pθ^(p)∥ , verified against known bifurcations.
system
split
PHLieNet
JEPAWG
VAE
AE
Lorenz, M=800
interp ρ∈[10,35]
20.23±0.89
20.19±0.15
18.69±0.27
18.31±0.64
below ρ∈[5,9]
15.53±4.18
3.61±0.25
5.58±0.81
4.31±1.82
above ρ∈[36,40]
2.04±0.22
2.03±0.09
2.48±0.11
2.21±0.22
Rössler, M=1000
interp c∈[3,9]
50.16±1.93
47.42±0.90
35.84±2.58
30.10±2.49
below c∈[1.8,2.8]
6.98±2.25
42.66±12.53
27.64±12.20
10.44±8.30
above c∈[9.2,10.2]
15.72±1.71
9.06±0.92
7.86±0.51
6.09±0.54
Table 1: Forecasting performance: TtT0.2 on held-out interpolation parameters and on each side of the training range, mean ± std over 5 seeds. Bold: best mean and all entries whose interval overlaps the best entry’s. Rows in which all entries overlap carry no bold. Continuous-time systems in time units, the discrete Hénon map in iterations; Hénon extrapolation is omitted: for discrete maps one inaccurate iteration exits the training region, so TtT collapses to ≈0 for every method. Per-experiment transient handling is given in Table 6 and the full test sets in Table 7 .
Figure 2: Extrapolation TtT0.2 of JEPAWG and PHLieNet, by regime beyond the training range. Left: limit cycles that keep changing; for Rössler the cycle shrinks as c decreases below the training range ( c<3 ), for the Brusselator it grows and moves as b increases above it ( b>4 ). Middle: a stable fixed point, with initial transients included for Lorenz and excluded for the Brusselator. Right: chaotic dynamics.
JEPAWG
PHLieNet
side
regime
λmax
TtT0.2
bounded (%)
TtT0.2
bounded (%)
Rössler, c∈[1.8,2.8]
limit cycle
0
42.7±12.5
66±8
7.0±2.3
49±11
Brusselator, b∈[4.2,4.8]
limit cycle
0
5.7±3.3
100±0
3.2±1.1
70±27
Lorenz, ρ∈[5,9]
stable fixed pt.
−0.65
3.6±0.3
54±7
15.5±4.2
97±6
Brusselator, b∈[1.1,1.4]
stable fixed pt.
−0.58
0.8±0.2
0±0
5.4±6.0
0±0
Rössler, c∈[9.2,10.2]
chaotic
0.11
9.1±0.9
100±0
15.7±1.7
93±7
Table 2: Extrapolation by regime of the true dynamics beyond the training range: TtT0.2 (short-horizon tracking) and bounded fraction (share of rollouts that stay within 1.5× the extent of the true test trajectories over the full horizon), mean ± std over 5 seeds of per-side means; bold as in Table 1 . Intervals as in Table 7 ; λmax is the largest Lyapunov exponent over the range, 0 on limit cycles and the leading Jacobian eigenvalue at fixed points. JEPAWG leads on both measures where the limit cycle keeps changing; PHLieNet holds a fixed point that JEPAWG loses on Lorenz; in chaos the measures split: PHLieNet tracks longer for Rössler with c∈[9.2,10.2] , but JEPAWG’s rollouts always stay bounded.
Figure 3: Lorenz. (a) Raw-weights derivative ∥dθ/dρ∥ across the ρ -sweep, with a Gaussian-smoothed version as the thick line (top); derivative G(ρ)=∥dθ^/dρ∥ of the JEPAWG-decoded ( M=800 ) and PHLieNet-generated weights, 5 seeds each, with their mean as the thick curves (middle); the ground-truth bifurcation diagram (bottom). The Hopf point ρH , at which the fixed points lose their stability, is shown as a red dashed line in all panels. (b) Weight-space response G(σ,ρ)=∥∇σ,ρθ^∥ for a zoo trained on (σ,ρ)∈[5,20]×[10,40] ( M=1024 , β=8/3 ), shown for ρ≤35 because the finite-difference derivative is unreliable at the edge of the training range; the analytic Hopf curve is in red, and the blue curve marks the level set λmax=0 of the largest Lyapunov exponent. (c) Zoom above the Hopf curve with parameter ranges (σ,ρ)∈[8.5,12]×[28.5,33.5] ( M=1024 , β=8/3 ): a ring of elevated response surrounds the marked Bykov T-point at (σ,ρ)≈(10.17,30.87) . White lines: the principal homoclinic curves of the origin, obtained independently of the model by integrating its unstable manifold Wu(0) on a 6002 grid and locating the parameter values at which its winding number around the nearest equilibrium before escaping to the other wing jumps by one turn. They form a single spiral winding onto the T-point [ 20 , 21 ] ; the ring of G lies on its innermost turns, where successive curves are closest, while no fixed point changes stability anywhere in the window.
Figure 4: Brusselator. (a) Weight-space response G(a,b)=∥∇a,bθ^∥ with the analytic Hopf curve b=1+a2 in red; white dots mark the M=1024 Sobol sampled training parameters. (b) Latent space (PCA-2D with 86% explained variance), coloured by the signed distance s=b−(1+a2) to the curve; circled points ( ∣s∣<0.12 ) correspond to the parameters nearest to the curve and trace a coherent boundary in every seed. Parameter ranges are (a,b)∈[0.5,2]×[1,5] ( M=1024 ).
Figure 7
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
grid
method
TwoNN
k=5
k=10
k=20
Lorenz
raw θ
19.9
17.9
15.1
12.3
( dimp=1 )
PCA-16
4.7
4.1
3.5
3.3
AE
7.9
7.3
6.5
5.7
VAE
11.4
10.4
9.7
8.9
JEPAWG
2.9
2.3
1.7
1.4
Rössler
raw θ
2.8
8.9
14.2
17.1
Appendix
Table 4: Intrinsic-dimension estimates across neighbourhood scale (TwoNN, then Levina–Bickel with k=5,10,20 ) for the weight grids of four families and their embeddings (all bottleneck dz=16 ; PCA: rank 16). The JEPAWG latent dimension is closest to dimp in 15 of 16 columns (exception: TwoNN on Rössler); PCA, AE and VAE are fitted without access to p .
TtT0.2
one-step MSE
system
base θ0
raw θi
decoded θ^
base θ0
raw θi
decoded θ^
Lorenz
20.64
15.79
20.19±0.15
1.8e-3
2.5e-3
6.7e-4
Rössler
42.08
15.33
47.42±0.90
2.1e-3
6.6e-4
1.1e-4
Brusselator
17.44
24.25
25.40±0.51
9.7e-4
6.8e-4
8.6e-5
Hénon
856
733
772±85
9.6e-6
1.3e-5
1.5e-5
Appendix
Table 5: Shared p -conditional base network, raw fine-tuned grid weights and JEPAWG-decoded weights on the held-out interpolation test sets. The base θ0 is the single conditional TCNN-CD every zoo member was fine-tuned from; raw weights use the nearest grid point’s θi ; both are deterministic given the zoo. Decoded TtT entries are mean ± std over 5 hypernetwork seeds; decoded MSE entries are seed means. TtT0.2 in time units (Hénon: iterations); one-step MSE in the standardised training space. The decoder has the lowest one-step error on the three continuous systems and recovers the Rössler horizon the per-parameter fine-tunes lose; on Hénon fine-tuning gives no improvement over the base.
Figure 10
Figure 8: Random parameters. k random coordinates ξ∼U(0,1)k , with no effect on the dynamics, are appended to p=(ρ,ξ) on the Lorenz zoo ( M=800 , weights unchanged), and both models are retrained. Shown is G(ρ) evaluated at a fixed random draw of ξ , averaged over 5 seeds; the red line marks ρH≈24.74 . (a) JEPAWG; the curves for k=0,2,5,10 coincide. (b) the random-feature control, with the same architecture but encoders frozen at initialisation. As k grows, its baseline rises, the Hopf peak broadens and loses prominence.
Figure 9: Predicted and true trajectories at the furthest extrapolation parameter in each regime (one seed). In the limit-cycle regimes, JEPAWG-generated forecasters follow the true cycle more closely, while PHLieNet-generated trajectories tend to diverge after a number of steps, as quantified by the bounded fraction in Table 2 .
Figure 10: Forecast trajectories of the Lorenz system at held-out ρ (solid: truth; dashed: JEPAWG) at σ=10 . From left to right: spiral trajectory into a stable fixed point, slowly decaying spiral just below the subcritical Hopf point, and the chaotic attractor. At these ρ inside the training range, the forecasts follow the true trajectory with reasonable precision.
Figure 11: Rössler weight-space response at M=1000 ( a=b=0.2 , c∈[3,9] ). Top: the response G(c) of JEPAWG-generated weights (Gaussian smoothed), per-seed curves and seed mean. Bottom: ground-truth diagram (successive local maxima of x(t) per parameter). The peaks of G correspond to features of the bifurcation diagram in the plot below: the 2→4 doubling at c=3.83 and the entry/exit boundaries of the periodic windows show clear signal; the finite parameter resolution limits the visibility of the doublings closer to chaos.
Figure 12: Hénon u -cascade. Top: weight-space response G(u) of the cascade zoo, 5 seeds in colour and their mean in black, over the full trained range u∈[−4.5,−0.12] (Fig. 5 shows the part u≥−3.2 ). Bottom: bifurcation diagram of the Hénon map ( b=0.3 ) in the same coordinate u=log10(a∞−a) , with a∞=1.05805 [ 33 ] . For each of 6000 parameter values, 512 iterates of the attractor are shown after a 2×105 -iterate transient. Red dashed lines mark the period-doubling bifurcations a1,…,a7 (labels on the top axis); a1 is exact, a2,…,a7 were located by bisection on the attractor period. The last two splittings of the orbit points are narrower than 10−3 in x and are not resolved at this scale. G(u) peaks at each of the first five doublings, shows a step at a6 and no feature at a7 , where the seeds begin to disagree.
Figure 13: Hénon. Left: ground-truth orbit period over (a,b) ; the boundaries between coloured bands are the bifurcation curves. Right: JEPAWG weight-space response G(a,b)=∥∇a,bθ^∥ on the same rectangle, mean over seeds. The ridge of G tracks the interior doubling curves and the chaotic band.
exp.
M
trans.
Δt
len (tr/te)
ICs
h
Δp
λmaxsup
Forecasting
Lorenz, ρ
800
0
0.02
20/30
5
3.1×10−2
3.1×10−2
1.06
Rössler, c
1000
300
0.1
20/100
10
7.5×10−3
6.0×10−3
0.09
Brusselator, b
1024
20
0.05
30/30
5
3.1×10−3
8×10−4
≤0
Hénon, a
1024
200 it.
1
800/1500
5
3.1×10−4
2.4×10−4
0.18
Interpretability
Appendix
Table 6: Per-experiment protocol. The M parameter values are drawn from scrambled Sobol sequences (seed 0; stratified per-band draws for the Brusselator b sweep and the Hénon u -cascade, joint draws for the two-parameter families). Transients are integrated in addition to the recorded lengths; a zero entry means transients are retained in the data. Δt : observation interval (maps: 1 iteration); recorded lengths as training/test. ICs: initial conditions per training parameter (test sets use 10; Lorenz 20). h : evaluation step of the G sweep (two-parameter rows list the per-axis pair (h1,h2) in the order of the experiment label), alongside the mean spacing Δp of the fine-tuned parameters (Brusselator b : spacing in the dense band b∈[2.35,2.85] around bc ; overall mean 2.4×10−3 ) ( u -cascade: both in u=log10(a∞−a) , on which it is trained). λmaxsup : largest Lyapunov exponent over the training range (per time unit; maps: per iteration); ≤0 marks ranges without chaotic parameters. Input noise 5% of the standardised feature scale. The Lorenz- ρ grid also underlies Fig. 3(a) .
experiment
train range
interp. test
extrapolation
Lorenz, ρ
[10,35]
10 fixed
36 – 40 ; low 5 – 9
Rössler, c
[3,9]
10 Sobol
1.8 – 2.8 ; 9.2 – 10.2
Brusselator, b
[1.5,4]
32 Sobol
1.1 – 1.4 ; 4.2 – 4.8
Hénon, a
[0.85,1.10]
32 Sobol
—
Appendix
Table 7: Interpolation and extrapolation test sets for the forecasting experiments. Interpolation test parameters are Sobol draws from the training range, except Lorenz- ρ , which uses 10 fixed values. Fixed parameters: Lorenz σ=10 , β=8/3 ; Rössler a=b=0.2 ; Brusselator a=1.25 ; Hénon b=0.3 .
Lattice field theory is the workhorse of non-perturbative physics, used to simulate phenomena from the strong nuclear force to critical phenomena in materials. Its Boltzmann distributions are parametrized analytically by coupling constants, but these bare parameters are weak predictors of observables -- extracting physics typically requires extensive simulation. While normalizing flows have emerged as effective samplers at fixed couplings, it remains difficult to interpret what these networks have learned. This raises a natural question: can the physics be read off directly from the flow network parameters themselves, and can those parameters be generated for unseen theories? We propose lattice field theory as a testbed for neural network interpretability: because the target physics is qualitatively well-understood and smoothly varying, it provides ideal synthetic data with known ground truth. To this end, we introduce JEPAWG, a Joint-Embedding Predictive Architecture-based Weight Generator that maps couplings directly to flow weights via a learned latent space. On a scalar theory at lattices of size 62 to 112, the JEPAWG latent space recovers the correct intrinsic dimension of the underlying manifold, locates the phase transition, and encodes a finite-size shift aligned with the 2D Ising exponent ν≈1, allowing us to uncover physical structure by studying the network weights alone. This suggests the fascinating idea of treating the network weights as a new type of physical observable. As a generator, JEPAWG also interpolates and extrapolates to unseen couplings effectively and remains robust to weight-space discontinuities introduced by multi-seed training data, outperforming PCA, AE, and VAE baselines.
Tobias Göbel, Julian R. Ebelt, Zier Mensch +2
Institute of Physics, University of Amsterdam, Amsterdam, Netherlands · Department of Physics, National Taiwan University, Taipei, Taiwan · Department of Physics, Massachusetts Institute of Technology, Cambridge, Massachusetts, USA +2
Deep learning surrogates have become powerful tools for simulating and forecasting complex dynamical systems, yet their utility remains limited by catastrophic error accumulation during long-term autoregressive rollouts. This behavior is partly tied to the nature of the underlying systems: chaotic spatiotemporal systems visit phase space unevenly, with dynamics dominated by recurrent, low-dimensional quiescent states and characterized by rare and dynamically complex regime transitions. Trained under a sample-wise uniform objective, standard neural surrogates allocate their finite capacity to the statistically more numerous low-dimensional quiescent states, systematically under-representing the transient regimes that trigger disproportionate, localized errors. Existing imbalanced-regression methods tackle this issue by reweighting samples according to target-space density. However, statistical target-space rarity does not coincide with the intrinsic dynamical rarity encoded in the recurrence geometry of the attractor. To address this, we introduce Dynamics-Aware Weighting (DAW), a data-centric objective reweighting framework. Using the local dimension d from dynamical systems theory as an a priori measure of a state's dynamical complexity, DAW reshapes the loss landscape to allocate representational capacity toward the sparse, high-d regimes where forecast errors are systematically large. On the chaotic KS equation, DAW consistently outperforms uniform training as well as weighting based on target-space rarity, and its randomly permuted ablation, reducing long-term autoregressive error relative to all baselines. Event-level analysis shows that DAW achieves this by suppressing the localized error amplifications incurred during sharp jumps in the local dimension d, which typically accompany complex physical processes such as wave-merging in the KS system.
Zhou Fang, Gianmarco Mengaldo
Department of Mechanical Engineering National University of Singapore Singapore
JEPA-family world models use a static predictor whose weights do not adapt when test-time dynamics diverge from training. We compare two mechanisms for incorporating accumulated experience into a JEPA predictor under distribution shift: operand-side injection, where a compressed experience representation is added as a residual to the predictor's hidden state (EI-JEPA), and operator-side modulation, where the same representation generates low-rank weight deltas via LoRA applied to the predictor's weights (EPM-JEPA). On a pre-registered comparison (Moving MNIST, gravity shift), EPM-JEPA (D_shift^{n=50} = 0.7848 +/- 0.0078, three seeds) differs from EI-JEPA (0.8238) by delta = 4.74% - Outcome C: a null result - by our stated criterion, a valid outcome. As a secondary, non-pre-registered observation, EPM-JEPA improves 1.90% over a no-memory baseline (0.8000), consistently across seeds, while EI-JEPA underperforms the baseline, indicating the benefit is specific to weight-level modulation. Our primary contribution is a mechanism analysis: the D_shift^{n=50} trajectory reflects three independent dynamical processes - buffer cycling, EMA target drift, and an intrinsic LoRA settling transient of +0.021 - rather than convergence to equilibrium. These findings motivate PEM-JEPA, a physics-grounded successor addressing this dynamical-peak limitation.
Vedant Pandya
School of Artificial Intelligence and Data Engineering (SAIDE) Indian Institute of Technology Jodhpur