PaperlayerView exact revision on arXiv

Source-linked AI summary

Tracing the Unlabeled Storm: Cross-Variable Transfer in a Lagrangian Atmospheric JEPA Framework

K M Anirudh, S Sandeep, Hariprasad Kodamana

arXiv:2608.22358v1cs.LGphysics.geo-ph

TL;DR

Monsoon precipitation is difficult to model directly because rainfall is sparse and heavy-tailed, while continuous proxies more coherently express convective organization. M-JEPA pretrains a frozen latent rollout on five proxies over moving convective patches, then transfers it to rainfall forecasting, achieving a small resolved CRPS advantage over the operational ensemble while retaining clear scope and rollout limitations.

  • Problem

    Rainfall’s sparse, intermittent, localized, and heavy-tailed statistics make it a difficult target for learning atmospheric dynamics, despite the importance and limited predictability of South Asian monsoon rainfall.

  • Method

    M-JEPA learns seven-day autoregressive latent dynamics from five continuous proxy fields on Lagrangian patches, then decodes the frozen rollout into probabilistic and deterministic rainfall forecasts.

  • Results

    M-JEPA attains fair CRPS 6.81 versus 6.89 mm/day for the 51-member operational ensemble, with positive Brier skill (+0.05 versus −0.04) and gains concentrated at heavy-rain thresholds and fine spatial scales.

  • Takeaways & Limitations

    Cross-variable proxy pretraining captures transferable subseasonal dynamics across disjoint atmospheric targets and supports competitive monsoon precipitation forecasting.

  • Takeaways & Limitations

    The latent rollout saturates by lead 5, and the Lagrangian benchmark covers only 337 windows and 169 matched days, limiting broader generalization.

Abstract

from arXiv · show

Deep atmospheric convection governs South Asian monsoon variability, yet attempting to learn its latent world model directly from zero-inflated, heavy-tailed precipitation yields suboptimal predictive representations. Continuous atmospheric proxies, such as outgoing longwave radiation (OLR), express this convective organization far more coherently. We address this mismatch with \emph{cross-variable proxy learning}: M-JEPA, a multiscale Monsoon Joint-Embedding Predictive Architecture, is pretrained on five continuous proxy fields over Lagrangian patches tracking moving convective systems---without rainfall supervision at any point. The resulting frozen representation is transferred to daily precipitation forecasts through a shared decoder trunk featuring parallel probabilistic and deterministic branches. Because rainfall is strictly unobserved during pretraining, downstream skill directly measures the predictive information captured in the latent rollout. A frozen-backbone probing framework with two controls (an identical architecture trained on rainfall alone, and a randomly initialized backbone) attributes the transfer specifically to proxy pretraining: direct rainfall training exhibits $36\%$ higher CRPS error ($7.52$ vs.\ $5.54$\,mm/day). Against the 51-member operational ECMWF ensemble, the transferred model attains a statistically resolved CRPS advantage ($6.81$ vs.\ $6.89$\,mm/day) and higher Brier skill ($+0.05$ vs.\ $-0.04$) using $15.4$M parameters on a single consumer GPU, concentrated at heavy-rain thresholds and fine spatial scales, while the ensemble retains an advantage in neighborhood skill and deterministic references on point metrics. The result provides a competitive monsoon precipitation forecast grounded in intraseasonal dynamics and a diagnostic framework for evaluating transferred atmospheric representations.

1 Introduction

M-JEPA addresses the difficulty of learning monsoon dynamics from sparse, heavy-tailed rainfall by pretraining on continuous atmospheric proxies and transferring the frozen representation to precipitation forecasting.

  • South Asian monsoon rainfall is consequential yet remains difficult for both numerical and data-driven forecasting systems.
  • Precipitation’s sparsity, intermittency, localization, and heavy tails make it unlike the smooth continuous fields on which machine-learned weather models perform strongly.
  • Continuous fields, especially OLR, express MISO convective organization more coherently than rainfall, motivating cross-variable representation learning.
  • The frozen representation is evaluated through precipitation transfer, with controls separating proxy-learning benefits from architectural inductive bias and decoder effects.
  • M-JEPA learns seven-day latent dynamics from five proxy fields on Lagrangian patches, with rainfall absent from pretraining.

2 Related Work

