PaperlayerView exact revision on arXiv

Source-linked AI summary

S2S-JEPA: Predicting the Predictable at Subseasonal-to-Seasonal Timescales

Chenyu Dong, Gianmarco Mengaldo

arXiv:2610.03106v1physics.ao-phcs.AIcs.LG

TL;DR

S2S forecasting is a difficult but important window between medium-range weather and seasonal prediction, while many AI models lose skill beyond two weeks. S2S-JEPA predicts latent representations of future weeks rather than all grid-level details. It becomes increasingly competitive with ECMWF at longer leads, matching it on some deterministic measures and exceeding it on another at weeks 5–6.

  • Problem

    S2S forecasts cover a critical two-week-to-two-month decision window, but remain challenging between atmospheric initial-condition guidance and slower climate signals.

    Source: §1 · p. 1

  • Method

    S2S-JEPA forecasts future weekly Earth-system states in latent space, using a two-stage approach that progresses from one-step prediction to multi-step rollouts.

    Source: §1 · p. 3, §2.1 · p. 4, §2.2 · p. 6

  • Results

    S2S-JEPA becomes increasingly competitive with ECMWF as lead time increases; at weeks 5–6 it matches ECMWF for MSLP and TP and achieves significantly higher TCC for T2M.

    Source: §4 · p. 15

  • Takeaways & Limitations

    The results demonstrate the feasibility of using JEPA for S2S prediction, with competitive skill at longer lead times.

    Source: §4 · pp. 14–15

    S2S-JEPA is substantially less skillful than ECMWF at weeks 1–3, and its weekly aggregation may discard fine-grained information useful at shorter leads.

    Source: §4 · p. 15

Abstract

from arXiv · show

