Bifurcations from Weight Space: JEPAWG on Parametric Dynamical Systems

AI summary

Overview

Research area: Machine learning for the physical sciences; specifically weight-space learning, hypernetworks, and data-driven detection of bifurcations in parametric dynamical systems.

Technical level: Advanced. The paper assumes familiarity with dynamical-systems theory (Hopf bifurcations, period doubling, invariant manifolds, global bifurcations), joint-embedding predictive architectures, and hypernetwork training.

Scope: The paper builds parametric model zoos for four dynamical systems and tests whether JEPAWG, a joint-embedding hypernetwork mapping physical parameters to network weights, can both forecast at unseen parameters and expose known bifurcations through the parameter derivative of its decoded weights.

What This Paper Is About

Trained network weights are increasingly studied as data, but there is rarely a ground truth for what a population of models should encode, which makes weight-space methods hard to validate. The authors argue that parametric dynamical systems are an ideal test-bed, because the dynamics change qualitatively on bifurcation submanifolds whose locations are known analytically or numerically for well-studied families. They build model zoos for four such systems and ask whether JEPAWG's learned parameter-to-weight map is competitive at forecasting and, more importantly, whether its weight-space geometry localises the known bifurcations.

Key Contributions

1. A weight-space benchmark with independent ground truth. The authors construct four parametrised model zoos — Lorenz (M = 800), Rössler (M = 1000), Brusselator (M = 1024), and Hénon (M = 1024) — each with a known bifurcation structure to validate weight-space diagnostics against: the Brusselator Hopf curve b_c = 1 + a², the Lorenz Hopf locus ρ_H(σ) at β = 8/3, and the period-doubling cascades of the Rössler oscillator and Hénon map, whose accumulation should recover the Feigenbaum constant δ ≈ 4.6692.

2. Detection of local bifurcations from weight space. The norm of the parameter derivative of the decoded weights peaks where the dynamics are most parameter-sensitive. For the Lorenz system (σ = 10, β = 8/3) the peak lies at ρ = 24.51 ± 0.23, inside the bistable window preceding the subcritical Hopf point ρ_H ≈ 24.74, and it tracks ρ_H(σ) when both parameters are swept. The response also locates the Brusselator's supercritical Hopf curve and resolves the first five Hénon period doublings accurately enough for a qualitative estimate of the Feigenbaum constant.

3. Detection of a global bifurcation. The same derivative detects the Lorenz Bykov T-point, a global bifurcation at which every fixed point retains its stability and which therefore has no eigenvalue signature — invisible to linearisation and traditionally requiring invariant-manifold computation or attractor-level scans.

4. The decoded map as a regulariser of the weight grid. Fine-tuned weights carry per-run training noise that dominates their raw parameter derivative; the decoder, fitted on a fixed parameter embedding, does not reproduce this noise. Decoded weights forecast accurately for longer than the raw fine-tuned weights on all four systems, and the JEPA-style latent is organised by dynamical regime, has an intrinsic dimension close to dim p, and is insensitive to parameters that do not enter the dynamics.

Main Findings

• Competitive interpolation forecasting without a forecasting loss. JEPAWG is tied with the end-to-end PHLieNet hypernetwork on three systems (Lorenz: 20.19 ± 0.15 vs 20.23 ± 0.89; Rössler: 47.42 ± 0.90 vs 50.16 ± 1.93; Brusselator: 25.40 ± 0.51 vs 27.34 ± 0.23) and within two time units on the Brusselator, despite being supervised only on the parameter–weight correspondence. It beats the VAE and AE baselines on all four systems, including a clear margin on Hénon (772 ± 85 vs 594 ± 90 and 471 ± 116).

• Lorenz Hopf localisation. The weight-space response G(p) peaks at ρ = 24.51 ± 0.23, inside the bistable window preceding the subcritical Hopf point at ρ_H ≈ 24.74, and tracks the analytic curve ρ_H(σ) = σ(σ+β+3)/(σ−β−1) in the two-parameter (σ, ρ) plane.

• Global bifurcation detected with no eigenvalue signal. The Lorenz Bykov T-point is registered by the same derivative, demonstrating that information inaccessible to a Jacobian computation can be read off the geometry of JEPAWG-generated weights.

• Period-doubling cascade resolved. In the Hénon map, the first five period doublings are resolved accurately enough for a qualitative estimate of the Feigenbaum constant δ ≈ 4.6692; the Rössler cascade is tracked closely enough for a qualitative estimate as well.

• Decoded weights beat raw fine-tuned weights on forecast horizon. The denoised decoded weights achieve better time-to-threshold than the raw weights they were trained to reconstruct on all four systems, most strongly on Rössler — though not always better than the shared conditional base network.

