Less Tokens, Better Forecasts: Sparse Residual Routing for Efficient Weather Prediction

# Less Tokens, Better Forecasts: Sparse Residual Routing for Efficient Weather Prediction
Source: [https://arxiv.org/html/2607.02829](https://arxiv.org/html/2607.02829)
11institutetext:Tulane University22institutetext:Northeastern University33institutetext:University of Alabama at Birmingham44institutetext:Oak Ridge National Laboratory###### Abstract

Existing ViT\-based weather forecasting models apply uniform computation across all spatial tokens, even though nearby atmospheric grid points often contain similar values and large regions evolve smoothly over time\. This makes much of the intermediate per\-token computation redundant\. Standard token\-efficiency methods, such as pruning or merging, reduce cost by removing or fusing tokens\. However, weather forecasting is a spatiotemporal dense prediction problem in which a history of atmospheric states must be mapped to future values on the original latitude\-longitude grid\. Thus, every grid cell must retain a physically meaningful representation, especially under autoregressive rollout\. We introduceSparse\-Reslim, a parameter\-free plug\-in routing module that makes sparse token processing compatible with this fixed\-grid requirement\. Sparse\-Reslim routes only 25% of spatial tokens through the expensive middle transformer blocks and treats those blocks as residual updates: it computes the change produced for the routed tokens and scatters only this delta back to the full sequence\. Unselected tokens keep their pre\-routing representations exactly, so no grid cell is dropped or replaced by a mask token, and no fusion layer or additional parameters are introduced\. Across ERA5 resolutions up to the operational 0\.25° standard and two model families, a deterministic Transformer and a diffusion model, Sparse\-Reslim improves forecast accuracy on every evaluated variable while substantially reducing cost: training is about2\.5×2\.5\\timesfaster in the main settings and reaches3\.18×3\.18\\timesspeedup at 0\.25°, with over2\.2×2\.2\\timeslower peak memory\. A controlled decomposition shows that the accuracy gain comes primarily from sparse routing itself, while random token selection provides an additional regularization benefit without selector overhead\. Code is available at[https://github\.com/janet\-sw/Sparse\-Reslim](https://github.com/janet-sw/Sparse-Reslim)\.

## 1Introduction

Data\-driven weather forecasting has reached a turning point\[weyn2019can,scher2018toward,dueben2018challenges\]\. Models trained on reanalysis data now match or exceed the accuracy of operational numerical weather prediction \(NWP\) systems that took decades to develop, and they do so orders of magnitude faster at inference time\[bi2023accurate,lam2023learning,bodnar2025foundation\]\. These models operate on high\-dimensional gridded atmospheric fields \(ERA5 at 0\.25° comprises over one million grid points per variable\) and must produce spatially dense outputs where every grid cell carries a physically meaningful forecast\. As the community pushes toward higher resolutions and larger model capacities, training cost becomes the primary bottleneck, because a finer grid produces more spatial tokens and self\-attention scales quadratically with token count, while larger models add further compute and activation memory per token\.

A striking pattern across the leading weather models \(Pangu\-Weather\[bi2023accurate\], ClimaX\[nguyen2023climax\], Stormer\[nguyen2024scaling\], FuXi\[chen2023fuxi\], Aurora\[bodnar2025foundation\], and their generative counterparts GenCast\[price2023gencast\]and CorrDiff\[mardani2025residual\]\) is that they all rely on the Vision Transformer \(ViT\) or its variants\. Self\-attention treats every spatial token identically, spending the sameO​\(L2​D\)O\(L^\{2\}D\)computation on a token over calm open ocean as on one at the core of an extratropical cyclone\. Yet atmospheric fields are highly autocorrelated: large regions of the globe carry smooth, slowly varying conditions whose intermediate representations are largely redundant\. This suggests that substantial computation can be saved by processing only a representative subset of tokens in the most expensive layers\.

Token pruning\[rao2021dynamicvit,liang2022not,kong2022spvit,xu2022evo\], token merging\[bolya2022token\], and window\-based attention\[liu2021swin\]have been widely explored for ViT efficiency in the vision community\. More recently, Sprint\[park2025sprint\]proposed a dense\-sparse\-dense block schedule for diffusion transformers that routes only a random subset of tokens through the middle blocks, achieving large training speedups for image generation\. Directly applying these methods to weather forecasting is nontrivial because weather prediction is a spatiotemporal dense prediction problem: the model conditions on past atmospheric states and must produce a physically meaningful future value at every grid point, often for autoregressive rollout\. Token pruning and merging can remove or alter spatial positions, while Sprint replaces skipped tokens with learned mask embeddings and later fuses them through a linear projection\. This is acceptable for static image generation, but in weather forecasting it creates a mismatch: late blocks operate on a mixture of true atmospheric representations and synthetic placeholder tokens at fixed grid positions, without explicitly preserving the temporal conditioning structure of the forecast problem\.

In this work, we proposeSparse\-Reslim, a parameter\-free plug\-in routing module for ViT\-based weather forecasters that makes sparse token processing compatible with dense prediction\. Sparse\-Reslim partitions the transformer blocks into dense early, sparse middle, and dense late stages\. At the entrance to the sparse middle stage, it randomly selects a fractionrrof theLLspatial tokens, yieldingK=⌊r​L⌋K=\\lfloor rL\\rfloorrouted tokens, and processes only this subset through the expensive middle blocks\. Crucially, instead of replacing the full sequence with processed tokens and placeholders, Sparse\-Reslim computes the*residual delta*produced by the sparse middle blocks and scatters only this delta back to the original token positions\. Unselected tokens receive zero delta and therefore retain their pre\-routing representations exactly, ensuring that allLLspatial token positions remain valid throughout the model\. The module requires no mask embeddings, fusion layers, or additional learnable parameters, and can be enabled in an existing ViT block loop with a single configuration flag\.

We validate Sparse\-Reslim on ERA5\[hersbach2020era5\]at multiple resolutions and on two architecture families: \(1\) a deterministic Res\-Slim\-ViT adapted from the ORBIT\-2 exascale system\[wang2025orbit\], and \(2\) a diffusion\-based generative model following the EDM framework\[karras2022elucidating\]\. For the generative variant, we use an asymmetric attention design in which self\-attention is sparsified while cross\-attention retains the full conditioning sequence\. Across all settings, a keep ratio ofr=0\.25r\{=\}0\.25yields∼2\.5×\{\\sim\}2\.5\\timestraining speedup\. Notably, Sparse\-Reslim also improves forecast accuracy over the dense baseline across all evaluated variables and metrics\. Our ablations show that this gain is not explained by stochastic regularization alone: deterministic sparse routing already captures most of the improvement, while random token selection provides a smaller additional benefit without selector parameters or scoring overhead\. Our main contributions are:

- •We propose Sparse\-Reslim, a sparse routing module for spatiotemporal weather forecasting that preserves all spatial token positions while respecting temporal conditioning\. It adds no learnable components and integrates into ViT\-based weather forecasters as a plug\-in module\.
- •We provide an empirical study spanning two architecture families \(deterministic and generative\), three ERA5 grid resolutions, and four atmospheric variables\. Sparse\-Reslim consistently improves accuracy across evaluated settings, and its benefits become more pronounced at higher resolution\.
- •We show that sparse routing is orthogonal to other efficiency strategies, including mixed\-precision training and resolution downsampling, and can be combined with them for further gains\. To our knowledge, this is the first application of sparse token routing to efficient training of ViT\-based weather models in both deterministic and generative settings\.

![Refer to caption](https://arxiv.org/html/2607.02829v1/x1.png)Figure 1:Overview of Sparse\-Reslim\. \(a\) Three\-stage token routing: dense early \(n1n\_\{1\}blocks\), sparse middle \(n2n\_\{2\}blocks onK=⌊r​L⌋K\{=\}\\lfloor rL\\rfloortokens\), and dense late \(n3n\_\{3\}blocks\), with unselected tokens following an identity path\. \(b, c\) Integration into the deterministic Res\-Slim\-ViT and generative Res\-Slim\-EDM backbones\.
## 2Related Work

### 2\.1AI Weather Forecasting Models

Data\-driven weather forecasting has emerged as a compelling alternative to traditional NWP, reaching comparable or better accuracy at a fraction of the computational cost\. The task is to forecast future atmospheric states𝐗T∈ℝV×H×W\\mathbf\{X\}\_\{T\}\\in\\mathbb\{R\}^\{V\\times H\\times W\}from initial conditions𝐗0\\mathbf\{X\}\_\{0\}, whereVVdenotes physical variables,H×WH\\times Wthe spatial grid, andTTthe target lead time\. Among deterministic approaches, Pangu\-Weather\[bi2023accurate\]introduced a 3D Earth\-Specific Transformer operating on pressure\-level volumes, while GraphCast\[lam2023learning\]adopted a mesh\-based graph neural network for message passing on the sphere\. ClimaX\[nguyen2023climax\]proposed a variable\-agnostic ViT trained across heterogeneous climate datasets, and Stormer\[nguyen2024scaling\]combined shifted window attention with weather\-specific inductive biases\. FengWu\[chen2023fengwu\]uses a multi\-modal, multi\-task cross\-modal fusion transformer to push skillful forecasts beyond 10 lead days, and FuXi\[chen2023fuxi\]and Aurora\[bodnar2025foundation\]extended the frontier with cascade and foundation model designs, respectively\.

On the probabilistic side, GenCast\[price2023gencast\]applied diffusion models to produce calibrated ensemble forecasts, while OmniCast\[nguyen2025omnicast\]and CorrDiff\[mardani2025residual\]explored generative approaches for modeling the full distribution of future weather states, addressing the blurring and error accumulation that affect deterministic autoregressive rollout\. Despite their architectural diversity, all these models share one trait: they apply uniform computation across all spatial tokens, regardless of local informational complexity\. This becomes increasingly wasteful at higher resolutions, where quadratic attention cost dominates training time yet much of the grid is smooth and redundant\. Res\-Slim\-ViT, introduced as part of the ORBIT\-2 exascale climate downscaling system\[wang2025orbit\], showed strong scalability through a dual\-path design combining a ViT backbone with a convolutional residual path, making it a natural testbed for efficient token routing\.

### 2\.2Efficient Vision Transformers

The quadratic cost of self\-attention has motivated extensive work on reducing the computational burden of ViTs, and decreasing the number of tokens processed is among the most studied strategies\. Attention\-based token pruning removes tokens with low attention scores\[liang2022not,kong2022spvit,xu2022evo,fayyaz2022adaptive\]\. Several methods go further by training scoring modules that learn task\-dependent pruning criteria\[rao2021dynamicvit,rao2023dynamic\]\. Token merging offers a complementary direction: ToMe\[bolya2022token\]condenses similar tokens via bipartite matching, reducing sequence length while preserving representational content\. Other approaches operate at the attention level itself, including window\-based attention\[liu2021swin\]and learned sparse attention patterns\[wei2023sparsifiner,chen2021chasing\]\.

A more recent line of work targets generative model training specifically\. Sprint\[park2025sprint\]introduced a dense\-sparse\-dense block schedule for diffusion transformers: only a fraction of randomly selected tokens pass through the middle blocks, while dropped tokens are replaced by learnable mask embeddings and later fused with shallow dense features via a linear projection\. While effective for image generation, Sprint is not directly suited to weather forecasting\. Weather models condition on past atmospheric states and must return a complete future state on the original grid\. In autoregressive use, this output becomes the next input\. Replacing unprocessed positions with mask embeddings therefore does more than add a small projection layer: it breaks the continuity of the atmospheric state representation at fixed grid locations\. Our residual delta formulation addresses this gap by keeping unselected tokens as their actual pre\-routing representations and scattering back only the updates computed for routed tokens\. In the generative setting, we further preserve temporal conditioning by sparsifying self\-attention over target tokens while retaining full cross\-attention to the conditioning sequence\.

## 3Methods

This section describes Sparse\-Reslim, a drop\-in sparse token routing module for ViT\-based weather models\. Figure[1](https://arxiv.org/html/2607.02829#S1.F1)provides a visual overview of Sparse\-Reslim\. We first cover the shared tokenization and base architectures \(Sec\.[3\.1](https://arxiv.org/html/2607.02829#S3.SS1)\), then the routing module \(Sec\.[3\.2](https://arxiv.org/html/2607.02829#S3.SS2)\), and its integration into deterministic \(Sec\.[3\.3](https://arxiv.org/html/2607.02829#S3.SS3)\) and generative \(Sec\.[3\.4](https://arxiv.org/html/2607.02829#S3.SS4)\) forecasting pipelines\.

### 3\.1Preliminaries

Tokenization\.Both architectures share a per\-variable tokenization scheme\. Each variablevvis independently patchified by a dedicated convolutional embedding𝙿𝚊𝚝𝚌𝚑𝙴𝚖𝚋𝚎𝚍v\\mathtt\{PatchEmbed\}\_\{v\}, producing tokens of dimensionDD\. A learnable variable embedding𝐞v\\mathbf\{e\}\_\{v\}is added, and variables are aggregated via cross\-attention with a shared query𝐪∈ℝD\\mathbf\{q\}\\in\\mathbb\{R\}^\{D\}:

𝐳i=CrossAttn​\(𝐪,\{𝐩i,v\+𝐞v\}v∈V\)∈ℝD,\\mathbf\{z\}\_\{i\}=\\text\{CrossAttn\}\\\!\\bigl\(\\mathbf\{q\},\\;\\\{\\mathbf\{p\}\_\{i,v\}\+\\mathbf\{e\}\_\{v\}\\\}\_\{v\\in V\}\\bigr\)\\in\\mathbb\{R\}^\{D\},\(1\)where𝐩i,v\\mathbf\{p\}\_\{i,v\}is the patch token at spatial positioniifor variablevv\. This yields a sequence𝐙∈ℝL×D\\mathbf\{Z\}\\in\\mathbb\{R\}^\{L\\times D\}withL=\(h/p\)​\(w/p\)L=\(h/p\)\(w/p\)tokens, augmented with sinusoidal positional and spatial resolution embeddings\.

Deterministic forecasting\.The deterministic modelfθf\_\{\\theta\}directly regresses the forecast target\. Given input history𝐗\\mathbf\{X\}, the history dimension is flattened into the channel dimension and the model predicts:

Y^=fθ​\(𝐗\)=CNNskip​\(𝐗\)\+Unpatchify​\(MLPhead​\(ViT​\(𝐙\)\)\),\\hat\{Y\}=f\_\{\\theta\}\(\\mathbf\{X\}\)=\\text\{CNN\}\_\{\\text\{skip\}\}\(\\mathbf\{X\}\)\+\\text\{Unpatchify\}\\bigl\(\\text\{MLP\}\_\{\\text\{head\}\}\(\\text\{ViT\}\(\\mathbf\{Z\}\)\)\\bigr\),\(2\)whereViT​\(⋅\)\\text\{ViT\}\(\\cdot\)denotesNNtransformer blocks andCNNskip\\text\{CNN\}\_\{\\text\{skip\}\}is a convolutional residual path \(Conv→\\toPixelShuffle→\\toConv\) that bypasses the transformer\. The dual\-path design lets the lightweight CNN capture dense local patterns, freeing the transformer to focus on long\-range dependencies\. Training minimizes a latitude\-weighted mean squared error:

ℒdet=1\|Vout\|​∑vωv⋅∑icos⁡\(ϕi\)​\(Y^v,i−Yv,i\)2∑icos⁡\(ϕi\),\\mathcal\{L\}\_\{\\text\{det\}\}=\\frac\{1\}\{\|V\_\{\\text\{out\}\}\|\}\\sum\_\{v\}\\omega\_\{v\}\\cdot\\frac\{\\sum\_\{i\}\\cos\(\\phi\_\{i\}\)\\bigl\(\\hat\{Y\}\_\{v,i\}\-Y\_\{v,i\}\\bigr\)^\{2\}\}\{\\sum\_\{i\}\\cos\(\\phi\_\{i\}\)\},\(3\)whereωv\\omega\_\{v\}are per\-variable weights andϕi\\phi\_\{i\}is the latitude at grid pointii\.

Generative forecasting\.The diffusion model follows the EDM framework\[karras2022elucidating\], parameterizing a denoiserDθD\_\{\\theta\}via preconditioning scalars derived from the noise levelσ\\sigma:

Dθ​\(𝐲;σ,𝐜\)=cskip​\(σ\)⋅𝐲\+cout​\(σ\)⋅Fθ​\(cin​\(σ\)⋅𝐲;σ,𝐜\),D\_\{\\theta\}\(\\mathbf\{y\};\\sigma,\\mathbf\{c\}\)=c\_\{\\text\{skip\}\}\(\\sigma\)\\cdot\\mathbf\{y\}\+c\_\{\\text\{out\}\}\(\\sigma\)\\cdot F\_\{\\theta\}\\\!\\bigl\(c\_\{\\text\{in\}\}\(\\sigma\)\\cdot\\mathbf\{y\};\\;\\sigma,\\;\\mathbf\{c\}\\bigr\),\(4\)where𝐲=Y\+σ​ϵ\\mathbf\{y\}=Y\+\\sigma\\boldsymbol\{\\epsilon\}is the noisy target withϵ∼𝒩​\(𝟎,𝐈\)\\boldsymbol\{\\epsilon\}\\sim\\mathcal\{N\}\(\\mathbf\{0\},\\mathbf\{I\}\),𝐜\\mathbf\{c\}represents the conditioning fields \(history frames\), andcskip,cout,cinc\_\{\\text\{skip\}\},c\_\{\\text\{out\}\},c\_\{\\text\{in\}\}are analytic functions ofσ\\sigmaand the data standard deviationσdata\\sigma\_\{\\text\{data\}\}\. The raw networkFθF\_\{\\theta\}is a ViT whose blocks contain noise\-conditioned FiLM modulation, noise\-gated cross\-attention, and a CNN skip path:

Fθ​\(⋅\)=Unpatchify​\(MLPhead​\(ViTedm​\(𝐙x,𝐙c,𝐞σ\)\)\)\+CNNskip​\(cin⋅𝐲\)\.F\_\{\\theta\}\(\\cdot\)=\\text\{Unpatchify\}\\bigl\(\\text\{MLP\}\_\{\\text\{head\}\}\(\\text\{ViT\}\_\{\\text\{edm\}\}\(\\mathbf\{Z\}\_\{x\},\\mathbf\{Z\}\_\{c\},\\mathbf\{e\}\_\{\\sigma\}\)\)\\bigr\)\+\\text\{CNN\}\_\{\\text\{skip\}\}\(c\_\{\\text\{in\}\}\\cdot\\mathbf\{y\}\)\.\(5\)Here𝐙x\\mathbf\{Z\}\_\{x\}and𝐙c\\mathbf\{Z\}\_\{c\}are token sequences from the noisy target and conditions respectively, and𝐞σ=MLP​\(ln⁡σ/4\)\\mathbf\{e\}\_\{\\sigma\}=\\text\{MLP\}\(\\ln\\sigma/4\)is the noise embedding\. Each EDM block consists of FiLM modulation \(*scale*,*shift*from𝐞σ\\mathbf\{e\}\_\{\\sigma\}\), self\-attention,σ\\sigma\-gated cross\-attention \(gate=sigmoid​\(𝐖𝐞σ\)\\text\{gate\}=\\text\{sigmoid\}\(\\mathbf\{W\}\\mathbf\{e\}\_\{\\sigma\}\)\), and MLP\. Training minimizes the denoising loss:

ℒedm=𝔼σ∼p​\(σ\)​\[λ​\(σ\)​‖Dθ​\(𝐲;σ,𝐜\)−Y‖2\],\\mathcal\{L\}\_\{\\text\{edm\}\}=\\mathbb\{E\}\_\{\\sigma\\sim p\(\\sigma\)\}\\Bigl\[\\lambda\(\\sigma\)\\bigl\\\|D\_\{\\theta\}\(\\mathbf\{y\};\\sigma,\\mathbf\{c\}\)\-Y\\bigr\\\|^\{2\}\\Bigr\],\(6\)whereln⁡σ∼𝒩​\(Pmean,Pstd2\)\\ln\\sigma\\sim\\mathcal\{N\}\(P\_\{\\text\{mean\}\},P\_\{\\text\{std\}\}^\{2\}\)andλ​\(σ\)=\(σ2\+σdata2\)/\(σ⋅σdata\)2\\lambda\(\\sigma\)=\(\\sigma^\{2\}\+\\sigma\_\{\\text\{data\}\}^\{2\}\)/\(\\sigma\\cdot\\sigma\_\{\\text\{data\}\}\)^\{2\}is the weighting\.

### 3\.2Sparse\-Reslim: Sparse Token Routing

Motivation\.Self\-attention costsO​\(L2​D\)O\(L^\{2\}D\)per block\. For ERA5 at 1\.0° with patch size 2,L=16,200L\{=\}16\{,\}200tokens, and theN=12N\{=\}12transformer blocks dominate training time\. Atmospheric fields, however, exhibit strong spatial autocorrelation: nearby grid points carry highly redundant information, especially in the smooth interior of large\-scale circulation features\. Every token must appear in the final output \(forecasts are spatially dense\), but the*intermediate representations*in the middle layers can be approximated by processing only a subset of tokens and broadcasting their updates back via a residual delta\.

Three\-stage block partition\.GivenNNtransformer blocks, Sparse\-Reslim partitions them as:

Block1,…,Blockn1⏟Dense Early,Blockn1\+1,…,Blockn1\+n2⏟Sparse Middle,Blockn1\+n2\+1,…,BlockN⏟Dense Late,\\underbrace\{\\text\{Block\}\_\{1\},\\ldots,\\text\{Block\}\_\{n\_\{1\}\}\}\_\{\\text\{Dense Early\}\},\\quad\\underbrace\{\\text\{Block\}\_\{n\_\{1\}\+1\},\\ldots,\\text\{Block\}\_\{n\_\{1\}\+n\_\{2\}\}\}\_\{\\text\{Sparse Middle\}\},\\quad\\underbrace\{\\text\{Block\}\_\{n\_\{1\}\+n\_\{2\}\+1\},\\ldots,\\text\{Block\}\_\{N\}\}\_\{\\text\{Dense Late\}\},\(7\)wheren1\+n2\+n3=Nn\_\{1\}\+n\_\{2\}\+n\_\{3\}=N\. The dense early stage \(n1n\_\{1\}blocks\) builds rich initial representations through full token interaction\. The sparse middle stage \(n2n\_\{2\}blocks\) processes only a subset of tokens, yielding the bulk of the speedup\. The dense late stage \(n3n\_\{3\}blocks\) refines all tokens before the prediction head, restoring spatial coherence\.

Random token selection\.At the boundary between the dense early and sparse middle stages, we randomly selectK=⌊L⋅r⌋K=\\lfloor L\\cdot r\\rfloortokens per sample, wherer∈\(0,1\]r\\in\(0,1\]is the*keep ratio*\. Letℐ=\{i1,…,iK\}⊂\{1,…,L\}\\mathcal\{I\}=\\\{i\_\{1\},\\ldots,i\_\{K\}\\\}\\subset\\\{1,\\ldots,L\\\}be the index set, drawn uniformly without replacement and sorted to maintain spatial ordering\. Random selection, rather than attention\-score\-based selection, has three advantages: \(1\) it requires no scoring overhead, \(2\) it acts as a stochastic regularizer during training \(analogous to dropout on the token dimension\), and \(3\) it avoids systematic bias that could cause certain geographic regions to be consistently underrepresented\.

Residual delta formulation\.Let𝐱∈ℝB×L×D\\mathbf\{x\}\\in\\mathbb\{R\}^\{B\\times L\\times D\}denote the full token sequence after the dense early stage\. We save a copy𝐱0=𝐱\\mathbf\{x\}\_\{0\}=\\mathbf\{x\}, then gather the selected tokens:𝐡0=Gather​\(𝐱,ℐ\)∈ℝB×K×D\.\\mathbf\{h\}\_\{0\}=\\text\{Gather\}\(\\mathbf\{x\},\\;\\mathcal\{I\}\)\\in\\mathbb\{R\}^\{B\\times K\\times D\}\.TheseKKtokens pass through then2n\_\{2\}sparse middle blocks:

𝐡n2=Blockn1\+n2∘⋯∘Blockn1\+1​\(𝐡0\)\.\\mathbf\{h\}\_\{n\_\{2\}\}=\\text\{Block\}\_\{n\_\{1\}\+n\_\{2\}\}\\\!\\circ\\cdots\\circ\\text\{Block\}\_\{n\_\{1\}\+1\}\(\\mathbf\{h\}\_\{0\}\)\.\(8\)The residual delta is then computed and scattered back:

𝚫=𝐡n2−𝐡0,𝐱′=𝐱0\+Scatter​\(𝚫,ℐ\),\\boldsymbol\{\\Delta\}=\\mathbf\{h\}\_\{n\_\{2\}\}\-\\mathbf\{h\}\_\{0\},\\qquad\\mathbf\{x\}^\{\\prime\}=\\mathbf\{x\}\_\{0\}\+\\text\{Scatter\}\(\\boldsymbol\{\\Delta\},\\;\\mathcal\{I\}\),\(9\)whereScatterplaces theKK\-dimensional deltas into anLL\-dimensional zero tensor at the selected indices\. For theL−KL\{\-\}Kunselected tokens,𝚫i=0\\boldsymbol\{\\Delta\}\_\{i\}=0, so they retain their pre\-routing representation𝐱0,i\\mathbf\{x\}\_\{0,i\}exactly, forming an implicit identity skip connection\. This is a key departure from Sprint\[park2025sprint\], which replaces dropped tokens with learnable mask embeddings and requires an extra linear projection to fuse them back\. The mask embeddings never pass through the middle transformer blocks, so they carry a fundamentally different representation distribution than the processed tokens; the fusion layer must learn to compensate for this mismatch\. Our formulation sidesteps the issue entirely: unselected tokens retain their actual dense\-early representations, the delta scatter is additive and parameter\-free, and no fusion layer is needed\. We verify this advantage empirically in Table[5](https://arxiv.org/html/2607.02829#S4.T5), where even within our framework, deterministic selection strategies underperform random selection, confirming that the stochastic regularization from our zero\-parameter design is a feature, not a compromise\.

Computational savings\.In the sparse middle stage, self\-attention operates onKKtokens instead ofLL, reducing cost fromO​\(n2​L2​D\)O\(n\_\{2\}L^\{2\}D\)toO​\(n2​r2​L2​D\)O\(n\_\{2\}r^\{2\}L^\{2\}D\)\. With the defaultr=0\.25r\{=\}0\.25andn1=n3=2n\_\{1\}\{=\}n\_\{3\}\{=\}2, the theoretical FLOPs reduction for the transformer blocks is:

Speedup≈Nn1\+n2​r2\+n3=122\+8⋅0\.0625\+2=124\.5≈2\.67×\.\\text\{Speedup\}\\approx\\frac\{N\}\{n\_\{1\}\+n\_\{2\}r^\{2\}\+n\_\{3\}\}=\\frac\{12\}\{2\+8\\cdot 0\.0625\+2\}=\\frac\{12\}\{4\.5\}\\approx 2\.67\\times\.\(10\)In practice, we observe∼2\.46×\{\\sim\}2\.46\\timesspeedup on AMD MI250X GPUs, with the gap from theory explained by the overhead ofgather/scatteroperations and non\-transformer components \(patch embedding, prediction head, CNN skip path\)\.

### 3\.3Integration with Deterministic Forecasting

For Res\-Slim\-ViT, integration is straightforward: each block contains only self\-attention and MLP with pre\-norm residual connections\. The sparse middle blocks process onlyKKtokens through the standard block:

𝐡←𝐡\+Self\-Attn​\(LN​\(𝐡\)\),𝐡←𝐡\+MLP​\(LN​\(𝐡\)\),\\mathbf\{h\}\\leftarrow\\mathbf\{h\}\+\\text\{Self\-Attn\}\(\\text\{LN\}\(\\mathbf\{h\}\)\),\\quad\\mathbf\{h\}\\leftarrow\\mathbf\{h\}\+\\text\{MLP\}\(\\text\{LN\}\(\\mathbf\{h\}\)\),\(11\)with attention costO​\(K2​D\)O\(K^\{2\}D\)per block\. All other components \(patch embedding, variable aggregation, positional embedding, MLP head, CNN skip path, and unpatchify\) remain unchanged and operate on the fullLLtokens\. The complete procedure is given in Algorithm[1](https://arxiv.org/html/2607.02829#alg1)\.

Algorithm 1:Sparse\-Reslim for Deterministic Forecasting \(Res\-Slim\-ViT\)1:Input

𝐗∈ℝB×\(H⋅\|Vin\|\)×h×w\\mathbf\{X\}\\in\\mathbb\{R\}^\{B\\times\(H\\cdot\|V\_\{\\text\{in\}\}\|\)\\times h\\times w\}, blocks

\{Blocki\}i=1N\\\{\\text\{Block\}\_\{i\}\\\}\_\{i=1\}^\{N\}, keep ratio

rr, block split

\(n1,n2,n3\)\(n\_\{1\},n\_\{2\},n\_\{3\}\)
2:Forecast

Y^∈ℝB×\|Vout\|×h×w\\hat\{Y\}\\in\\mathbb\{R\}^\{B\\times\|V\_\{\\text\{out\}\}\|\\times h\\times w\}
3:

𝐬←CNNskip​\(𝐗\)\\mathbf\{s\}\\leftarrow\\text\{CNN\}\_\{\\text\{skip\}\}\(\\mathbf\{X\}\)⊳\\trianglerightParallel CNN residual path

4:

𝐙←VarAgg​\(PatchEmbed​\(𝐗\)\)\+𝐩pos\+𝐩res\\mathbf\{Z\}\\leftarrow\\text\{VarAgg\}\(\\text\{PatchEmbed\}\(\\mathbf\{X\}\)\)\+\\mathbf\{p\}\_\{\\text\{pos\}\}\+\\mathbf\{p\}\_\{\\text\{res\}\}⊳\\triangleright𝐙∈ℝB×L×D\\mathbf\{Z\}\\in\\mathbb\{R\}^\{B\\times L\\times D\}

5:% Stage 1: Dense early \(allLLtokens\)

6:for

i=1i=1to

n1n\_\{1\}do

7:

𝐙←Blocki​\(𝐙\)\\mathbf\{Z\}\\leftarrow\\text\{Block\}\_\{i\}\(\\mathbf\{Z\}\)
8:endfor

9:% Stage 2: Sparse middle \(KKrandom tokens\)

10:

K←⌊L⋅r⌋K\\leftarrow\\lfloor L\\cdot r\\rfloor;

ℐ←sort\(randperm\(L\)\[:K\]\)\\mathcal\{I\}\\leftarrow\\texttt\{sort\}\(\\texttt\{randperm\}\(L\)\[\\,:K\\,\]\)
11:

𝐱0←𝐙\.clone\(\)\\mathbf\{x\}\_\{0\}\\leftarrow\\mathbf\{Z\}\.\\texttt\{clone\(\)\};

𝐡0←gather​\(𝐙,ℐ\)\\mathbf\{h\}\_\{0\}\\leftarrow\\texttt\{gather\}\(\\mathbf\{Z\},\\mathcal\{I\}\);

𝐡←𝐡0\.clone\(\)\\mathbf\{h\}\\leftarrow\\mathbf\{h\}\_\{0\}\.\\texttt\{clone\(\)\}
12:for

i=n1\+1i=n\_\{1\}\{\+\}1to

n1\+n2n\_\{1\}\{\+\}n\_\{2\}do

13:

𝐡←Blocki​\(𝐡\)\\mathbf\{h\}\\leftarrow\\text\{Block\}\_\{i\}\(\\mathbf\{h\}\)⊳\\trianglerightSelf\-Attn cost:O​\(K2​D\)O\(K^\{2\}D\)

14:endfor

15:

𝚫←𝐡−𝐡0\\boldsymbol\{\\Delta\}\\leftarrow\\mathbf\{h\}\-\\mathbf\{h\}\_\{0\};

𝐙←𝐱0\+scatter​\(𝟎,ℐ,𝚫\)\\mathbf\{Z\}\\leftarrow\\mathbf\{x\}\_\{0\}\+\\texttt\{scatter\}\(\\mathbf\{0\},\\mathcal\{I\},\\boldsymbol\{\\Delta\}\)
16:% Stage 3: Dense late \(allLLtokens\)

17:for

i=n1\+n2\+1i=n\_\{1\}\{\+\}n\_\{2\}\{\+\}1to

NNdo

18:

𝐙←Blocki​\(𝐙\)\\mathbf\{Z\}\\leftarrow\\text\{Block\}\_\{i\}\(\\mathbf\{Z\}\)
19:endfor

20:

Y^←Unpatchify​\(MLPhead​\(LN​\(𝐙\)\)\)\+𝐬\\hat\{Y\}\\leftarrow\\text\{Unpatchify\}\(\\text\{MLP\}\_\{\\text\{head\}\}\(\\text\{LN\}\(\\mathbf\{Z\}\)\)\)\+\\mathbf\{s\}

### 3\.4Integration with Generative Forecasting

The EDM architecture presents a more involved integration because each block contains cross\-attention with the conditioning sequence𝐙c∈ℝB×L×D\\mathbf\{Z\}\_\{c\}\\in\\mathbb\{R\}^\{B\\times L\\times D\}and noise\-conditioned FiLM modulation\. During the sparse middle stage, we adopt an*asymmetric attention design*:

Self\-attentionoperates on only theKKselected tokens \(Q,K,V∈ℝK×D\\text\{Q\},\\text\{K\},\\text\{V\}\\in\\mathbb\{R\}^\{K\\times D\}\), reducing cost fromO​\(L2​D\)O\(L^\{2\}D\)toO​\(K2​D\)O\(K^\{2\}D\)\.

Cross\-attentionuses the selected tokens as queries but retains the*full*conditioning sequence as keys and values \(Q∈ℝK×D\\text\{Q\}\\in\\mathbb\{R\}^\{K\\times D\},K,V∈ℝL×D\\text\{K\},\\text\{V\}\\in\\mathbb\{R\}^\{L\\times D\}\), reducing cost fromO​\(L2​D\)O\(L^\{2\}D\)toO​\(K​L⋅D\)O\(KL\\cdot D\)\.

This asymmetry is essential\. The conditioning fields𝐙c\\mathbf\{Z\}\_\{c\}encode the full atmospheric context \(all history frames, all input variables\), and restricting keys/values to onlyKKpositions would lose meteorological information needed for denoising\. The noise embedding𝐞σ\\mathbf\{e\}\_\{\\sigma\}is a global vector broadcast to all tokens and needs no modification; FiLM modulation and theσ\\sigma\-gate operate element\-wise and naturally accommodate the reduced sequence length\.

The per\-block cost in the sparse middle stage becomes:

O​\(K2​D\)⏟sparse self\-attn\+O​\(K​L​D\)⏟sparse cross\-attn\+O​\(K​D2\)⏟MLP \+ FiLM,\\underbrace\{O\(K^\{2\}D\)\}\_\{\\text\{sparse self\-attn\}\}\+\\underbrace\{O\(KLD\)\}\_\{\\text\{sparse cross\-attn\}\}\+\\underbrace\{O\(KD^\{2\}\)\}\_\{\\text\{MLP \+ FiLM\}\},\(12\)compared toO​\(L2​D\)\+O​\(L2​D\)\+O​\(L​D2\)O\(L^\{2\}D\)\+O\(L^\{2\}D\)\+O\(LD^\{2\}\)in the dense baseline\. The complete procedure is given in Algorithm[2](https://arxiv.org/html/2607.02829#alg2)\.

Algorithm 2:Sparse\-Reslim for Generative Forecasting \(EDM\)1:Noisy target

𝐲∈ℝB×\|Vout\|×h×w\\mathbf\{y\}\\in\\mathbb\{R\}^\{B\\times\|V\_\{\\text\{out\}\}\|\\times h\\times w\}, conditions

𝐜∈ℝB×H×\|Vin\|×h×w\\mathbf\{c\}\\in\\mathbb\{R\}^\{B\\times H\\times\|V\_\{\\text\{in\}\}\|\\times h\\times w\}, noise level

σ\\sigma, blocks

\{Blockiedm\}i=1N\\\{\\text\{Block\}\_\{i\}^\{\\text\{edm\}\}\\\}\_\{i=1\}^\{N\}, keep ratio

rr, split

\(n1,n2,n3\)\(n\_\{1\},n\_\{2\},n\_\{3\}\)
2:Denoised output

Dθ​\(𝐲;σ,𝐜\)D\_\{\\theta\}\(\\mathbf\{y\};\\sigma,\\mathbf\{c\}\)
3:

cin,cout,cskip←EDM\_precond​\(σ\)c\_\{\\text\{in\}\},c\_\{\\text\{out\}\},c\_\{\\text\{skip\}\}\\leftarrow\\text\{EDM\\\_precond\}\(\\sigma\)
4:

𝐬←CNNskip​\(cin⋅𝐲\)\\mathbf\{s\}\\leftarrow\\text\{CNN\}\_\{\\text\{skip\}\}\(c\_\{\\text\{in\}\}\\cdot\\mathbf\{y\}\)⊳\\trianglerightParallel CNN residual path

5:

𝐙x←VarAgg​\(PatchEmbed​\(cin⋅𝐲\)\)\+𝐩pos\+𝐩res\+𝐩time\\mathbf\{Z\}\_\{x\}\\leftarrow\\text\{VarAgg\}\(\\text\{PatchEmbed\}\(c\_\{\\text\{in\}\}\\cdot\\mathbf\{y\}\)\)\+\\mathbf\{p\}\_\{\\text\{pos\}\}\+\\mathbf\{p\}\_\{\\text\{res\}\}\+\\mathbf\{p\}\_\{\\text\{time\}\}
6:

𝐙c←TemporalProj​\(VarAgg​\(PatchEmbed​\(𝐜\)\)\)\+𝐩pos\+𝐩res\\mathbf\{Z\}\_\{c\}\\leftarrow\\text\{TemporalProj\}\(\\text\{VarAgg\}\(\\text\{PatchEmbed\}\(\\mathbf\{c\}\)\)\)\+\\mathbf\{p\}\_\{\\text\{pos\}\}\+\\mathbf\{p\}\_\{\\text\{res\}\}⊳\\triangleright𝐙c∈ℝB×L×D\\mathbf\{Z\}\_\{c\}\\in\\mathbb\{R\}^\{B\\times L\\times D\}

7:

𝐞σ←MLP​\(SiLU​\(MLP​\(ln⁡σ/4\)\)\)\\mathbf\{e\}\_\{\\sigma\}\\leftarrow\\text\{MLP\}\(\\text\{SiLU\}\(\\text\{MLP\}\(\\ln\\sigma/4\)\)\)⊳\\trianglerightNoise embedding

8:% Stage 1: Dense early \(allLLtokens\)

9:for

i=1i=1to

n1n\_\{1\}do

10:

𝐙x←Blockiedm​\(𝐙x,𝐙c,𝐞σ\)\\mathbf\{Z\}\_\{x\}\\leftarrow\\text\{Block\}\_\{i\}^\{\\text\{edm\}\}\(\\mathbf\{Z\}\_\{x\},\\;\\mathbf\{Z\}\_\{c\},\\;\\mathbf\{e\}\_\{\\sigma\}\)
11:endfor

12:% Stage 2: Sparse middle \(KKrandom tokens\)

13:

K←⌊L⋅r⌋K\\leftarrow\\lfloor L\\cdot r\\rfloor;

ℐ←sort\(randperm\(L\)\[:K\]\)\\mathcal\{I\}\\leftarrow\\texttt\{sort\}\(\\texttt\{randperm\}\(L\)\[\\,:K\\,\]\)
14:

𝐱0←𝐙x\.clone\(\)\\mathbf\{x\}\_\{0\}\\leftarrow\\mathbf\{Z\}\_\{x\}\.\\texttt\{clone\(\)\};

𝐡0←gather​\(𝐙x,ℐ\)\\mathbf\{h\}\_\{0\}\\leftarrow\\texttt\{gather\}\(\\mathbf\{Z\}\_\{x\},\\mathcal\{I\}\);

𝐡←𝐡0\.clone\(\)\\mathbf\{h\}\\leftarrow\\mathbf\{h\}\_\{0\}\.\\texttt\{clone\(\)\}
15:for

i=n1\+1i=n\_\{1\}\{\+\}1to

n1\+n2n\_\{1\}\{\+\}n\_\{2\}do

16:

𝐡←Blockiedm​\(𝐡,𝐙c,𝐞σ\)\\mathbf\{h\}\\leftarrow\\text\{Block\}\_\{i\}^\{\\text\{edm\}\}\(\\mathbf\{h\},\\;\\mathbf\{Z\}\_\{c\},\\;\\mathbf\{e\}\_\{\\sigma\}\)⊳\\trianglerightSelf\-Attn:O​\(K2\)O\(K^\{2\}\); Cross\-Attn Q=sparse, KV=full

17:endfor

18:

𝚫←𝐡−𝐡0\\boldsymbol\{\\Delta\}\\leftarrow\\mathbf\{h\}\-\\mathbf\{h\}\_\{0\};

𝐙x←𝐱0\+scatter​\(𝟎,ℐ,𝚫\)\\mathbf\{Z\}\_\{x\}\\leftarrow\\mathbf\{x\}\_\{0\}\+\\texttt\{scatter\}\(\\mathbf\{0\},\\mathcal\{I\},\\boldsymbol\{\\Delta\}\)
19:% Stage 3: Dense late \(allLLtokens\)

20:for

i=n1\+n2\+1i=n\_\{1\}\{\+\}n\_\{2\}\{\+\}1to

NNdo

21:

𝐙x←Blockiedm​\(𝐙x,𝐙c,𝐞σ\)\\mathbf\{Z\}\_\{x\}\\leftarrow\\text\{Block\}\_\{i\}^\{\\text\{edm\}\}\(\\mathbf\{Z\}\_\{x\},\\;\\mathbf\{Z\}\_\{c\},\\;\\mathbf\{e\}\_\{\\sigma\}\)
22:endfor

23:

Fθ←Unpatchify​\(MLPhead​\(LN​\(𝐙x\)\)\)\+𝐬F\_\{\\theta\}\\leftarrow\\text\{Unpatchify\}\(\\text\{MLP\}\_\{\\text\{head\}\}\(\\text\{LN\}\(\\mathbf\{Z\}\_\{x\}\)\)\)\+\\mathbf\{s\}
24:

Dθ←cskip⋅𝐲\+cout⋅FθD\_\{\\theta\}\\leftarrow c\_\{\\text\{skip\}\}\\cdot\\mathbf\{y\}\+c\_\{\\text\{out\}\}\\cdot F\_\{\\theta\}

Training and inference\.A fresh random permutation is drawn at every forward pass, providing stochastic regularization analogous to dropout on the token dimension\. The same random selection is used at inference; averaging multiple runs yields implicit ensembling\. For the EDM model, routing is applied identically across all denoising steps\.

Architecture\-agnostic design\.Sparse\-Reslim requires no additional learnable parameters\. It introduces only three hyperparameters: the keep ratiorrand the block split\(n1,n3\)\(n\_\{1\},n\_\{3\}\)\(withn2=N−n1−n3n\_\{2\}=N\-n\_\{1\}\-n\_\{3\}\)\. The module wraps the existing block loop and is activated by a single configuration flag\. When disabled \(r=1r\{=\}1\), the model reverts exactly to the dense baseline with zero overhead\.

## 4Experiments

![Refer to caption](https://arxiv.org/html/2607.02829v1/x2.png)Figure 2:Training convergence for Res\-Slim\-ViT at 1\.40625° and 1\.0°\. Sparse\-Reslim \(solid\) converges faster than the dense baseline \(dashed\) in both wall\-clock time and training loss, reaching lower final values at both resolutions\.### 4\.1Experimental Setup

Dataset\.We train and evaluate on ERA5 reanalysis\[hersbach2020era5,hersbach2023era5,rasp2020weatherbench,rasp2024weatherbench\], the most widely used benchmark for data\-driven weather forecasting\. Following the protocol of ClimaX\[nguyen2023climax\], we use 1979–2017 for training, 2018–2019 for validation, and 2020 for testing\. The input consists of 40 atmospheric variables: six variables \(geopotential, temperature, specific humidity, u\-wind, v\-wind, relative humidity\) at seven pressure levels \(50, 250, 500, 600, 700, 850, 925 hPa\), plus four surface and constant fields \(2m temperature, 10m u\-wind, 10m v\-wind, land\-sea mask\)\. Evaluation targets four variables standard in NWP verification: geopotential at 500 hPa \(Z500\), temperature at 850 hPa \(T850\), 2\-meter temperature \(T2m\), and 10\-meter zonal wind \(U10\)\. Z500 and T850 are upper\-air variables used for synoptic\-scale assessment; T2m and U10 are surface variables directly relevant to human activities\. Unless noted otherwise, experiments use a 5\-day \(120\-hour\) forecast lead time\. We run experiments at three spatial resolutions: 1\.40625° \(128×256128\\times 256grid,L=8,192L\{=\}8\{,\}192tokens with patch size 2\), 1\.0° \(181×360181\\times 360grid,L=16,200L\{=\}16\{,\}200tokens with patch size 2\), and 0\.25° \(720×1440720\\times 1440grid,L=16,200L\{=\}16\{,\}200tokens with patch size 8\)\. The 0\.25° setting uses the same 1979–2017 training period as the coarser\-resolution experiments\.

Model configurations\.We evaluate Sparse\-Reslim on two backbone architectures:*Deterministic model \(Res\-Slim\-ViT\)\.*Adapted from ORBIT\-2\[wang2025orbit\], this model uses a dual\-path architecture with a ViT backbone and a convolutional residual skip path\. At 1\.40625°, the model has embed\_dim=256\{=\}256, depth=12\{=\}12, four attention heads, totaling 12\.18M parameters\. At 1\.0°, the same architecture on the larger grid yields 14\.23M parameters\. At 0\.25°, we use patch size 8 so that the routed token count matches the 1\.0° setting\. All deterministic settings use block split\(n1,n2,n3\)=\(2,8,2\)\(n\_\{1\},n\_\{2\},n\_\{3\}\)\{=\}\(2,8,2\)with keep ratior=0\.25r\{=\}0\.25\.*Generative model \(EDM\)\.*Following the EDM framework\[karras2022elucidating\], this model uses a ViT backbone with FiLM\-conditioned noise modulation,σ\\sigma\-gated cross\-attention, and a CNN skip path\. The architecture has patch\_size=6\{=\}6, embed\_dim=1024\{=\}1024, depth=8\{=\}8, eight attention heads, totaling∼\{\\sim\}180M parameters, with block split\(n1,n2,n3\)=\(2,4,2\)\(n\_\{1\},n\_\{2\},n\_\{3\}\)\{=\}\(2,4,2\)andr=0\.25r\{=\}0\.25\.

Implementation details\.All models are trained with Fully Sharded Data Parallel \(FSDP\) and activation checkpointing\. The deterministic model runs on 2 nodes with 16 AMD MI250X GPUs, using AdamW with learning rate5×10−45\{\\times\}10^\{\-4\}\(scaled bynum\_gpus0\.8\\text\{num\\\_gpus\}^\{0\.8\}\), cosine schedule with 5 epochs of linear warmup, and batch size 32 \(1\.40625°\) or 16 \(1\.0° and 0\.25°\)\. The generative model is trained on 8 nodes with 64 AMD MI250X GPUs\. Both use bfloat16 mixed precision\. The latitude\-weighted MSE \(Eq\.[3](https://arxiv.org/html/2607.02829#S3.E3)\) and EDM denoising loss \(Eq\.[6](https://arxiv.org/html/2607.02829#S3.E6)\) are used for the deterministic and generative models, respectively\.

Baselines\.Our primary comparison is the dense baseline, i\.e\., the identical architecture with sparse routing disabled \(r=1r\{=\}1\)\. This isolates the effect of sparse routing from other design choices\. We refer to the dense baseline asReslimand the sparse variant asSparse\-Reslim, trained for the same number of epochs\.

Evaluation metrics\.Forecast quality is measured with three complementary metrics: \(1\)RMSE\(root mean squared error\), the standard pointwise accuracy metric; \(2\)ACC\(anomaly correlation coefficient\), which measures correlation between predicted and observed anomalies relative to climatology and is the primary skill score in operational NWP; and \(3\)SSIM\(structural similarity index\), which captures perceptual similarity and penalizes spatial blurring\. For efficiency, we report wall\-clock training time, throughput \(samples per second per GPU\), and peak VRAM\.

### 4\.2Results on Deterministic Forecasting

![Refer to caption](https://arxiv.org/html/2607.02829v1/x3.png)Figure 3:Training efficiency across resolutions for the deterministic model\. Sparse\-Reslim achieves consistent wall\-clock speedup at both resolutions, with larger gains at 1\.0° due to the higher token count \(L=16,200L\{=\}16\{,\}200vs\.8,1928\{,\}192\)\.Table 1:Forecast quality on the ERA5 test set, reported as single\-step, non\-autoregressive 120\-hour forecasts for the deterministic Res\-Slim\-ViT\. Sparse\-Reslim usesr=0\.25r=0\.25with block split\(2,8,2\)\(2,8,2\)\. The0\.25∘0\.25^\{\\circ\}setting uses the full 1979–2017 training period with patch size 8\.Bold= better within each resolution\.Table[1](https://arxiv.org/html/2607.02829#S4.T1)compares forecast quality between the dense Reslim baseline and Sparse\-Reslim across all four target variables and all three resolutions\. Sparse\-Reslim outperforms the dense baseline on every metric and every variable, often by a substantial margin\. At 1\.0°, Z500 RMSE drops from 868\.7 to 768\.3 and ACC rises from 0\.551 to 0\.651\. T2m shows a similar pattern: RMSE improves from 3\.19 to 2\.68 and ACC from 0\.797 to 0\.859\. The gains are consistent at 1\.40625° as well, though slightly smaller in magnitude, which is expected given the lower token count\. The 0\.25° experiment further confirms the high\-resolution behavior under the full 39\-year training split: Sparse\-Reslim improves all four variables while using the same routed token count as the 1\.0° patch\-size\-2 setting\.

Efficiency methods typically aim to*preserve*accuracy, not improve it\. Our ablations indicate a dual mechanism: token\-level sparsity itself reduces redundant intermediate computation and can improve generalization, while random selection adds a further dropout\-like effect without selector parameters or scoring overhead\. Figure[2](https://arxiv.org/html/2607.02829#S4.F2)confirms that Sparse\-Reslim converges faster and reaches lower final loss than the dense baseline, ruling out under\-training as an explanation\. Figure[3](https://arxiv.org/html/2607.02829#S4.F3)shows that speedup is consistent across resolutions and grows at higher resolutions where the larger token count amplifies quadratic savings\.

Table 2:Peak VRAM per AMD MI250X GCD for the deterministic model\. Sparse\-Reslim usesr=0\.25r\{=\}0\.25and split\(2,8,2\)\(2,8,2\)\.Table[2](https://arxiv.org/html/2607.02829#S4.T2)reports empirical peak VRAM at fixed token counts on 16 AMD MI250X GCDs\. Sparse\-Reslim reduces peak activation memory by more than2\.2×2\.2\\timesacross all tested settings\. The measured reduction is close to the theoretical middle\-block estimate12/\(4\+8r2\)=2\.67×12/\(4\+8r^\{2\}\)=2\.67\\timesfor split\(2,8,2\)\(2,8,2\)atr=0\.25r\{=\}0\.25; the remaining gap comes from full\-sequence components such as patch embedding, dense early/late blocks, the prediction head, and the CNN skip path\.

Table 3:Autoregressive rollout stability on the 2020 test year\. Both models are trained for 24\-hour prediction at 1\.40625° and rolled out for 1–10 days\. Values are T2m ACC\.Table[3](https://arxiv.org/html/2607.02829#S4.T3)addresses the stability of stochastic routing under autoregressive deployment\. Sparse\-Reslim remains above the dense baseline throughout the 10\-day rollout, and the skill decays smoothly rather than showing compounding instability\. The result is consistent with the residual identity path: an unselected token is not replaced by noise or a mask embedding, but simply receives zero middle\-block delta before being refined by the dense late blocks\.

### 4\.3Results on Generative Forecasting

We apply Sparse\-Reslim to the EDM\-based generative model at 1\.0° with the asymmetric attention design described in Sec\.[3\.4](https://arxiv.org/html/2607.02829#S3.SS4)\. The EDM backbone has∼180\{\\sim\}180M parameters, depth 8 with split\(2,4,2\)\(2,4,2\), and uses patch size 6 with embed dim 1024\. All metrics are computed on the ensemble mean of 4 samples\.

Figure 4:Forecast quality of the EDM generative model on the ERA5 test set \(year 2020, 120\-hour lead time\) at 1\.0°\. Sparse\-Reslim usesr=0\.25r\{=\}0\.25with block split\(2,4,2\)\(2,4,2\)\.Bold= better\.Table[4](https://arxiv.org/html/2607.02829#S4.F4)shows that Sparse\-Reslim transfers to the generative setting: Z500 RMSE drops from 951\.3 to 889\.7 and ACC rises from 0\.517 to 0\.571\. T2m and T850 follow the same pattern with RMSE reductions of 6\.6% and 6\.4%, respectively\. The improvement margin is smaller than in the deterministic case, which is expected: the EDM backbone uses split\(2,4,2\)\(2,4,2\)with only 4 sparse middle blocks instead of 8, reducing the computational savings and the strength of stochastic regularization\. Sparse\-Reslim EDM trains at1\.82×1\.82\\timesthe wall\-clock speed of the dense EDM baseline \(20\.8 vs\. 11\.4 samples/s/GPU\), consistent with the smallern2n\_\{2\}and the fact that cross\-attention retains full\-length conditioning tokens\.

### 4\.4Ablation Studies

Ablations use the deterministic Res\-Slim\-ViT at 1\.0° unless stated otherwise\.

Effect of keep ratiorr\.We varyr∈\{0\.10,0\.25,0\.50,0\.75,1\.0\}r\\in\\\{0\.10,0\.25,0\.50,0\.75,1\.0\\\}with a fixed split\(2,8,2\)\(2,8,2\)\. Figure[5](https://arxiv.org/html/2607.02829#S4.F5)plots Z500 RMSE and wall\-clock speedup as a function ofrr, wherer=1\.0r\{=\}1\.0corresponds to the dense baseline\. The accuracy\-speedup tradeoff exhibits a clear U\-shape:r=0\.25r\{=\}0\.25achieves the lowest RMSE \(768\.3\) and the highest ACC \(0\.651\), outperforming both the dense baseline \(r=1\.0r\{=\}1\.0, RMSE 868\.7\) and the more aggressiver=0\.10r\{=\}0\.10\(RMSE 801\.5\)\. Atr=0\.10r\{=\}0\.10, the model loses too many tokens and intermediate representations become impoverished; atr≥0\.50r\{\\geq\}0\.50, the stochastic regularization effect diminishes and accuracy degrades toward the dense baseline\. The speedup ranges from1\.27×1\.27\\timesatr=0\.75r\{=\}0\.75to2\.91×2\.91\\timesatr=0\.10r\{=\}0\.10, withr=0\.25r\{=\}0\.25yielding2\.46×2\.46\\times\.

![Refer to caption](https://arxiv.org/html/2607.02829v1/x4.png)Figure 5:Keep ratiorrvs\. Z500 RMSE and speedup\.r=0\.25r\{=\}0\.25is the sweet spot\.Effect of block split\(n1,n2,n3\)\(n\_\{1\},n\_\{2\},n\_\{3\}\)\.Table[4](https://arxiv.org/html/2607.02829#S4.T4)compares four block partitions for a 12\-block model atr=0\.25r\{=\}0\.25\. The split\(2,8,2\)\(2,8,2\)produces the best forecast accuracy\. Reducing either end to a single dense block, as in\(1,10,1\)\(1,10,1\), degrades RMSE by 5\.7% despite yielding the highest speedup \(3\.12×3\.12\\times\): one dense early block is insufficient for building rich initial token representations, and one dense late block cannot adequately restore spatial coherence\. Conversely, allocating more blocks to the dense stages, as in\(4,4,4\)\(4,4,4\), leaves too few sparse middle blocks for meaningful regularization and limits speedup\.

Table 4:Effect of block split\(n1,n2,n3\)\(n\_\{1\},n\_\{2\},n\_\{3\}\)on Z500 at 1\.0°\.r=0\.25r\{=\}0\.25,N=12N\{=\}12total blocks\.Random vs\. learned token selection\.A natural question is whether informed token selection can outperform random sampling\. Table[5](https://arxiv.org/html/2607.02829#S4.T5)compares three strategies: \(1\) random selection \(our default\), \(2\) top\-K by self\-attention score from the last dense early block, and \(3\) a learned difficulty head \(a lightweight MLP that predicts per\-token importance\)\. Random selection achieves the best Z500 RMSE while adding zero parameters and zero computational overhead\. Attention\-score\-based selection introduces a systematic spatial bias, consistently prioritizing dynamically active regions \(jet streams and frontal boundaries\) and underrepresenting quiescent tropical areas, which hurts generalization\. The learned head partially mitigates this bias but remains deterministic, forgoing the regularization benefit that random selection provides\.

Table 5:Token selection strategy comparison on Z500 at 1\.0°\.r=0\.25r\{=\}0\.25, split\(2,8,2\)\(2,8,2\)\.Scaling behavior with model size\.We train three model sizes \(Small: embed dim 128, depth 8, 3\.1M params; Base: embed dim 256, depth 12, 12\.2M; Large: embed dim 512, depth 16, 48\.6M\) with split\(2,N−4,2\)\(2,N\{\-\}4,2\)andr=0\.25r\{=\}0\.25\. Figure[6](https://arxiv.org/html/2607.02829#S4.F6)shows that both speedup and accuracy improvement grow with model size\. The Large model achieves2\.85×2\.85\\timesspeedup \(vs\.1\.72×1\.72\\timesfor Small\) because deeper networks have a larger fraction of their compute in the sparse middle blocks\. The RMSE improvement over the dense baseline also grows:−4\.0%\-4\.0\\%for Small,−11\.6%\-11\.6\\%for Base, and−12\.6%\-12\.6\\%for Large, indicating that larger models benefit more from the stochastic regularization provided by random token dropping\.

![Refer to caption](https://arxiv.org/html/2607.02829v1/x5.png)Figure 6:Scaling behavior: both speedup and RMSE improvement grow with model size\.

## 5Conclusion

We introduced Sparse\-Reslim, a parameter\-free sparse token routing module that partitions ViT blocks into dense early, sparse middle, and dense late stages, processing only a random 25% of tokens in the dominant middle layers while a residual delta formulation maintains spatially dense outputs\. Experiments on ERA5 at three resolutions show∼2\.5×\{\\sim\}2\.5\\timeswall\-clock training speedup and more than2\.2×2\.2\\timespeak VRAM reduction alongside improved forecast accuracy: Z500 RMSE drops by 11\.6% at 1\.0° relative to the dense baseline, with similar gains for T2m, T850, and U10\. We attribute this to a combination of sparse routing and stochastic regularization: deterministic sparse routing already captures much of the gain, while random token selection adds a smaller but consistent benefit without selector parameters or scoring overhead\. The method generalizes across two architecturally distinct backbones, a deterministic Res\-Slim\-ViT and a diffusion\-based generative model, remains stable in 10\-day autoregressive rollout, and adds no learnable parameters\. A natural extension is adaptive token selection that allocates computation to dynamically active regions such as storm fronts and jet streams\.

Acknowledgements\.This manuscript has been authored by UTBattelle, LLC, under contract DE\-AC05\-00OR22725 with the US Department of Energy \(DOE\)\. This research was supported by the ORNL’s AI Initiative sponsored by the Director’s Research and Development Program at ORNL\. This work’s computing time was supported through the Oak Ridge Leadership Computing Facility\.

## References

Appendix

## Appendix 0\.AImplementation Details

### 0\.A\.1Architecture Specifications

Tables[6](https://arxiv.org/html/2607.02829#Pt0.A1.T6)and[7](https://arxiv.org/html/2607.02829#Pt0.A1.T7)consolidate the full architectural specifications for the deterministic and generative models, respectively\.

Table 6:Architecture specification for the deterministic Res\-Slim\-ViT\.Table 7:Architecture specification for the generative EDM model at1\.0∘1\.0^\{\\circ\}\.A key design distinction: the EDM model uses*continuous noise levels*sampled from a log\-normal distribution, not discrete diffusion steps\. There is no fixed “number of training steps” or DDIM\-style sampler\. At inference, an ODE integrator \(deterministic or stochastic\) produces samples; averaging 4 ensemble members yields the reported metrics\.

### 0\.A\.2Training Configuration

Table[8](https://arxiv.org/html/2607.02829#Pt0.A1.T8)consolidates all training hyperparameters\.

Table 8:Full training configuration for the main 120\-hour experimental settings\.All models are trained on the ORNL Frontier supercomputer using AMD MI250X GPUs \(64 GB HBM2e each\)\. Fully Sharded Data Parallel \(FSDP\) distributes model state and gradients across GPUs, and activation checkpointing reduces peak memory at the cost of recomputation during the backward pass\. Both models use bfloat16 mixed precision\. The dense baseline and Sparse\-Reslim are trained for the same number of epochs to ensure a fair comparison\.

### 0\.A\.3Scaling Experiment Configurations

Table[9](https://arxiv.org/html/2607.02829#Pt0.A1.T9)details the three model sizes used in the scaling experiment \(Section[4\.4](https://arxiv.org/html/2607.02829#S4.SS4), Figure[6](https://arxiv.org/html/2607.02829#S4.F6)\)\. The block split follows the\(2,N−4,2\)\(2,\\;N\{\-\}4,\\;2\)pattern, so larger models have a greater fraction of their blocks in the sparse middle stage\.

Table 9:Model configurations for the scaling experiment \(deterministic Res\-Slim\-ViT,1\.0∘1\.0^\{\\circ\},r=0\.25r\{=\}0\.25\)\.
### 0\.A\.4Gather / Scatter Implementation

The token selection and residual delta reconstruction are implemented using standard PyTorch indexing\. Indices are sorted after random sampling to preserve spatial ordering, as specified in Algorithm[1](https://arxiv.org/html/2607.02829#alg1)\(line 6\)\.

```
# Gather: select r*L tokens (sorted for spatial ordering)
K = int(r * L)
I = torch.randperm(L, device=x.device)[:K].sort().values
x0 = x.clone()                   # save pre-routing state
h0 = x[:, I, :]                  # (B, K, D)  gather
h  = h0.clone()

# ... sparse middle blocks operate on h (B, K, D) ...

# Scatter: residual delta formulation
delta = h - h0                    # (B, K, D)  change only
x_out = x0                       # start from saved state
x_out[:, I, :] += delta          # additive scatter
```

A new random permutation is drawn at every forward pass during training, providing stochastic regularization\. At inference, the same random selection mechanism is used; averaging multiple forward passes with different seeds yields implicit ensembling \(Section[3\.2](https://arxiv.org/html/2607.02829#S3.SS2)\)\. The gather/scatter overhead is negligible because both operations are simple indexing without additional computation\.

## Appendix 0\.BAdditional Experiments

### 0\.B\.1Additional Variables and Backbones

Table[10](https://arxiv.org/html/2607.02829#Pt0.A2.T10)evaluates total precipitation, a more localized and intermittent variable than the four main targets\. Sparse\-Reslim remains better than the dense baseline, indicating that the sparse routing gain is not limited to smooth upper\-air fields\.

Table 10:Total precipitation at 1\.40625° on the 2020 test year\. RMSE is daily\-aggregated mm/day\.Table[11](https://arxiv.org/html/2607.02829#Pt0.A2.T11)removes the CNN\-skip confound by integrating Sparse\-Reslim into a ClimaX\-style pure ViT backbone with no convolutional residual path\. The sparse variant improves all evaluated variables, supporting the claim that Sparse\-Reslim functions as a drop\-in component beyond the Res\-Slim\-ViT architecture\.

Table 11:Sparse\-Reslim as a drop\-in component for a ClimaX\-style pure ViT at 1\.0°\. Values are RMSE / ACC on the 2020 test year\.
### 0\.B\.2Mechanism Ablations

Table[12](https://arxiv.org/html/2607.02829#Pt0.A2.T12)separates generic stochastic regularization from token\-level sparse routing\. Stochastic depth improves the dense baseline, and deterministic Top\-K sparse routing improves further, showing that sparsity itself accounts for a large part of the gain\. Random Sparse\-Reslim remains best while adding no selector parameters or scoring overhead\.

Table 12:Decomposing efficiency and regularization at 1\.40625°\. Values are RMSE / ACC on the 2020 test year\.We also evaluated two parameter\-free content\-aware token selectors at 1\.40625°: a gradient\-based sampler using a Laplacian score and an activation\-norm sampler from the first dense block\. As shown in Table[13](https://arxiv.org/html/2607.02829#Pt0.A2.T13), neither matches random selection, which avoids persistent geographic bias while requiring no scoring pass\.

Table 13:Additional token\-selection strategies on Z500 at 1\.40625°\.

## Appendix 0\.CLimitations and Future Work

Data and evaluation scope\.This work focuses on ERA5 reanalysis\[hersbach2020era5\], a controlled gridded approximation of atmospheric state rather than raw observational data\. Performance on operational forecast systems with irregular observation networks \(e\.g\., radiosondes, satellite swaths\) remains to be validated\. Our main forecast\-quality table focuses on single\-step 120\-hour prediction, and our autoregressive analysis is limited to T2m over 1–10 rollout days\. A broader rollout study across all variables and both deterministic and generative models would further clarify long\-horizon deployment behavior\.

Resolution and variable gap\.Our experiments include a 0\.25° high\-resolution setting, but the output space remains limited to a small set of target variables\. State\-of\-the\-art systems such as GraphCast\[lam2023learning\]and Pangu\-Weather\[bi2023accurate\]operate at0\.25∘0\.25^\{\\circ\}with 200\+ variables\. Extending Sparse\-Reslim to this broader variable set, especially humidity and precipitation across more pressure levels, remains an important direction\.

Adaptive token selection\.Random token selection is optimal in our experiments \(Table[5](https://arxiv.org/html/2607.02829#S4.T5)\), but this may not hold at higher resolutions where the information density is more spatially heterogeneous\. Adaptive strategies that allocate computation to dynamically active regions \(storm fronts, jet streams\) — noted as a natural extension in Section[5](https://arxiv.org/html/2607.02829#S5)— deserve further investigation\.

Subseasonal\-to\-seasonal forecasting\.Our evaluation focuses on medium\-range forecasting \(1–10 day lead times\)\. Subseasonal\-to\-seasonal \(S2S\) forecasting at 2–6 week horizons may require different sparsity schedules, since the relevant spatial scales and predictability sources \(e\.g\., tropical modes, stratospheric coupling\) differ substantially\.

Broader efficiency and adaptation context\.Sparse\-Reslim targets train ing\-time efficiency for a fixed forecasting task, but a complementary challenge is efficiently adapting a pretrained model when the target domain, resolution, or observation network shifts after deployment\. Parameter\-efficient and prompt\-based adaptation\[xiao2025prompt,zhang2025otvp,zhang2025dpcore,xiao2026not\], including in black\-box or model\-as\-a\-service settings\[zhang2025prime,zhang2026adapting\], offers lightweight alternatives to full fine\-tuning for this purpose\. Feedback\-guided generative modeling\[wang2025doctor\]similarly suggests a route toward improving the physical plausibility of generative weather forecasts\. Combining Sparse\-Reslim with such adaptation\-time techniques is a promising direction for future work\.