Dynamics-Aware Weighting for Deep Learning Forecasts of Chaotic Systems

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.