Dynamics-Aware Weighting for Deep Learning Forecasts of Chaotic Systems

Zhou Fang; Gianmarco Mengaldo

cs.LGAug 23, 2026

Dynamics-Aware Weighting for Deep Learning Forecasts of Chaotic Systems

Authors: Zhou Fang, Gianmarco Mengaldo

Organizations: Department of Mechanical Engineering National University of Singapore Singapore

Abstract

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.

Figures & tables

Figure 1: Visualization of DAW weight distribution. The left axis displays the empirical probability density $P(d)$ of the local dimension $d$ , highlighting the severe imbalance between highly recurrent, low-complexity states and the sparse, right-skewed tail of extreme regime transitions. The right axis shows the corresponding loss weights allocated by different strategies. The Standard baseline applies a naive uniform weight across all regimes, whereas DAW anchors the loss landscape to the underlying dynamical complexity. As the intensity parameter $\alpha$ increases from DAW ( $\alpha=0.0$ ) to DAW ( $\alpha=1.0$ ), the density-weighting component is progressively engaged on top of the $\tilde{d}$ -tilt, concentrating loss mass onto the high- $d$ tail and increasingly amplifying rare regime transitions relative to the recurrent low- $d$ bulk.

Figure 2: One-step forecast error binned by local dimension $d$ . The test set is partitioned into five equal-mass bins of increasing $d$ (each containing $20\%$ of states); bars show the mean one-step MAE of each method within each bin. For every training objective the error increases monotonically from the low- $d$ bulk to the high- $d$ tail, validating the premise of Section 2.3 that one-step error concentrates at high- $d$ states. DAW attains the lowest error in every bin and reduces error across the full range of $d$ – including the low- $d$ bulk – relative to Standard and RandomWeight.

Step	Method	MAE $\downarrow$	Correlation $\uparrow$
	DAW	0.0032 $\pm$ 0.0000	1.0000 $\pm$ 0.0000
	DenseWeight	0.0032 $\pm$ 0.0001	1.0000 $\pm$ 0.0000
	RandomWeight	0.0047 $\pm$ 0.0001	1.0000 $\pm$ 0.0000
1 (0.01 LT)	Standard	0.0058 $\pm$ 0.0006	1.0000 $\pm$ 0.0000
	DAW	0.1109 $\pm$ 0.0084	0.9876 $\pm$ 0.0020
	DenseWeight	0.1320 $\pm$ 0.0195	0.9825 $\pm$ 0.0050

Table 1: Model evaluation results (mean

\pm

std over three independent runs). The best-performing method for each metric at every forecast horizon is highlighted in bold .

Figure 3: Autoregressive forecasting performance on the test set. DAW (red) is compared against Standard (black), DenseWeight (orange), and RandomWeight (blue). Left: Mean Absolute Error (MAE) at each autoregressive step, averaged over the full test set. Right: Spatial Pearson Correlation Coefficient at each step, averaged over the test set. In both panels, the primary horizontal axis (bottom) gives the discrete step index; the secondary axis (top) indicates the corresponding physical timescale in Lyapunov Time (LT). Solid lines are per-step means across the test set; shaded bands denote $\pm 1$ standard deviation across test-set trajectories at each step. In the correlation panel, a horizontal dashed grey line marks the $0.5$ predictability threshold; coloured vertical dotted lines indicate the step at which each model’s mean correlation first crosses this threshold.

Figure 4: Event-level error analysis and spatiotemporal visualization of autoregressive forecasting errors on the KS test set. (a) Left three panels: pairwise log-log scatter plots of the per-event Cumulative Error (CE), defined in Section 3.4 as the sum of absolute point-wise errors over a $\pm 3$ -step window centered on the peak- $d$ time index of each identified high- $d$ event window. Each scatter plot compares Standard CE ( $y$ -axis) against the CE of one competing method ( $x$ -axis); points above the dashed $y=x$ diagonal indicate events where Standard incurs higher CE (the competing method wins), while points below indicate events where Standard incurs lower CE. Right panel: overall win-rate pie chart reporting, across all 21,146 identified high- $d$ event windows, the fraction of windows in which each method achieves the lowest CE among all four competitors. (b) Spatiotemporal MAE contour maps for four representative test trajectories (columns). The top row shows the ground-truth KS field; subsequent rows show the spatial MAE of Standard, DAW, DenseWeight, and RandomWeight. In all panels, the primary horizontal axis (bottom) gives the discrete step index and the secondary axis (top) gives the corresponding LT. The primary vertical axis gives the spatial coordinate $x$ . Overlaid dashed curves trace the temporal evolution of the local dimension $d$ against the secondary vertical axis (right); grey-shaded intervals mark the identified high- $d$ event windows (exceeding the 75th-percentile threshold) at which CE is evaluated.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure A1: Joint probability density of $(u_{xx},u_{x})$ along the autoregressive rollout. Rows correspond to the lead times $k=1$ , $20$ , $40$ , $60$ and $80$ steps ( $0.01$ to $1.0$ LT); columns give the reference data followed by the four training objectives. Each panel is the density of the pushforward measure $\nu_{k}$ , estimated by pooling all test initial conditions at the fixed step index $k$ , with derivatives computed spectrally under the band limit $n_{c}=16$ applied identically to the reference data and to every surrogate. All panels share the same bin edges and the same logarithmic colour scale, normalised to unit peak density. The annotation in each panel reports the sliced $2$ -Wasserstein distance (WD) to the reference measure at the same lead time, defined in Eq. ( 21 ). The four objectives are indistinguishable at $k=1$ and separate monotonically thereafter, with DAW holding the smallest distance at every subsequent horizon.

