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

K M Anirudh; S Sandeep; Hariprasad Kodamana

Figure 1: The two stages of M-JEPA. Stage I (top): context encoder

f_{\theta}

, target encoder

f_{\xi}

, autoregressive predictor

g_{\phi}

, pretrained on proxies with no rain objective. Stage II (bottom): the frozen rollout

\hat{z}_{t+1:t+7}

, rain history

y_{t-4:t}

and static geography

s

condition one shared decoder trunk whose two branches emit a

16

-member probabilistic forecast and a deterministic field.

	Members	CRPS $\downarrow$	BSS $\uparrow$	FSS10 $\uparrow$	Spread/RMSE $\to 1$	Rank dev. $\downarrow$
M-JEPA Ens	16	6.81	+0.05	0.53	0.73	0.60
IFS Ens	51	6.89	$-0.04$	0.63	0.45	0.95
IFS (det.)	1	9.76	$-0.36$	0.73	0.00	0.70

Table 1: Forecast skill on matched windows against the

51

-member operational IFS ensemble (

n{=}169

, fair CRPS in mm/day; the IFS (det.) row is on its own single-member convention). The CRPS difference against IFS Ens is

-0.08

mm/day,

95\%

paired-bootstrap CI

[-0.16,-0.001]

over

10^{4}

resamples. Boldface marks the better ensemble. Methodology: supplementary Section C.

Figure 2: Probability-branch diagnostics, matched

n{=}169

population. (a) Member-FSS at IMD Heavy/Very Heavy thresholds by lead. (b) Radial power spectral density, M-JEPA Ens vs. IFS Ens vs. observed.

System $\uparrow$	1	2	3	4	5	6	7
M-JEPA (IFS subset)	0.555	0.534	0.516	0.503	0.503	0.492	0.475
IFS	0.540	0.477	0.460	0.452	0.462	0.463	0.454
M-JEPA (AIFS/GC subset)	0.516	0.532	0.521	0.493	0.487	0.480	0.457
AIFS	0.526	0.497	0.479	0.471	0.443	0.427	0.430
GraphCast	0.424	0.414	0.400	0.387	0.378	0.366	0.377

Table 2: Dry-area specificity (correct-negative rate on truly-dry cells,

2

mm threshold) by lead. AIFS and GraphCast share one matched subset, so a single M-JEPA row serves both.

Figure 3: Five-model spatial comparison on a matched target day (2024-08-04). Columns: GPM observations, M-JEPA, IFS, AIFS, GraphCast; rows: forecasts issued

1

,

3

,

5

,

7

days ahead. Operational fields are single deterministic members cropped at the same Lagrangian patch, so this compares placement, not probabilistic skill; dashed box marks the native patch extent.

Pretraining signal	Params	CRPS $\downarrow$	BSS $\uparrow$
Five proxies	9.7M	5.54	+0.18
Rainfall only, same recipe	9.6M	7.52	$-0.04$
U-Net (rain history)	351K	7.38	$-0.62$
ConvLSTM (rain history)	22K	8.33	$-1.07$
ViT (rain history)	836K	8.65	$-0.99$

Table 3: Isolating the source of transfer. The first two rows share architecture, recipe and probability head, so only the pretraining signal differs; the last three are conventional networks given the same rainfall history. CRPS in mm/day, no recalibration, on the full test set (

n{=}337

) except the rain-only control, which rainfall sparsity leaves a smaller held-out population (

n{=}118

). Table 1 ’s

6.81

is this same model on the

n{=}169

windows the IFS ensemble also covers.

Figure S1: Lagrangian tracking of one long-lived convective envelope (track 2766, genesis 2021-09-07, ocean, equatorial Indian Ocean; 45-day duration; net northward displacement

+25.6^{\circ}

). ( a ) The tracked centroid trajectory from genesis to terminus, colored by elapsed time; diamonds mark the six days shown at right. ( b ) The

20^{\circ}\times 20^{\circ}

bandpassed OLR anomaly patch centered on the moving centroid at each marked day, its border colored to match that day on the trajectory, illustrating how the Lagrangian frame re-centers on the system as it propagates from the equatorial Indian Ocean to the Bay of Bengal. This track is illustrative rather than typical, since the retained test population’s median net northward displacement is