The subseasonal-to-seasonal (S2S) timescale, roughly from two weeks to two months ahead, is a critical forecast window for sectors such as agriculture, energy, and water management. Yet, it is widely known as the `predictability desert'. Recent AI weather models excel up to two weeks ahead but deteriorate beyond, largely because they are trained to predict fine-scale details that are neither predictable nor essential at S2S timescales. We argue that a more physically grounded objective is to forecast only the slowly varying components that remain predictable. Computer vision reached the same conclusion with the Joint-Embedding Predictive Architecture (JEPA), which predicts in latent space, discarding unpredictable details. In this work, we introduce S2S-JEPA, which brings the JEPA paradigm to S2S forecasting. It is tailored to this task through design elements from state-of-the-art AI weather models. S2S-JEPA achieves comparable skill to the gold-standard ECMWF physics-based ensemble and surpasses it on multiple metrics at weeks 5 to 6.

1 Introduction

S2S forecasting matters for weather-sensitive decisions but remains difficult because it falls between the useful influence of atmospheric initial conditions and slower climate signals. S2S-JEPA instead predicts future Earth-system representations, achieving skill competitive with ECMWF at weeks 5–6 while running on a single GPU.

1 Introduction

  • S2S forecasting spans roughly two weeks to two months, a critical decision window known as the “predictability desert.”It lies beyond the timescale of atmospheric initial-condition guidance but before slower climate signals fully dominate.Source: §1 · p. 1
  • Physics-based S2S systems are physically grounded but computationally expensive, particularly because ensembles are needed to isolate weak signals from atmospheric variability.This computational burden limits operational use at these lead times.Source: §1 · p. 2
  • Many AI weather models lose skill beyond medium-range leads because full-state reconstruction spends capacity on fine-scale details that become unpredictable.The paper argues for forecasting slowly varying signals such as intraseasonal oscillations and persistent land and ocean anomalies instead.Source: §1 · pp. 2–3
  • JEPA predicts hidden-target representations from context rather than reconstructing raw observations, aligning with S2S predictability concentrated in a few slow, large-scale modes.The paper applies this principle to global subseasonal forecasting, following earlier JEPA work in vision and video.Source: §1 · p. 3
  • At weeks 5–6, S2S-JEPA delivers skill competitive with the ECMWF operational S2S ensemble while training and running on a single GPU.Operational S2S forecasts require ensemble integrations on dedicated supercomputing infrastructure.Source: §1 · pp. 3–4
  • S2S-JEPA reformulates global subseasonal forecasting as latent-space prediction, training representations to capture slowly evolving, predictable Earth-system signals.It adapts the V-JEPA encoder–predictor architecture to 79 atmospheric, land-surface, and oceanic variables.Source: §1 · p. 3

2 Method

S2S-JEPA forecasts six future weeks by predicting latent representations of weekly Earth-system states. Its two-stage training first learns one-step prediction and then trains multi-step rollouts, with separate readouts and evaluation against ECMWF forecasts.

2.1 Problem setup

  • The task forecasts six weeks of weekly-mean temperature and sea-level pressure and weekly accumulated precipitation from 79 standardized-anomaly variables, using a seven-day context week.Training, validation, and test periods are 1979–2018, 2019–2021, and 102 initializations from December 2023 to November 2025; anomaly climatology uses training years only.Source: §2.1 · pp. 4–5

2.2 Joint-embedding predictive formulation and two-stage training

  • Training proceeds from one-step latent prediction to six-step rollouts, matching predictions to EMA-target-encoder embeddings without reconstructing physical fields.The rollout stage trains on the predictor’s own outputs; each predicted state is detached before the next step, while the context encoder receives gradient through the first step only.Source: §2.2 · pp. 5–7, Figure 1 · p. 5

2.3 Model architecture

  • The encoder tokenizes 79 variables into 4,000 tokens across five latent levels, while the predictor uses global cross-attention and windowed self-attention.Climatology and top-of-atmosphere insolation condition tokens through FiLM layers.Source: §2.3 · p. 7

2.4 Deterministic and probabilistic readout

  • Separate readouts produce deterministic forecasts and probabilistic ensembles from latent states without sending gradients back to the encoder or predictor.The deterministic readout is a lightweight convolutional decoder; the probabilistic readout uses diffusion to sample residuals, with 50 samples per target week.Source: §2.4 · pp. 7–8

2.5 Evaluation protocol

  • Evaluation compares deterministic TCC and probabilistic fair RPSS against ECMWF forecasts on 102 shared test initializations, using ERA5 weekly aggregates as ground truth.Quintile boundaries use each system’s own 20-year hindcast climate; temperature forecasts and observations are detrended before boundary computation.Source: §2.5 · p. 8

3 Results

Across weeks 5–6, S2S-JEPA matches the ECMWF ensemble across deterministic variables and shows strong long-range probabilistic skill for temperature and precipitation. Ablations favor autoregressive forecasting, latent prediction for rollout skill, and two-stage training.

3.1 Deterministic skill

  • By weeks 5–6, S2S-JEPA matches ECMWF across all three deterministic variables and significantly outperforms it for 2-meter temperature.ECMWF leads at week 3; S2S-JEPA matches or exceeds it for temperature from week 4 and for pressure from week 5, while ECMWF remains ahead on precipitation across weeks 3–6.Source: §3.1 · pp. 8–9, Table 1 · p. 9
  • At weeks 5–6, the systems show similar temperature predictability patterns, while S2S-JEPA has higher skill over South America, Africa, India, and Indochina.For pressure, their area averages and spatial patterns largely agree; ECMWF retains an overall advantage for scattered precipitation skill and significant skill over a broader area.Source: §3.1 · p. 9, Figure 2 · p. 10

3.2 Probabilistic skill

  • For probabilistic skill, S2S-JEPA matches or exceeds ECMWF for temperature from week 4 and significantly outperforms it for precipitation at weeks 5–6, while ECMWF retains higher pressure RPSS.ECMWF has a clear advantage across all three variables at week 3; its pressure advantage is no longer significant at the two longest leads.Source: §3.2 · pp. 9–10, Table 2 · p. 11
  • At weeks 5–6, area-averaged precipitation RPSS is 0.017 for S2S-JEPA versus 0.006 for ECMWF.For temperature, S2S-JEPA shows additional significant skill in regions including eastern Canada and the Arabian Peninsula; detrending favors S2S-JEPA, but the reported detrended precipitation improvement is cautioned because trend estimates are noisy.Source: §3.2 · pp. 10–12, Figure 3 · p. 11

3.3 Sources of forecast skill

  • Among forecasting formulations, autoregressive forecasting performs better overall: lead-conditioned models perform substantially worse, while per-lead experts lead only at week 1.Per-lead experts score 0.620 versus 0.609 at week 1, then fall behind from week 2 onward; the authors attribute differences partly to trade-offs in timescale consistency and pretraining-data diversity.Source: §3.3 · p. 12
  • After rollout training, S2S-JEPA has substantially better skill than direct supervised regression from week 2 onward, although supervised regression is better at week 1.The authors interpret the rollout gap as evidence for latent prediction emphasizing more predictable components; they note this does not make direct field reconstruction unsuitable for S2S generally.Source: §3.3 · pp. 13–14
  • Two-stage training performs best at S2S leads; curriculum learning shifts skill toward earlier weeks, while mixed-horizon rollout training is slightly worse overall.The gap between mixed-horizon training and two-stage training becomes small at weeks 5–6, and the authors adopt two-stage training in the final model.Source: §3.3 · p. 14

4 Discussion

S2S-JEPA predicts evolving Earth-system structure in latent space rather than reconstructing every detail, with skill becoming more competitive against ECMWF at longer leads. It remains weaker at weeks 1–3, and the discussion identifies variable-specific predictability and further development as open concerns.

4 Discussion

  • S2S-JEPA shifts the objective from reproducing every future-state detail to predicting predictable Earth-system structure in a learned representation.The motivation is that individual weather systems become less predictable with lead time while slower, larger-scale components can retain useful predictability.Source: §4 · p. 14
  • At weeks 5–6, S2S-JEPA matches ECMWF on deterministic MSLP and TP and has higher TCC for T2M; probabilistically, it matches or exceeds ECMWF for T2M and TP but not MSLP.Probabilistic improvements are significant for T2M at week 6 and TP at weeks 5–6.Source: §4 · p. 15
  • The framework remains substantially less skillful than ECMWF at weeks 1–3 and does not outperform it across all variables and metrics at longer leads.Weekly aggregation and emphasis on predictable latent information may discard fine-grained information; retained representations may also align unevenly with variable-specific predictability.Source: §4 · pp. 15–16
  • As a first exploration of JEPA for S2S prediction, the framework leaves room for development through larger datasets, physics-aware masking, and interpretation of learned representations.The discussion also identifies alternative JEPA formulations as potential building blocks for future models.Source: §4 · p. 16

A.1 Variable List

S2S-JEPA uses daily-mean ERA5 fields on a global 1.5° grid, including atmospheric, land-surface, and oceanic inputs intended to represent slowly varying sources of subseasonal predictability.

A.1 Variable List

  • S2S-JEPA uses 79 time-varying ERA5 fields on a global 1.5° grid, alongside static fields and insolation.The forecast targets are t2m, msl, and tp.Source: §A.1 · p. 22, Table 4 · p. 22
  • Inputs extend beyond common atmospheric and surface variables to include soil temperature and moisture, ocean surface conditions, and stratospheric coverage up to 10 hPa.These additions are intended to represent slowly varying components and boundary conditions relevant to S2S predictability.Source: §A.1 · p. 22
  • The variable table identifies which inputs enter the encoder and which enter the predictor, with insolation assigned to the predictor.The static fields enter the encoder, while the listed forecast targets are t2m, msl, and tp.Source: Table 4 · p. 22

A.2 Data Preprocessing

Preprocessing converts daily ERA5 inputs into standardized anomalies on a 120 × 240 grid, using training-period climatology and a precipitation-specific transformation.

A.2 Data Preprocessing

  • ERA5 fields from 1979–2025 are conservatively regridded from 0.25° to 1.5°; accumulated precipitation and top net thermal radiation are daily totals, and other fields are daily means.The −90° row is dropped, leaving a 120 × 240 grid for forecasts and verification.Source: §A.2 · p. 23
  • Anomalies are standardized against a fixed day-of-year climatology computed only from training years 1979–2018, then clipped at ±10.The climatology pools nearby day offsets and smooths day-of-year statistics with the first four annual harmonics.Source: §A.2 · p. 23
  • Daily precipitation is transformed as log(1 + x/ε), with ε set to the 10th percentile of positive daily totals in the training years.The stated purpose is to account for precipitation's heavy-tailed distribution; ε is 1.3×10−5 m.Source: §A.2 · p. 23

B Model Architecture Detail

The encoder organizes weather inputs into physically motivated groups and latent levels, while the predictor uses global and local attention to advance the latent state one week at a time.

Encoder.

  • The encoder partitions 79 variables into 16 physically motivated groups organized across five latent levels, from boundary and surface variables to pressure-level regimes.Seven context days are stacked channel-wise and embedded with 6 × 6 patch convolutions; group embeddings within each level are concatenated and projected to one token per patch.Source: Encoder · p. 23, Table 5 · p. 24
  • Tokens of width 512 are processed by six Swin Transformer blocks with three-dimensional windows spanning all five levels and 4 × 4 patches, resolving vertical coupling in every layer.Source: Encoder · p. 24

Predictor.

  • The predictor uses global cross-attention to the encoded state and windowed self-attention among learned output queries to predict the next weekly latent state.It advances one week per step, conditions on input-week climatology, and its complete output can feed the next autoregressive step; each step is trained against the target week's teacher representation.Source: Predictor · pp. 24–25

B.2 Deterministic readout head

A shared lightweight convolutional head decodes the latent state into physical fields at each rollout step. It predicts all 79 variables, providing a diagnostic of information retained in the latent representation.

B.2 Deterministic readout head

  • The head predicts standardized weekly-mean anomalies for all 79 variables, not just the three forecast targets, to diagnose information retained in the latent representation.Precipitation is predicted as a weekly total in logarithmic space because averaging daily log-transformed values does not preserve the weekly total.Source: §B.2 · p. 25
  • At each patch, tokens from five latent levels pass through 1 × 1, 3 × 3, and final 1 × 1 convolutions, then pixel shuffle restores the 120 × 240 grid.The convolutional head is shared across rollout steps and uses GELU activations.Source: §B.2 · p. 25
  • The head is trained jointly as a probe, with gradients stopped at the latent state so its ground-truth loss does not shape the learned representation.Its latitude-weighted mean-squared-error loss is computed against ground truth.Source: §B.2 · p. 25
  • At inference, the shared head produces deterministic forecasts at 7–42-day leads by decoding after each of six rollout steps.Source: §B.2 · p. 25

B.3 Diffusion-based generative model

A conditional diffusion model generates ensemble members by learning residuals to the deterministic forecast while the forecast backbone and readout head remain frozen. It uses forecast fields and their source latent state as conditioning, then calibrates ensemble spread before evaluation.

B.3 Diffusion-based generative model

  • The conditional diffusion model learns the residual between verifying truth and deterministic forecast, with the backbone and readout head frozen.A separate model is trained for each of the three target variables.Source: §B.3 · p. 25
  • Conditioning combines decoded target forecasts, their source latent state, forecast lead, and target-week day of year.At the coarsest U-Net stages, field tokens cross-attend to all 800 latent tokens on the 20 × 40 patch grid.Source: §B.3 · p. 25
  • Models are trained on 1979–2018 and selected by lowest fair CRPS on 2019–2021; 20-step Heun sampling produces ensembles calibrated by variable and lead before the test period.Test ensembles contain 50 members; hindcasts for the model’s own climatology contain 11.Source: §B.3 · p. 26

C Training details

The paper uses the Stage-2 checkpoint selected for best weeks 3–6 validation skill, alongside staged training of the forecast model and separate diffusion heads. System differences are assessed with paired moving-block bootstrap intervals, which the authors caution are optimistic given only two test years.

C Training details

  • Stage 1 trains encoder, predictor, and readout head with a one-step joint-embedding objective; Stage 2 continues from iteration 10,000 and trains six rollout steps against teacher encodings.Stage 2 detaches latent states between steps, freezes the teacher, and uses a backbone learning rate 25 times smaller than Stage 1.Source: C Training details · p. 26
  • The Stage-2 snapshot at iteration 2,000 had the best validation skill at weeks 3–6 and is used throughout the paper.Source: C Training details · p. 27
  • The forecast model uses the same architecture, data, and optimizer in both stages, with 91.0M trainable parameters including the readout head.Source: Table 6 · p. 27
  • Each diffusion head is trained on rollouts of frozen S2S-JEPA with residuals normalized per lead and variable; the lowest-validation-fair-CRPS snapshot is retained.The three heads share settings except for validation-chosen noise-distribution mean and spread; each has 9.3M parameters.Source: C Training details · p. 27, Table 7 · p. 27

D Significance Test

  • Paired moving-block bootstrap resamples eight consecutive initializations over 102 dates, with 1,000 resamples and significance when the 95% interval excludes zero.Because there are only two test years, the authors say the intervals should be read as optimistic; blocks address subseasonal autocorrelation but not interannual variability.Source: D Significance Test · pp. 27–28

E Trend removal in the probabilistic evaluation

The probabilistic evaluation also examines removing linear trends from MSLP and precipitation before quintile categorization. Under that variant, scores shift in S2S-JEPA’s favor, which the authors attribute to noise in extrapolated precipitation trends rather than forecast skill.

E Trend removal in the probabilistic evaluation

Related papers

Loading 2610.03106v1…