Figure A2: Density ratio relative to the reference measure. Base-ten logarithm of the ratio $\rho_{\mathrm{pred}}/\rho_{\mathrm{true}}$ on the bin grid of Fig. A1 , with rows and columns as in that figure and the reference column omitted. Red indicates an excess of probability mass relative to the reference and blue a deficit; bins that are empty in either density are left blank. The colour scale is symmetric about zero and shared across all panels. Whereas Fig. A1 reports how far each measure has drifted, this figure localises the drift: under uniform training the excess mass accumulates in the interior of the reference distribution while the outer arcs are depleted, the signature of a map that has smoothed the sharp gradient structure of the attractor.

Figure B1: Convergence and Robustness of the $d$ Estimation. (a) Probability density functions of the standardized local dimension $d$ estimated from trajectories of increasing length (from 50 LT to 1200 LT); in the axis labels, $d_{L}$ denotes $d$ estimated from a trajectory of the indicated length. (b) Consistency comparison of the $d$ estimated at 850 LT versus 1200 LT. A perfect fit is shown with the dashed diagonal line.

Figure C1: Hyperparameter sensitivity analysis of the weighting intensity parameter $\alpha$ . The plots display the Mean Absolute Error (MAE) evaluated at four autoregressive forecasting horizons: Step 1, Step 20 (0.25 LT), Step 40 (0.5 LT), and Step 80 (1.0 LT). The red solid line represents the DAW model (mean $\pm$ std over three independent runs). The horizontal grey dashed line and orange dash-dotted line indicate the mean performance of the Standard baseline and the DenseWeight ( $\alpha=1.0$ ) baseline, respectively; shaded bands around each line denote $\pm 1$ standard deviation over the three independent runs. At extended forecast horizons (0.5 LT and 1.0 LT), DAW demonstrates a broad and robust performance basin, consistently outperforming both baselines across a wide range of $\alpha$ values, confirming that its efficacy is not contingent upon narrow hyperparameter tuning.

Step	Method	MAE $\downarrow$	Correlation $\uparrow$
	$d(x)$	0.0033 $\pm$ 0.0001	1.0000 $\pm$ 0.0000
	$d(x,y)$	0.0032 $\pm$ 0.0000	1.0000 $\pm$ 0.0000
1	$d(y)$	0.0033 $\pm$ 0.0001	1.0000 $\pm$ 0.0000
	$d(x)$	0.1171 $\pm$ 0.0021	0.9867 $\pm$ 0.0001
	$d(x,y)$	0.1109 $\pm$ 0.0084	0.9876 $\pm$ 0.0020
20 (0.25 LT)	$d(y)$	0.1210 $\pm$ 0.0059	0.9856 $\pm$ 0.0010

Appendix

Table D1: Factor-space ablation results (mean

\pm

std over three independent runs). The best-performing method for each metric at every forecast horizon is highlighted in bold .

Figure E1: Autoregressive forecasting performance on the test set ( $q=0.995$ ). DAW (red) is compared against Standard (black), DenseWeight (orange), and RandomWeight (blue). Left: Mean Absolute Error (MAE) at each autoregressive step, averaged over the full test set. Right: Spatial Pearson Correlation Coefficient at each step, averaged over the test set. In both panels, the primary horizontal axis (bottom) gives the discrete step index; the secondary axis (top) indicates the corresponding physical timescale in Lyapunov Time (LT). Solid lines are per-step means across the test set; shaded bands denote $\pm 1$ standard deviation across test-set trajectories at each step. In the correlation panel, a horizontal dashed grey line marks the $0.5$ predictability threshold; coloured vertical dotted lines indicate the step at which each model’s mean correlation first crosses this threshold.