+3.5^{\circ}

. This patch geometry, not the raw OLR values shown here, is what the tracker supplies; the encoder itself consumes unfiltered fields (Section A).

Figure S2: Lagrangian trajectories of all 149 convective OLR objects tracked in the held-out test population (genesis years 2018–2025, JJAS), the population underlying every forecast-skill number reported in this paper. Black circles mark genesis and red crosses mark terminus, with genesis concentrated across the equatorial Indian Ocean and Bay of Bengal; net northward displacement is positive on balance (median

+3.5^{\circ}

, Section A). The 1979–2017 self-supervised pretraining population (711 tracks, of which the 2014–2017 subset also serves as validation) shows the same qualitative spatial pattern.

	Stage 1 (pretraining)	Stage 2 (readout)
Objective	Eq. (4)	Eq. (9)
Optimizer	AdamW	AdamW
Learning rate	$5\times 10^{-4}$	$10^{-3}$
Weight decay	$0.05$	$10^{-4}$
Warmup	$10$ epochs	none
Epochs	$100$	$80$

Table S1: Optimization settings for both training stages. Equation numbers refer to the main paper. Stage 1 sees no rainfall at any point; Stage 2 trains only the trunk and the two branches, on the frozen Stage-1 rollout.

Model	Architecture	Training
M-JEPA	Multiscale transformer context encoder ( $4.8$ M) + autoregressive predictor ( $4.9$ M), both frozen at readout; shared decoder trunk with two branches ( $5.7$ M)	Self-supervised on proxies only, no rain objective; trunk and branches trained on paired rain data, $667$ train windows
Rain-only control	Identical architecture, backbone $9.6$ M against M-JEPA’s $9.7$ M (only the tokenizer’s input-channel count differs)	Same recipe as M-JEPA, applied to rainfall alone; $226$ pretrain windows
U-Net	Encoder-decoder, 2 downsample / 2 upsample conv stages, skip connections, 32 base channels (351K params)	AdamW, lr $10^{-3}$ , weight decay $10^{-4}$ , 60 epochs, rain-history input only
ConvLSTM	Single ConvLSTM cell, 24 hidden channels, closed-loop autoregressive rollout (22K params)	AdamW, lr $10^{-3}$ , weight decay $10^{-4}$ , 60 epochs, rain-history input only
ViT	Patch embedding (patch $4$ , dim $128$ ), 4-layer/4-head transformer encoder (836K params)	AdamW, lr $10^{-3}$ , weight decay $10^{-4}$ , 60 epochs, rain-history input only
IFS	Physics-based NWP, operational	Not trained by us; public open-data archive

Table S2: Model and training configuration for every system compared in this paper. All M-JEPA, rain-only-control, U-Net, ConvLSTM and ViT models are trained on the identical Lagrangian patches and year-block splits (Section A) so that comparisons isolate the modeling choice, not the data. IFS, AIFS, and GraphCast are external references, evaluated but not trained.

Figure S3: Rain decodability (

R^{2}

, linear probe) at each lead from the rolled latent (forecast) vs. the frozen present-state control (no forecast). The rolled latent decodes future rainfall better at every lead, by the widest margin at lead 3. Shaded region is the paired rolled-minus-frozen margin, whose pooled value and

95\%

bootstrap interval (

10^{4}

resamples) are annotated; this paired margin, not the two marginal curves separately, is the statistic the claim rests on.

Lead (days)	Rolled $R^{2}$ $\uparrow$	Static $R^{2}$ $\uparrow$	Margin $\uparrow$	Margin 95% CI
1	0.341	0.261	+0.079	$[+0.004,+0.169]$
2	0.198	0.088	+0.110	$[-0.003,+0.246]$
3	0.105	$-0.034$	+0.139	$[-0.003,+0.292]$
4	0.092	0.022	+0.070	$[-0.076,+0.227]$
5	0.169	0.086	+0.083	$[-0.051,+0.226]$
6	0.169	0.070	+0.099	$[-0.048,+0.263]$

Table S3: Rolled vs. static-latent R 2 per lead. Margin peaks at lead 3, consistent with a forecast whose advantage over persistence erodes at long range. The

95\%

confidence interval is from a

10^{4}

-replicate paired bootstrap over test windows; the pooled (lead-averaged) margin excludes zero, while individual leads overlap zero given the modest test-window count.