• Extrapolation is regime-dependent. JEPAWG extrapolates better where dynamics feature moving limit cycles (Rössler with c ∈ [1.8, 2.8]: 42.66 ± 12.53 vs PHLieNet's 6.98 ± 2.25), while PHLieNet is better where the dynamics settle onto a fixed point (Lorenz with ρ ∈ [5, 9]: 15.53 ± 4.18 vs 3.61 ± 0.25). In the chaotic Rössler range c ∈ [9.2, 10.2], PHLieNet leads on time-to-threshold (15.72 ± 1.71 vs 9.06 ± 0.92) but keeps bounded rollouts slightly less often (93 ± 7% vs 100%). For Lorenz with ρ ∈ [36, 40], all methods reach threshold after about two to three Lyapunov times.

• Accurate forecasts do not imply a usable weight-space signal. PHLieNet matches JEPAWG on forecasting yet its generated weights show no visible response at the Lorenz Hopf point, because its parameter-to-weight map is constrained only by forecast accuracy.

Methodology in Plain English

The pipeline has three stages.

Build the model zoo. For each dynamical system, the authors sample M parameter values with scrambled Sobol sequences, generate trajectories (RK4 at solver step δt = 10⁻³ for continuous systems, direct iteration for discrete maps), and train a single parameter-conditional temporal convolutional network with causal dilated convolutions. That network takes the parameter appended to each state as an extra input channel, and is trained with 5% additive Gaussian input noise on the input window only, so it learns to contract its own small errors during iterated rollout. This shared base network is then fine-tuned separately at each parameter value for 30 epochs with Adam at learning rate 10⁻³, producing a family of weight vectors indexed by the physical parameter. The common initialisation matters: the loss is invariant under a large symmetry group (channel permutations, rescalings across ReLU layers), so independently trained networks at neighbouring parameters would differ mostly by which symmetry copy they landed in rather than by the dynamics — a common start acts as an approximate gauge fixing.

Fit JEPAWG. The parameter–weight pairs are treated as two views of the same object. A weight encoder, a parameter encoder, and a latent predictor are trained jointly with a JEPA/VICReg objective (inverse-variance and variance terms, with the covariance term switched off), with latent dimension d_z = 16 chosen deliberately larger than the parameter dimension so that any concentration onto fewer dimensions is learned rather than imposed. A decoder back to weight space is then trained with the encoders and predictor frozen. At inference, weights at a new parameter come from composing parameter encoder, predictor, and decoder. Encoder and decoder have two hidden layers of 512; the parameter encoder and predictor each use a single width-64 hidden layer; the parameter input includes 40 sinusoidal encodings per axis (128 for the Hénon u-cascade). Training uses Adam at 3 × 10⁻⁴, batch size 256, 5000 steps, and 5 seeds per configuration.

Evaluate two ways. As a hypernetwork, decoded weights are tested at unseen parameters using time-to-threshold: the longest initial duration over which the initial-condition-averaged NRMSE stays below 0.2, computed in original state units and averaged over 5 seeds. As a probe of the dynamics, the authors compute G(p) = ‖∇_p θ̂(p)‖_F by central differences on a dense grid independent of the fine-tuning grid, averaged over seeds, and ask whether it peaks at the known bifurcations. The intuition: the forecaster is fitted on the attractor's invariant measure, which reorganises when the attractor changes, so the weights pinned to its support must change rapidly there. Differentiating the decoded weights rather than the raw ones suppresses per-run training noise.

Why This Matters

Impact on research. Weight-space methods have lacked a setting with independent ground truth. Parametric dynamical systems supply exactly that: cheap data generation plus an analytically or numerically known partition of parameter space. This turns "does this weight-space diagnostic work?" from an open-ended question into a testable one, and gives hypernetwork research a benchmark where the correct answer is known in advance. The result that a global bifurcation with no eigenvalue signature is legible in weight space is a genuinely new capability relative to Jacobian-based and equilibrium-based data-driven methods, several of which explicitly exclude global bifurcations from scope.

Real-world applications:

• Scientific model management: maintaining and auditing large families of simulation or surrogate models indexed by physical parameters, where a bifurcation flag from weights alone would signal that a model needs retraining or replacement.
• Control and safety monitoring: detecting when a controlled physical system (power grids, chemical reactors modelled by the Brusselator, fluid systems) approaches a qualitative change in behaviour, including transitions that linear stability analysis misses.
• Reduced-order modelling in engineering: using a parameter-to-weights map as a compact, differentiable surrogate across a design space, with the weight derivative serving as a cheap sensitivity indicator.
• Climate and geophysical model ensembles: organising populations of trained models by dynamical regime rather than by raw weight distance, using a latent whose intrinsic dimension tracks the physical parameter dimension.

Industry relevance. The approach offers a compression-and-denoising story: decoded weights outperform the noisy fine-tuned weights they were trained to reconstruct, meaning a learned parameter-to-weight map can act as a regulariser over a model zoo. For organisations holding many fine-tuned checkpoints, this suggests a way to store, interpolate, and quality-check model populations with an interpretable low-dimensional latent.

Future Directions

• Closing the extrapolation gap. JEPAWG and PHLieNet fail in complementary regimes — moving limit cycles versus fixed points, chaotic versus periodic exteriors. A hybrid that constrains decoded weights outside the training range, or an architecture that combines JEPAWG's direct parameter–weight supervision with PHLieNet's anchor-based saturation, is a natural next step.

• Pushing resolution limits. Detection resolution is bounded below by the average parameter spacing in the zoo, so transitions closer together than that effective width cannot be separated. How the method scales with grid density, and whether active sampling near detected peaks can refine them, is unresolved.

• Quantitative rather than qualitative constants. The cascade analysis gives only a qualitative estimate of the Feigenbaum constant. Turning G(p) peaks into calibrated numerical estimates — as the original lattice field theory application did for the 2D Ising critical exponent — would strengthen the case considerably.

• Transients as a second signature. The authors note that retaining rather than discarding initial transients makes the data sample relaxation onto the attractor, which slows down critically near a bifurcation, giving the weights a second signature beyond the change in the invariant measure. This deliberate per-experiment choice is flagged but not systematically exploited.

• Transfer beyond dynamical systems. Whether the same weight-space response localises transitions in systems without a clean parameter indexing — heterogeneous public model repositories, for instance — remains open.

Target Audience

Researchers working on weight-space learning, model zoos, and hypernetworks who need a validation setting with known ground truth; physicists and applied mathematicians interested in data-driven bifurcation detection, particularly for transitions that linearisation cannot see; and practitioners building families of surrogate or reduced-order models across a design or control parameter space. Readers without a background in dynamical-systems theory will find the bifurcation discussion demanding, but the forecasting benchmark and the denoising result are accessible on their own.

Authors’ abstract

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.