Figure E2: Autoregressive forecasting performance on the test set ( $q=0.98$ ). DAW (red) is compared against Standard (black), DenseWeight (orange), and RandomWeight (blue). Left: Mean Absolute Error (MAE) at each autoregressive step, averaged over the full test set. Right: Spatial Pearson Correlation Coefficient at each step, averaged over the test set. In both panels, the primary horizontal axis (bottom) gives the discrete step index; the secondary axis (top) indicates the corresponding physical timescale in Lyapunov Time (LT). Solid lines are per-step means across the test set; shaded bands denote $\pm 1$ standard deviation across test-set trajectories at each step. In the correlation panel, a horizontal dashed grey line marks the $0.5$ predictability threshold; coloured vertical dotted lines indicate the step at which each model’s mean correlation first crosses this threshold.

Figure F1: Autoregressive forecasting performance on the KS test set with a 1D UNet backbone. DAW (red) is compared against Standard (black), DenseWeight (orange), and RandomWeight (blue). Left: Mean Absolute Error (MAE) at each autoregressive step, averaged over the full test set. Right: Spatial Pearson Correlation Coefficient at each step, averaged over the test set. In both panels, the primary horizontal axis (bottom) gives the discrete step index; the secondary axis (top) indicates the corresponding physical timescale in Lyapunov Time (LT). Solid lines are per-step means across the test set; shaded bands denote $\pm 1$ standard deviation across test-set trajectories at each step. In the correlation panel, a horizontal dashed grey line marks the $0.5$ predictability threshold; coloured vertical dotted lines indicate the step at which each model’s mean correlation first crosses this threshold.

Figure G1: Autoregressive forecasting performance on the test set for Lorenz. DAW (red) is compared against Standard (black), DenseWeight (orange), and RandomWeight (blue). Left: Mean Absolute Error (MAE) at each autoregressive step, averaged over the full test set. Right: Pearson correlation across the three state components at each step, averaged over the test set. In both panels, the primary horizontal axis (bottom) gives the discrete step index; the secondary axis (top) indicates the corresponding physical timescale in Lyapunov Time (LT). Solid lines are per-step means across the test set; shaded bands denote $\pm 1$ standard deviation across test-set trajectories at each step. In the correlation panel, a horizontal dashed grey line marks the $0.5$ predictability threshold; coloured vertical dotted lines indicate the step at which each model’s mean correlation first crosses this threshold.

Method	Win rate (%) [95% CI]	$\Delta$ vs. DAW (%) [95% CI]
DAW	$42.4\ [41.7,43.1]$	—
RandomWeight	$28.1\ [27.5,28.8]$	$+14.3\ [+13.0,+15.5]$
DenseWeight	$24.4\ [23.8,25.0]$	$+18.0\ [+16.8,+19.2]$
Standard	$\phantom{0}5.1\ [\phantom{0}4.8,\phantom{0}5.4]$	$+37.3\ [+36.5,+38.2]$

Appendix

Table H1: Cluster (block) bootstrap of the Overall Win Rate. Bootstrap of the Overall Win Rate is calculated over the

21{,}146

high-

d

event windows, resampling the

T=13{,}111

event-bearing test trajectories with replacement (

B=10{,}000

). Win rates are given with

95\%

percentile confidence intervals; the last column reports the paired difference

\mathrm{DAW}-\text{method}

with its

95\%

interval. All contrasts against DAW exclude zero.

Chaotic Time-Series Prediction

Chaotic time-series prediction studies how to forecast nonlinear systems whose sensitive dependence on initial conditions limits long-range accuracy.

See all 52 papers

May 14, 2026

Watch your neighbors: Training statistically accurate chaotic systems with local phase space information

Chaotic systems pose fundamental challenges for data-driven dynamics discovery, as small modeling errors lead to exponentially growing trajectory discrepancies. Since exact long-term prediction is unattainable, it is natural to ask what a good surrogate model for chaotic dynamics is. Prior work has largely focused either on reproducing the Jacobian of the underlying dynamics, which governs local expansion and contraction rates, or on training surrogate models that reproduce the ground-truth dynamics' long-term statistical behavior. In this work, we propose a new framework that aims to bridge these two paradigms by training surrogate dynamics models with accurate Jacobians and long-term statistical properties. Our method constructs a local covering of a chaotic attractor in phase space and analyzes the expansion and contraction of these coverings under the dynamics. The surrogate model is trained by minimizing the maximum mean discrepancy between the pushforward distributions of the coverings under the surrogate and ground-truth dynamics. Experiments show that our method significantly improves Jacobian accuracy while remaining competitive with state-of-the-art statistically accurate dynamics learning methods. Our code is fully available at https://anonymous.4open.science/r/neighborwatch.