Latent Dimension	Linear $R^{2}$ $\uparrow$	CV-tuned MLP $R^{2}$ $\uparrow$	Gap [95% CI]
$q$ (humidity)	0.697	0.644	$-0.052\ [-0.105,-0.007]$
$u$ (zonal wind)	0.668	0.642	$-0.026\ [-0.082,+0.029]$
$v$ (merid. wind)	0.679	0.668	$-0.011\ [-0.062,+0.039]$
$z$ (geopotential)	0.849	0.865	$+0.016\ [-0.005,+0.037]$
OLR (smoothed)	0.744	0.763	$+0.018\ [-0.021,+0.060]$

Table S4: Linear vs. CV-tuned nonlinear probe R 2 for the five proxy readout targets, with the paired-bootstrap (

10^{4}

-replicate)

95\%

CI on the tuned-minus-linear gap. Proxies are strongly and linearly decodable, and the tuned MLP does not significantly exceed the linear probe on any of them. Selected MLP configurations:

q,u

width

512

/depth

1

;

v

width

128

/depth

1

;

z

width

256

/depth

1

; OLR width

256

/depth

2

.

Lead (days)	CKA (context vs. rollout)
1	0.944
2	0.872
3	0.825
4	0.804
5	0.794
6	0.793

Table S5: CKA between the context latent and the rolled latent at each of the seven rollout steps (leads 1–7). Rapid initial drop (evolution), then plateau (saturation).

Figure S4: The readout draws on tokens around, not exactly on, the tracked convective centroid. ( a , b ) Token-level attribution (gradient of expected rainfall with respect to the rolled foreground latent tokens, native

5\times 5

token grid, no smoothing, arbitrary units, shared color scale) at leads 1 and 7; the star marks the peak token and the cross the patch’s geometric center; north is up and east is right. The peak sits one cell north-east of center at lead 1 and one cell east of center at lead 7. ( c ) Radially binned mean input-space saliency (gradient with respect to the raw

20\times 20

foreground field) against distance from the patch center, per lead: increases toward the patch edge rather than peaking at a fixed ring, a pattern we attribute to the decoder’s convolutional boundary handling rather than to convective structure.

Model	Params	CRPS (mm/day) $\downarrow$	FSS10 $\uparrow$	BSS $\uparrow$
M-JEPA (proxies)	9.7M	5.54	0.53	0.18
Rain-only control (same recipe)	9.6M	7.52	0.54	$-0.04$
U-Net (rain only)	351K	7.38	0.14	$-0.62$
ConvLSTM (rain only)	22K	8.33	0.00	$-1.07$
ViT (rain only)	836K	8.65	0.42	$-0.99$

Table S6: Rain-history baselines against the primary model and the same-recipe rain-only control, on the same tracked patches and year-block splits, case-averaged and without recalibration. Parameter counts are the self-supervised backbone for the two JEPA rows (context encoder plus predictor) and the full trained network for the three baselines. All rows are scored on the full test set (

n{=}337

windows) except the rain-only control, whose sparse rain-valid data leave it a smaller held-out population (

n{=}118

), a consequence of the very sparsity the control is designed to expose. The three standard architectures span two orders of magnitude in parameter count (22K–836K) yet all have negative Brier skill, below the climatological reference at the 1 mm threshold, so their failure is one of information, not capacity. The rain-only control’s Brier skill is closest to zero among the rain-only models, which supports the main-paper claim that the self-supervised recipe contributes substantively once given adequate information.

Lead (days)	Trained $R^{2}$ $\uparrow$	Random $R^{2}$ $\uparrow$	Ratio $\uparrow$	Gap $\uparrow$	Gap 95% CI
1	0.336	0.040	8.4×	+0.296	$[+0.183,+0.415]$
2	0.189	0.004	42.7×	+0.185	$[+0.078,+0.294]$
3	0.086	0.018	4.9×	+0.068	$[-0.060,+0.200]$
4	0.124	0.024	5.3×	+0.101	$[-0.018,+0.218]$
5	0.132	0.048	2.7×	+0.083	$[-0.037,+0.206]$
6	0.187	0.102	1.8×	+0.085	$[-0.030,+0.205]$