Prior atmospheric forecasting work emphasizes forecast scores or same-variable reconstruction, leaving the contents and cross-variable transferability of learned representations largely unexamined.

  • JEPA world-model research has focused mainly on video and robotics, where prediction targets share the encoder’s modality.
  • Weather-model benchmarks evaluate forecast skill while leaving open what representations encode and whether they transfer beyond training variables.
  • Monsoon prediction often compresses the system into indices or modes, potentially discarding zonal structure that determines rainfall location.
  • Existing atmospheric latents reconstruct their encoded variables, whereas this study transfers a self-supervised world-model latent to a held-out physical target.

3 Data and Problem Setup

The forecasting setup uses five ERA5 proxy fields and GPM rainfall over tracked, co-moving convective patches, allowing the model to focus on within-system evolution.

  • The model forecasts seven daily precipitation fields on a 1° grid from five daily proxies, with rainfall excluded from pretraining.
  • Lagrangian patches follow individual convective envelopes rather than using fixed Eulerian windows.
  • A retrospective, zero-phase OLR filter identifies and centers MISO envelopes but is never supplied as a model input.
  • Re-centering supplies bulk envelope translation, so the model predicts intensity and internal structure rather than future system position.
  • Inputs are OLR, q, u, v, and 850-hPa geopotential from ERA5, while GPM satellite rainfall is the verification target.

4 Method

M-JEPA first learns multiscale autoregressive latent dynamics from proxy-based Lagrangian patches, then freezes that rollout and decodes it into probabilistic and deterministic rainfall forecasts.

  • 4 Method: M-JEPA uses proxy-only JEPA pretraining followed by a rainfall readout from the frozen latent rollout.
  • 4 Method: Two streams represent foreground OLR at fine scale and a broader background containing humidity, winds, geopotential, and smoothed OLR.
  • 4 Method: The context encoder embeds each day, while a block-causal predictor autoregressively rolls the latent state forward using its own outputs.
  • 4.1 Stage 1: upstream pretraining on proxies: Stage 1 combines teacher-forced and free-running latent regression with VICReg regularization and auxiliary reconstruction; rainfall enters neither the inputs nor objective.
  • 4.2 Stage 2: downstream rainfall readout: Stage 2 conditions a shared decoder trunk on the rollout, rain history, and geography, then reads it with probabilistic and deterministic branches.
  • 4.2 Stage 2: downstream rainfall readout: The probabilistic branch transports Gaussian noise to rainfall in square-root space, while the deterministic branch produces one softplus-constrained field.
  • 4.2 Stage 2: downstream rainfall readout: Sequential branch training avoids jointly optimizing incompatible distributional and point-forecast functionals; the distributionally trained trunk reaches pooled ACC 0.292.
  • 4.2 Stage 2: downstream rainfall readout: Table 1 compares M-JEPA’s 16-member ensemble with the 51-member IFS ensemble using fair CRPS and other probabilistic and deterministic metrics on matched windows.

5 Results

M-JEPA’s probabilistic branch slightly outperforms the operational IFS ensemble on CRPS and Brier skill, while deterministic and neighborhood metrics reveal a sharpness-versus-overlap trade-off. Additional probes and controls show that the frozen proxy-trained latent encodes organized atmospheric structure, future rainfall information, and event evolution, although its rollout trajectory saturates at longer leads.

  • Probabilistic forecast skill: 6.81 vs. 6.89 mm/day CRPS gives M-JEPA Ens a slight resolved advantage over the 51-member IFS Ens, with Brier skill +0.05 vs. −0.04.The CRPS difference is Δ = −0.08, with 95% CI [−0.16, −0.001] over 1,183 paired cases.
  • Probabilistic forecast skill: 1.4× to 1.8× Heavy and 1.3× to 4.1× Very Heavy member-FSS margins increase from lead 1 to lead 7 over the matched population.The Extremely Heavy ordering is lead-dependent, with IFS ahead at the first two leads and M-JEPA ahead from lead 3 onward.
  • Deterministic skill and metric sensitivity: 16.24 vs. 23.46 mm/day RMSE and 0.31 vs. 0.18 spatial anomaly correlation show the deterministic branch exceeding persistence at every lead.Persistence retains a small advantage at the 3° and 5° neighborhood widths, illustrating the recurring sharpness-versus-overlap trade-off.
  • Deterministic skill and metric sensitivity: Dry-area specificity exceeds IFS and GraphCast at every lead and AIFS at six of seven, losing only to AIFS at lead 1.Figure 3 compares placement on a matched target day using deterministic operational members cropped to the same Lagrangian patch.
  • Latent variable structure: 0.67 to 0.85 R2 probe scores recover the five proxy patch means, while mean absolute off-diagonal cosine 0.15 indicates mostly distinct latent directions.The largest interpretable overlaps are geopotential with the OLR envelope (+0.38) and humidity with meridional wind (+0.30).
  • Predictive information in the latent rollout: 0.174 vs. 0.084 pooled R2 shows the rolled latent decodes future rainfall better than the reused context latent, with a +0.090 margin.The pooled result supports predictive information in the rollout, though per-lead intervals are wide and only the pooled margin and lead 1 exclude zero.

