Binned Spectral Power Loss for Improved Prediction of Chaotic Systems

TL;DR

The paper tackles spectral bias in data-driven forecasting of chaotic, multiscale dynamics by introducing Binned Spectral Power (BSP) Loss, a frequency-domain objective that aligns the predicted and true energy distributions across spatial scales via energy binning. BSP is designed to be architecture-agnostic and adds minimal computational overhead, improving both stability and spectral fidelity during long autoregressive rollouts. Through synthetic experiments and benchmarks on 2D and 3D turbulence, as well as related chaotic flows, BSP demonstrates superior preservation of energy across wavenumbers and better alignment with physical invariants compared to pointwise losses. The approach provides a practical, scalable path toward robust long-term predictions in chaotic systems, with noted limitations on unstructured grids and opportunities for grid-aware extensions.

Abstract

Forecasting multiscale chaotic dynamical systems with deep learning remains a formidable challenge due to the spectral bias of neural networks, which hinders the accurate representation of fine-scale structures in long-term predictions. This issue is exacerbated when models are deployed autoregressively, leading to compounding errors and instability. In this work, we introduce a novel approach to mitigate the spectral bias which we call the Binned Spectral Power (BSP) Loss. The BSP loss is a frequency-domain loss function that adaptively weighs errors in predicting both larger and smaller scales of the dataset. Unlike traditional losses that focus on pointwise misfits, our BSP loss explicitly penalizes deviations in the energy distribution across different scales, promoting stable and physically consistent predictions. We demonstrate that the BSP loss mitigates the well-known problem of spectral bias in deep learning. We further validate our approach for the data-driven high-dimensional time-series forecasting of a range of benchmark chaotic systems which are typically intractable due to spectral bias. Our results demonstrate that the BSP loss significantly improves the stability and spectral accuracy of neural forecasting models without requiring architectural modifications. By directly targeting spectral consistency, our approach paves the way for more robust deep learning models for long-term forecasting of chaotic dynamical systems.

Figures (16)

• Figure 1: (left) MSE over training iterations for BSP Loss (blue), MSE (orange), and FFT Loss (green), showing faster convergence of BSP. (right) Frequency domain plot of predictions across training: BSP (top) recovers high-frequency components of $g(k)$ better than MSE (bottom).
• Figure 3: Total variation (TV) distance between the predicted and true spectral component distributions across wavenumbers $k_x$ and $k_y$ for different loss functions. Among all methods, the model trained with the BSP loss exhibits the lowest TV distance, indicating the closest match to the true spectral distribution and the most effective mitigation of spectral bias.
• Figure 4: Velocity magnitude 3D plot for ground truth(left), UNet prediction(mid), and UNet + BSP loss prediction(right) after 5 auto-regressive rollouts. Clearly the UNet prediction has some blurring effect compared to other two.
• Figure 5: Comparison of energy spectra $E(k)$ as a function of wavenumber at different time steps ($T = 1, 15, 30$) and averaged over time. The plots show results from DNS (blue solid line), UNet (orange dashed line), and UNet model trained with BSP loss (green dashed line), along with the theoretical $k^{-5/3}$ scaling kolmogorov1941local (red solid line). The inclusion of BSP improves the spectral accuracy at high wavenumbers compared to the standalone UNet approach.
• Figure 6: Function approximation across training iterations. Top: BSP Loss; Bottom: MSE Loss. BSP better captures sharp transitions and high-frequency modes early in training.