Table S7: Trained vs. random-init backbone R 2 . The trained column is the same rolled-latent predictive-information probe reported in main-paper Section 5.5, re-run alongside its random-init control (see text for the run-to-run spread). Learning multiplies decodability by 5–43× over leads 1–4, converging to

\approx

1.8× at the longest leads as random-feature decodability of a smooth, seasonally autocorrelated target rises; residual random skill is not attributed to the representation. The gap

95\%

CI is a

10^{4}

-replicate paired bootstrap over test windows computed under the original pairing; the pooled gap and the near-term leads exclude zero.

Lead (days)	Rolled $R^{2}$ $\uparrow$	Static $R^{2}$ $\uparrow$	Margin $\uparrow$	Margin 95% CI
1	0.454	0.424	+0.030	$[-0.058,+0.121]$
2	0.163	$-0.017$	+0.180	$[+0.047,+0.322]$
3	$-0.145$	$-0.199$	+0.053	$[-0.124,+0.226]$
4	$-0.079$	$-0.170$	+0.090	$[-0.081,+0.256]$
5	$-0.125$	$-0.234$	+0.110	$[-0.096,+0.298]$
6	$-0.078$	$-0.293$	+0.215	$[-0.001,+0.419]$

Table S8: Rolled vs. static-latent R 2 for future-OLR anomaly (cross-target transfer). Margin positive at every lead; absolute skill negative beyond lead 2, but the relative advantage of the rollout persists and widens. The

95\%

CI is a

10^{4}

-replicate paired bootstrap over test windows; the pooled margin excludes zero.

Figure S5: Controls isolating the source of transfer (main-paper Section 5.6), shown visually. ( a ) Trained encoder vs. randomly initialized backbone of identical architecture: both tested on the same predictive-information probe (rolled-vs-static R 2 ). The trained model separates from the random one at every lead; random features retain residual long-lead decodability (

\approx 0.10

by leads 6–7, indicating target smoothness), while trained skill stays at least

1.4\times

higher at every lead, confirming the transfer is a product of learning. ( b ) Cross-target transfer: rolled-vs-static R 2 margin for a held-out second target (future OLR anomaly, never seen during rain-decoder training). Margin larger than rainfall (

+0.126

vs.

+0.090

) and broadens with lead (

+0.03

near-term,

+0.20

long-range), though absolute skill drops negative at long leads. Demonstrates the pretrained representation encodes general atmospheric state, not rain-specific shortcuts.

Proxy set	CRPS $\downarrow$	BSS $\uparrow$	FSS10 $\uparrow$
All five (full model)	5.78	0.18	0.60
$-$ humidity ( $q$ )	5.80	0.17	0.58
$-$ zonal wind ( $u$ )	5.81	0.17	0.59
$-$ meridional wind ( $v$ )	5.88	0.16	0.56
$-$ geopotential ( $z$ )	5.81	0.18	0.60

Table S9: Leave-one-proxy-out ablation, pooled convention, full test population. These runs predate the switch to case-averaged scoring and are internally comparable to each other, not to the case-averaged

5.54

of main-paper Table 3. Every removal costs CRPS, but the margins are modest (

0.5

–

1.8\%

) rather than large. Percentages here and in the text are computed from the unrounded scores, not from the two-decimal entries shown, which round to slightly different ratios. Meridional wind is the one consistently load-bearing variable, the only ablation that also costs BSS and FSS10 together; geopotential’s removal leaves BSS and FSS10 essentially unchanged (

-1.0\%

and

+0.1\%

).

System	RMSE $\,\downarrow$	ACC $\,\uparrow$	FSS $1^{\circ}$	FSS $3/5^{\circ}$
M-JEPA	16.24	0.31	0.35	0.47/0.53
Ensemble mean	16.48	0.29	0.34	0.48/0.55
Persistence	23.46	0.18	0.33	0.49 / 0.56

Table S10: Deterministic skill on the full test set (

n{=}337

), case-averaged over the seven-day lead window (each window/lead scored individually, then averaged; ACC: spatial anomaly correlation; FSS at the

10

mm threshold ( 35 ) ), M-JEPA scored with the deterministic branch. Ensemble mean averages the probability branch’s

16