6 Limitations and Scaling Outlook

The main limitations concern latent rollout dynamics and the restricted evaluation scope. The paper argues that improving long-horizon behavior requires richer information rather than simply increasing model size.

  • Temporal rollout dynamics: Lead-5 saturation causes decoded forecasts to under-predict the decay of long-lived systems.The paper attributes this damping to squared-error regression contracting toward the conditional mean as uncertainty grows; separating loss effects from representation limits would require an ensemble-valued rollout.
  • Benchmark scope: The benchmark covers only tracked convective patches, with operational references cropped into co-moving frames.The test set contains 337 windows and 169 matched days on a single seed, so broader generalization requires further evaluation.
  • Scaling outlook: Under a 17× reduction in pretraining data, the rain-only control still outperforms three conventional rain-history networks.The authors therefore characterize the regime as information-limited and propose denser, higher-frequency, deeper, and broader proxy inputs rather than more parameters.

7 Conclusion

M-JEPA learns transferable subseasonal dynamics from continuous atmospheric proxies without rainfall supervision. Its frozen rollout improves future-rainfall decoding and achieves small advantages over the ECMWF ensemble on selected probabilistic metrics, while retaining weaknesses on neighborhood and deterministic point metrics.

  • Transferable dynamics: A pooled R2 margin of +0.090 shows that the frozen latent rollout decodes future rainfall better than a static present-state baseline.Direct precipitation training has 36% higher CRPS error: 7.52 versus 5.54 mm/day.
  • Forecast comparison: Against the 51-member ECMWF ensemble, M-JEPA achieves fair CRPS of 6.81 versus 6.89 mm/day and Brier skill of +0.05 versus −0.04.The gains concentrate at heavy-rain thresholds and fine spatial scales, while the ensemble retains advantages in neighborhood skill and deterministic references on point metrics.

A Data and Problem Setup

The study forecasts seven daily precipitation fields on a 1° grid using five daily atmospheric proxies and Lagrangian patches centered on tracked convective systems. Preprocessing separates retrospective tracking geometry from model inputs and uses year-block splits to limit autocorrelation leakage.

  • Data and targets: The target is seven daily precipitation fields over 9.5°S–29.5°N and 60.5–99.5°E at 1° resolution.Inputs are OLR, specific humidity, zonal and meridional wind, and 850-hPa geopotential; the target is satellite precipitation.
  • Patch construction: Lagrangian patches follow tracked convective systems instead of using a fixed Eulerian window.Tracking is based on band-passed OLR, while all fields are regridded, deseasonalized, and z-scored using training-year statistics only.
  • Signal processing: The tracker filters daily OLR in the 20–90 day intraseasonal band with a zero-phase Butterworth filter used only to define patch geometry.The retrospective filter does not enter model inputs.
  • Detection: Deep-convective cores are selected below the daily 10th-percentile band-passed OLR anomaly, smoothed, connected, and filtered by a 30-cell minimum.Object locations use intensity-weighted centroids based on absolute OLR anomalies.
  • Tracking: Tracking uses greedy nearest-cost assignment with distance, fractional-area-change, and time-gap weights, while direction weight wn=0.Tracks can bridge up to four missed detections, with candidate matches restricted to a 5°/day × Δt search radius.
  • Retention: Tracks are retained when they span at least 12 days, covering the five-day context plus seven-day forecast horizon.Strict MISO physical filters are disabled, retaining a larger and more varied population.
  • Data splits: Year-block splits use 1979–2017 for pretraining, 1998–2013 for decoder training, 2014–2017 for validation, and 2018–2025 for testing.The held-out test set contains 337 windows, with ±2-grid-cell spatial-jitter augmentation during decoder training.