Joon-Hyuk Ko, Andrus Giraldo, Deok-Sun Lee

Center for AI and Natural Sciences Korea Institute for Advanced Study Seoul, Korea 02455 · School of Computational Sciences Korea Institute for Advanced Study Seoul, Korea 02455

Surrogate Modeling Nonlinear System Identification Chaotic Dynamical Systems

Ask PDF

May 29, 2026

The Dynamic-Probabilistic Consistency Gap in Chaotic Surrogate Modeling

Dynamical systems reconstruction (DSR) aims to learn surrogate models that capture the dynamics underlying time-series data. Reliably deploying these surrogates requires uncertainty estimates consistent with the learned dynamics. We expose a dynamic-probabilistic consistency (DPC) gap: the pursuit of finite-horizon probabilistic objectives can degrade dynamics or decouple predictive uncertainty from the local tangent dynamics it ought to reflect. We isolate three mechanisms behind this gap: core collapse, noise masking, and blind uncertainty. Specifically, we show that open-loop Gaussian rollout objectives can penalize Jacobian-generated covariance growth in chaotic systems, encouraging optimization shortcuts that weaken physical expansion or decouple uncertainty from it. To mitigate this gap, we propose KAFFEE (Kalman-Aware Framework For Ergodic Emulation), a differentiable extended Kalman filter-based training framework that evaluates likelihood on local predictive residuals (innovations) while transporting covariance through learned local Jacobians. On stochastic hyperchaotic Lorenz-96, KAFFEE reduces the identified failure modes, improves reconstruction of dynamical invariants relative to open-loop objectives, and maintains competitive predictive scores. We further show that the DPC gap appears when probabilistically adapting a DSR foundation model across 13 chaotic systems, where KAFFEE enables in-context Bayesian filtering while largely preserving zero-shot dynamics.

Andre Herz, Matthijs Pals, Daniel Durstewitz +1

1Interdisciplinary Center for Scientific Computing, Heidelberg University, Germany · Faculty of Mathematics and Computer Science, Heidelberg University, Germany · 3Dept. of Theoretical Neuroscience, Central Institute of Mental Health (CIMH), Mannheim, Germany +3

Dynamical Systems Surrogate Modeling Uncertainty Quantification Chaotic Dynamical Systems Extended Kalman Filtering

Ask PDF

Mar 26, 2026

A Distribution-to-Distribution Neural Probabilistic Forecasting Framework for Dynamical Systems

Probabilistic forecasts provide a principled representation of forecast uncertainty through predictive distributions. Directly propagating predictive distributions with conventional physical models requires solving distributional evolution equations, a task well-known to be computationally prohibitive for high-dimensional nonlinear systems. Consequently, multi-step predictive distributions are commonly approximated using ensembles or sampled trajectories. A novel distribution-to-distribution (D2D) neural probabilistic forecasting framework is proposed with compatible distributional encoding and decoding interfaces around a replaceable neural forecasting module. Kernel mean embeddings represent input distributions, while mixture density networks parameterise outputs, enabling recursive multi-step forecasting in distributional form. The framework is demonstrated on the Lorenz63 system. Results show that the D2D model captures nontrivial distributional evolution, produces skillful probabilistic forecasts without explicit ensemble simulation, and remains competitive with a naively initialised perfect-model benchmark. This provides proof-of-concept evidence that predictive distributions can be recursively propagated through a neural forecasting model and evaluated in terms of probabilistic forecast performance.

Tianlin Yang, Hailiang Du, Louis Aslett

Department of Mathematical Sciences, Durham University, Durham DH1 3LE, United Kingdom · School of Mathematics, East China University of Science and Technology, Shanghai 200237, China · Data Science Institute, The London School of Economics and Political Science, London WC2A 2AE, United Kingdom

Mixture Density Networks Probabilistic Time Series Forecasting

Ask PDF

Abstract

Chaotic Time-Series Prediction

Explore similar work

Watch your neighbors: Training statistically accurate chaotic systems with local phase space information

The Dynamic-Probabilistic Consistency Gap in Chaotic Surrogate Modeling

A Distribution-to-Distribution Neural Probabilistic Forecasting Framework for Dynamical Systems