raw members into a point forecast, scored identically, to check whether a dedicated point-forecast head earns its keep over simply averaging the generative samples. Boldface marks the best value. Matched-subset comparison against IFS, AIFS, and GraphCast: Table S12 ; per-lead detail below.

RMSE (mm/day) $\downarrow$
Lead (days)	M-JEPA (det)	Ensemble mean	Persistence
1	16.79	17.20	21.12
2	17.13	17.46	23.97
3	16.88	17.12	24.30
4	16.39	16.62	23.96
5	15.99	16.13	23.73

Table S11: Per-lead breakdown of Table S10 ’s comparison against persistence, same convention and routines throughout (

10

mm threshold,

5^{\circ}

neighborhood). Persistence is rebuilt here as the last observed context field at the fixed anchor; the lead-averaged values reproduce Table S10 ’s upper block exactly.

IFS matched subset ( $n{=}266$ ); M-JEPA re-scored
System	RMSE $\,\downarrow$	ACC $\,\uparrow$	FSS $1^{\circ}$	FSS $3/5^{\circ}$
M-JEPA	18.23	0.283	0.34	0.46/0.52
Ensemble mean	18.50	0.255	0.33	0.47/0.54
IFS	18.45	0.328	0.44	0.60/0.67
AIFS/GraphCast matched subset ( $n{=}127$ ); M-JEPA re-scored
M-JEPA	19.74	0.339	0.37	0.49/0.56

Table S12: Deterministic skill, case-averaged over the seven-day lead window, M-JEPA re-scored on each reference’s own matched-coverage subset (IFS:

n{=}266

; AIFS and GraphCast:

n{=}127

, an identical subset by construction, reported once). Ensemble mean as in Table S10 . Every entry is the mean of the corresponding per-lead column of Table S13 , each lead carrying equal window count, so the two tables reconcile term by term. The anomaly correlation uses a per-lead climatology throughout, as elsewhere in this paper; a climatology pooled across leads raises every ACC in this table by

0.005

–

0.010

without changing any ordering.

vs. IFS ( $n{=}161$ )
Lead	RMSE $\downarrow$			ACC $\uparrow$			FSS $\uparrow$
	det.	ens.	ref	det.	ens.	ref	det.	ens.	ref
1	19.03	19.49	20.18	0.410	0.366	0.367	0.725	0.723	0.759
2	19.27	19.64	19.49	0.326	0.284	0.355	0.624	0.636	0.717
3	18.97	19.19	18.98	0.282	0.259	0.349	0.535	0.555	0.695
4	18.41	18.72	18.57	0.248	0.220	0.319	0.481	0.492	0.666

Table S13: Cross-model skill by lead (per-lead detail behind Table S12 above), restricted to the cells M-JEPA’s tracked patches cover, the same patch-native frame as main-paper Table 1 and main-paper Figure 3. RMSE (mm/day)

\downarrow

, ACC

\uparrow

, FSS

\uparrow

(

10

mm,

5^{\circ}

), case-averaged. “det.” is the deterministic head; “ens.” is the probability branch’s

16

raw members averaged into a point forecast (Table S10 ’s ensemble-mean row). Window count per lead:

161

(IFS),

66

(AIFS/GraphCast); pooled across leads these give the

266

/

127

distinct target days of Table S12 . All three references lead on ACC and FSS at every lead. On RMSE, the deterministic head is lower than IFS at leads 1–6 (IFS regains it only at lead 7) but never leads AIFS or GraphCast; the ensemble mean’s RMSE advantage over IFS is narrower still (leads 1 and 5 only) and never leads AIFS/GraphCast.

Figure S6: Amplitude-stability comparison among actual forecasting systems (persistence excluded, it relays a stale, systematically too-large field and runs to

1.31

by lead 7, off this axis): damping ratio

\mathrm{std}(\mathrm{pred})/\mathrm{std}(\mathrm{obs})

by lead, patch frame, case-averaged. IFS retains the most amplitude (

\approx 0.55

–

0.60

, closest to

1

), M-JEPA and AIFS track a comparable, non-collapsing band (

\approx 0.44

–

0.55

), and GraphCast sits lowest throughout (

\approx 0.41

–

0.46

), consistent with its already-smooth-from-lead-1 signature. M-JEPA shows no distinct advantage on this diagnostic; the point is that it does not exploit the amplitude-collapse route to inflated ACC that AIFS’s declining ratio suggests.