B Model Architecture and Training

M-JEPA uses tracked Lagrangian patches and a two-stage frozen-backbone design to transfer proxy-learned latent dynamics to rainfall forecasting. Evaluation uses matched geometries, case-averaged scores, and controlled training procedures.

  • Architecture: M-JEPA follows tracked convective envelopes with Lagrangian patches, while the encoder consumes unfiltered atmospheric fields.The tracker supplies patch geometry and re-centering; the encoder does not consume the displayed bandpassed OLR values.
  • Component Sizes: Stage 1 contains a 9.7M-parameter context encoder and autoregressive predictor, with a target encoder updated by exponential moving average.The context encoder and predictor each use six transformer blocks, eight heads, and width 256.
  • Training Stages: Stage 2 freezes the Stage-1 stack and trains a shared trunk with probabilistic and deterministic rainfall branches.The deterministic branch is trained afterward on the frozen trunk, preventing its loss from updating the trunk.
  • Controls: The controlled comparison applies the same Stage-2 recipe and decoder to frozen proxy-pretrained and rainfall-only backbones.This design isolates the pretraining signal rather than decoder differences, while rain-history baselines are trained end to end.
  • Evaluation: Scores are computed on observation-defined cells within tracked patches, with each window-lead case weighted equally.Probabilistic metrics include locally climatology-referenced Brier skill and member-averaged neighborhood FSS.
  • Evaluation: ACC rewards spatial pattern agreement but is insensitive to amplitude, so smooth forecasts can score well despite losing fine-scale detail.This metric property motivates interpreting ACC alongside neighborhood and pointwise measures.

D.2 Variable Decodability and Subspace Structure: Extended Analysis

The latent representation stores proxy information in linearly accessible, distributed subspaces and evolves through a saturating, multiscale rollout. Attribution and controls connect this latent evolution to rainfall and cross-variable transfer.

  • Variable Decodability: 0.67–0.85 proxy R2 shows strong linear decodability, while tuned nonlinear probes do not significantly improve any proxy readout.The largest nonlinear point gain is only +0.02, and validation years overlap the pretraining span.
  • Subspace Structure: The five proxy directions have mean absolute off-diagonal cosine similarity 0.15, with physically interpretable overlaps among geopotential, OLR, humidity, and winds.The reported couplings include geopotential–OLR at +0.38 and humidity–meridional wind at +0.30.
  • Latent Dynamics: CKA falls from 0.944 at lead 1 to 0.794 at lead 5 before plateauing, indicating rapid latent evolution followed by saturation.The trajectory is computed independently of the downstream target.
  • Latent Dynamics: The top three latent-increment principal components explain 16%, 15%, and 11% of variance, totaling 42%.No single dominant increment axis appears, consistent with distributed evolution across weakly coupled modes.
  • Attribution: Rainfall attribution peaks one token north-east of the tracked center at short leads and one token east at the longest lead.The center token remains high but never maximal, reaching 91–99% of the peak across leads.
  • Controls: A rainfall-only backbone trails the proxy-pretrained model despite nearly identical parameter counts, with CRPS 7.52 against 5.54 mm/day.The comparison holds architecture and self-supervised recipe constant.
  • Cross-Target Transfer: The rolled latent outperforms the static latent for future OLR at every lead, but the widening margin reflects static-control degradation rather than rollout improvement.The static probe falls from −0.017 at lead 2 to −0.293 at lead 6 while the rolled probe remains approximately flat at later leads.

E.4 Leave-One-Proxy-Out Ablation

