Benchmark dataset KF256 for Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction

We present the dataset accompanying the ICML 2025 paper “Chaos Meets Attention: Transformers for Large-Scale Dynamical Prediction.” In this work, we proposed a novel architecture of attention for spatial correlation in modelling high-dimensional chaotic systems, with a focus on achieving faithful long-time statistical behaviour. This dataset addresses a benchmark on the 2D Kolmogorov flow, a canonical setting for studying turbulence and spatiotemporal chaos. The vorticity field data was generated using the high-fidelity pseudo-spectral solver introduced in (Dresdner et al. 2022). The dataset consists of 370 trajectories, each initialised with a randomised vorticity field and evolved for 500 time steps with a fixed stride of Δt = 0.00175 s. The spatial resolution is 256 × 256, and the Reynolds number regime is consistent with the moderate to strongly turbulent flows studied in prior work. We split the dataset into 300 training, 40 validation, and 30 test trajectories in .npy format. This dataset is intended to support research in machine learning for dynamical systems, particularly in the challenging regime of high-dimensional chaotic flows.