Figure S7: The ACC-mechanism, tested directly: every (window, lead) case’s ACC vs. that case’s own observed-field roughness, pooled over leads, with fitted trend and Pearson

r

per archive.

	IFS ( $n{=}145$ )		AIFS ( $n{=}63$ )		GraphCast ( $n{=}63$ )
Lead	RMSE	ACC	RMSE	ACC	RMSE	ACC
1	19.73	$-0.020$	20.52	0.020	19.24	0.032
2	19.78	$-0.017$	20.14	0.029	19.19	0.040
3	20.03	$-0.026$	19.96	0.025	19.19	0.035
4	20.26	$-0.035$	19.95	0.007	19.36	0.023
5	20.34	$-0.030$	20.21	0.003	19.89	0.017

Table S14: Full-domain (

40\times 40

) rescoring of the reference archives, case-averaged, on the Table S13 target days for which a complete full-domain reference field exists; that requirement drops roughly a tenth of the windows (

145/63/63

against

161/66/66

), since a patch-frame score needs only the patch cells whereas a full-domain score needs the whole grid. RMSE (mm/day)

\downarrow

, spatial anomaly correlation

\uparrow

. Unlike the patch-frame comparison, no patch is centered on the tracked system here, so any lifecycle-decay effect specific to the tracked system cannot appear.

RMSE (mm/day) $\downarrow$
Lead (days)	Ensemble mean	Persistence
1	17.20	21.12
2	17.46	23.97
3	17.12	24.30
4	16.62	23.96
5	16.13	23.73

Table S15: Ensemble mean vs. persistence by lead, full test set (

n{=}337

), case-averaged. Unlike the pooled convention, the ensemble mean’s ACC now exceeds persistence’s at every lead, including lead 1. FSS10 favors the ensemble mean at leads 2–3 only and persistence at leads 1 and 4–7, the same crossover pattern as the deterministic branch (Table S11 ).

Figure S8: GPM observations (left) and the M-JEPA deterministic field (right) at days

1,3,5,7

of the seven-day forecast for the wettest test window by observed day-3 rainfall (target days 2023-07-02 to 2023-07-08). The forecast places a maximum on the observed coastal cell at short range and broadens into a smoother band by day 7. Dashed box marks the native patch extent.

Figure S9: Pooled evolution diagnostics over the full test set (

n{=}337

), rain-centroid based (

1

cell

=1^{\circ}

;

95\%

window-bootstrap bands where applicable). ( a ) Northward/southward directional accuracy of the decoded centroid’s drift relative to lead 1 (sign match against observed), at three observed-drift thresholds: above the

50\%

chance line at every lead and every threshold, and rising as the threshold tightens (near-zero-displacement noise excluded). ( b ) Step-to-step latent rollout displacement

\|\hat{z}_{d}-\hat{z}_{d-1}\|_{2}

against lead (pooled foreground latent; the open marker and dashed segment mark the context

\to\hat{z}_{1}

step, a different regime from the subsequent

\hat{z}_{d-1}\to\hat{z}_{d}

steps). Displacement decays from

2.25

to a non-zero floor of

1.33

–

1.38

by leads 5–7 (

\approx 1.69\times

damping, early against late step): the latent rollout is damped but not frozen. ( c ) Composite intensity lifecycle: observed and decoded area-mean rain by day since genesis, pooled over

49

test tracks (window-bootstrap bands; days with fewer than

20

pooled samples omitted). Normalized to each series’ own day-5 value, the decay rates over days 5–16 are

-0.0247

and

-0.0253

per day on a dry bias of

2.1

mm/day; in the thinner tail of the longest-lived tracks (day

17

onward) the decoded curve stops declining while the observed curve continues to fall. Panels (b) and (c) are two signatures of the same rollout damping, on undecoded latent velocity and decoded amplitude respectively; panel (a) is a directional result and does not speak to damping.

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

Abstract

Figures & tables

Explore similar work

SimCast-S2S: A Computationally Efficient Diffusion Model for Subseasonal Precipitation Forecasting

Does Aurora Encode Atmospheric Structure? Latent Regime Analysis and Attribution

FlexCast: Adaptive Weather Forecasting from Arbitrary Field Sets