Leave-one-proxy-out experiments show that all retained proxies contribute non-redundant information, but the contribution is distributed rather than dominated by one variable. Meridional wind is the clearest jointly load-bearing component, while OLR cannot be ablated without changing the evaluation population.

  • Ablation Results: Every leave-one-proxy-out ablation increases CRPS, but the margins remain modest at 0.5–1.8% compared with the 36% five-proxy versus rain-only gap.The pattern supports redundant, distributed encoding rather than one proxy carrying the full signal.
  • Ablation Results: Meridional-wind removal is the only ablation that worsens CRPS, Brier skill, and FSS10 together.Its changes are +1.8% CRPS, −12.8% BSS, and −6.1% FSS10.
  • Physical Interpretation: The humidity–meridional-wind coupling links the ablation result to latent subspace structure, suggesting their joint representation carries moisture-transport information.The cited interpretation is that removing either weakens the signal organizing convection.
  • Ablation Results: Geopotential removal is mildest for BSS and FSS10, which remain unchanged, indicating greater redundancy with the other proxies.Its CRPS effect is +0.63%, compared with +0.50% for humidity.
  • Experimental Scope: OLR is not removed because it defines the tracking frame, so ablation would alter which windows enter evaluation.That would confound input removal with a changed evaluation population.
  • Robustness: Tracker-parameter sensitivity preserves the ordering against persistence and deterministic IFS across all ten configurations.The ensemble comparison was not rerun for each configuration.

F.2 The Archives’ Rising-ACC Pattern: Mechanism and a Full-Domain Check

The archives’ rising ACC with lead is explained by lifecycle-dependent smoothing in tracked patches, not improving long-lead skill. Full-domain rescoring removes this pattern, while latent rollouts retain damped, non-zero motion.

  • Mechanism: AIFS and GraphCast show ACC rising while RMSE falls with lead, a physically backwards forecast-skill pattern.AIFS ACC rises 0.376 → 0.458 as RMSE falls 21.48 → 16.63; GraphCast shows the same direction.
  • Mechanism: ACC is scale-invariant, so filtering fine-scale detail can raise correlation with the retained large-scale pattern while real skill declines.Damping and roughness diagnostics test whether increasing smoothness explains the lead trend.
  • Mechanism: Case-level correlations between ACC and observed roughness are negative, supporting lifecycle roughness decay as the direct explanation for the pooled trend.The relationship is tested for every window–lead case rather than only seven lead aggregates.
  • Full-domain check: Full-domain rescoring makes ACC flat and near zero while RMSE mildly increases with lead, eliminating the rising-ACC pattern seen in tracked patches.Full-domain ACC is ≲0.04, versus 0.29–0.46 in the patch frame.
  • Motion diagnostics: Decoded centroid drift matches observed northward or southward direction above chance across leads, providing a motion diagnostic beyond placement scores.Directional accuracy is evaluated over 337 held-out windows using rain-weighted centroids.
  • Rollout dynamics: Latent rollout displacement falls from 2.25 to a non-zero 1.33–1.38 floor, so CKA saturation reflects stabilized structure rather than a fixed latent state.The first-to-last displacement ratio is ≈1.69×, with the curve flat over leads 5–7.

G Scaling Hypothesis: Information-Limited Signal

The results suggest that transferable skill is information-limited: proxy window count and temporal cadence constrain what M-JEPA can encode. The model remains computationally compact, but comparisons with references confound cadence, initialization, and training-corpus scale.

  • Information-limited signal: The rainfall-only control uses 226 pretraining windows versus 3885 for five-proxy pretraining, a 17× reduction in information.Despite the smaller signal, its CRPS degrades gracefully rather than collapsing.
  • Information-limited signal: The control comparison supports a proxy-window-count hypothesis, not a measured scaling law across multiple pretraining sizes.A full scaling curve over window count, channel count, and temporal resolution remains future work.
  • Temporal resolution: M-JEPA discards sub-daily atmospheric structure by using daily proxy means and one-day latent steps.This includes diurnal convective initiation and nocturnal mesoscale-convective propagation.
  • Temporal resolution: Temporal cadence cannot be isolated as the cause of reference-model differences because initialization quality and training-corpus scale also differ.The direct test is a controlled ablation varying proxy cadence while matching initialization and data volume.
  • Computational cost: The model totals 15.4M parameters, with a seven-day deterministic batch forecast taking 204 ms on one GPU.The deterministic head requires a single forward pass rather than iterative sampling.
  • Evaluation scope: Table S14 reports full-domain RMSE and spatial anomaly correlation on reduced complete-field subsets, unlike patch-frame scores centered on tracked systems.The full-domain requirement reduces coverage to 145/63/63 windows versus 161/66/66.
Loading 2608.22358v1…