Open research Research only

Weather-JEPA

Learning the atmosphere before predicting it

Weather-JEPA is an open research program investigating whether reconstruction-free joint-embedding predictive learning can produce atmospheric representations that transfer across tasks, remain stable through time, characterize uncertainty, and preserve hazardous-weather structure under accessible compute constraints.

The long-term purpose is to contribute knowledge that may improve the anticipation and characterization of dangerous weather and ultimately support systems that protect human life.

Research-only notice: Weather-JEPA is not an operational forecast or warning system. Its outputs must not be used for emergency, aviation, marine, navigation, evacuation, or other life-safety decisions.

View the repository · Read the North-Star Plan · Documentation decision

The central question

Modern machine-learning weather systems can produce skillful forecasts. Weather-JEPA asks a narrower question:

Under matched architecture, data, parameters, and compute, does predicting latent atmospheric representations produce more transferable and dynamically useful knowledge than predicting raw future fields or reconstructing masked values?

The project compares three objectives:

Approach What the model predicts
Supervised forecasting Future physical weather fields
Masked autoencoding Missing physical values
Weather-JEPA Hidden or future target embeddings

This comparison is designed to isolate the learning objective. The project will not change the architecture, data, optimizer, resolution, and evaluation simultaneously and then congratulate the objective for the combined result.

Why representation learning matters

A forecast model may be accurate on average while remaining fragile under missing inputs, weak on transfer tasks, unstable during long rollout, poorly calibrated, or insensitive to a small but dangerous event.

Weather-JEPA therefore treats forecast error as one part of evaluation rather than the entire definition of success.

A useful atmospheric representation should support:

The project will also test the possibility that JEPA does not improve these capabilities. Negative results are part of the research record.

Scientific objectives

Representation quality

Determine whether JEPA representations improve frozen linear probes, shallow probes, few-shot learning, and transfer across seasons, regions, variables, or tasks.

Forecast skill

Measure field-level performance across variables, pressure levels, lead times, seasons, and regions.

Dynamical stability

Evaluate decoded and latent autoregression for error growth, variance behavior, spectral drift, front loss, unphysical extrema, and phase error.

Robustness

Stress the model with missing variables, masked areas, noisy inputs, temporal gaps, metadata errors, and distribution shift.

Hazard relevance

Evaluate selected events such as tropical cyclones, atmospheric rivers, blocking regimes, strong-wind systems, and extreme precipitation where labels and data support defensible analysis.

Compute efficiency

Report performance per parameter, training hour, peak memory, processed sample, and inference time.

Probabilistic utility

After deterministic stability is established, test whether the latent predictor can generate calibrated and spatially coherent alternative futures.

Mathematical model

The pressure-level atmosphere at time $t$ is represented as a tensor with dimensions for variables, vertical levels, latitude, and longitude:

\[\mathbf{X}_t \in \mathbb{R}^{C \times L \times H \times W}\]

Surface and static fields are represented separately:

\[\mathbf{S}_t \in \mathbb{R}^{C_s \times H \times W}, \qquad \mathbf{B} \in \mathbb{R}^{C_b \times H \times W}\]

A context window contains multiple recent atmospheric states. A trainable context encoder maps the visible history to latent tokens. An exponential-moving-average target encoder maps hidden or future atmospheric states to target tokens. A horizon-conditioned predictor estimates those target representations.

The core objective minimizes distance between predicted and target embeddings while stopping gradients through the target encoder:

\[\mathcal{L}_{\mathrm{JEPA}} = \operatorname{mean}_{k} \left\| \widehat{\mathbf{z}}_k - \operatorname{stop\_gradient} \left( \mathbf{z}^{\mathrm{target}}_k \right) \right\|_p\]

Collapse is not assumed away. The project measures latent variance, covariance, effective rank, singular values, token similarity, and downstream probe performance. Variance or covariance regularizers may be introduced only when those diagnostics justify them.

Global losses and metrics use area-aware latitude weighting. Data normalization uses training-split statistics only.

Reference architecture

atmospheric state
      |
      v
validation + normalization + temporal windows
      |
      v
variable / level / static / calendar / horizon embeddings
      |
      v
spatial patch embedding
      |
      v
factorized vertical-variable mixer
      |
      v
multiscale context encoder --------------------------.
      |                                             |
      v                                             v
horizon-conditioned JEPA predictor            EMA target encoder
      |                                             |
      '--------------- latent loss -----------------'
      |
      +--> physical decoder
      |
      +--> transfer probes and hazard heads

The initial architecture uses:

Model tiers progress from micro verification models to workstation-scale experiments. Larger is not a phase gate.

Data strategy

Primary sources

ERA5 is treated as a reanalysis produced by observations, physical modeling, and data assimilation—not as direct atmospheric truth.

Initial variables

Pressure-level variables:

Reference pressure levels:

1000, 925, 850, 700, 500, 300, 250, 200, 100 hPa

Surface and near-surface candidates include mean sea-level pressure, 2 m temperature, 10 m winds, total-column water vapor, surface pressure, and carefully defined accumulated precipitation.

Scale progression

synthetic tensors
    -> tiny geographic subset
    -> regional coarse resolution
    -> regional moderate resolution
    -> global coarse resolution
    -> higher resolution only after profiling

Provisional chronology

Split Years
Training 1979–2017
Validation 2018–2019
Test 2020–2022
Untouched challenge 2023 onward, subject to data maturity

There is no random frame-level split. Context windows cannot cross split boundaries. Complete event lifecycles remain within one split. Validation and test periods never determine normalization statistics.

Development roadmap

Foundation

Mathematical verification

Baselines

Weather-JEPA

Dynamics and hazards

Probability, scaling, and release

Read the full development gates

Evaluation

The project reports more than one convenient global average.

Forecast fields

Structure and dynamics

Representations

Hazards

Probability

Efficiency

Research and engineering standards

Scientific integrity

AI engineering

Data engineering

Software engineering

Current phase

Weather-JEPA is presently in the research-specification and repository-foundation phase.

Current governing artifacts:

The next executable gates are:

  1. initialize the repository;
  2. add license, security, contribution, citation, and changelog files;
  3. define the Python environment and package;
  4. configure formatting, linting, typing, tests, and CI;
  5. verify a PyTorch MPS smoke test;
  6. freeze the first research protocol;
  7. build a small, traceable ERA5 subset.

No model result will be presented before the data and software foundations can reproduce it.

Research Artifacts

Open science

Planned public artifacts include:

Large third-party datasets will not be committed to Git. Their manifests, retrieval instructions, transformations, and licensing requirements will be documented.

Safety and limits

Weather-JEPA is intended for atmospheric machine-learning research and education.

It is not intended for:

The project may investigate methods that could eventually support safer forecasting. It will not claim to have protected lives merely because that is the mission.

Participate

Researchers, meteorologists, data engineers, software engineers, reviewers, and educators are welcome to contribute after the repository standards are published.

A contribution should identify:

The project’s governing plan remains the final reference when a README summary and an implementation decision disagree.

License and citation

The code is intended for release under the Apache License 2.0. Data and third-party artifacts retain their original licenses.

A versioned CITATION.cff and archived DOI will accompany the first research release.

North star

Weather-JEPA will not be judged by the elegance of its latent projections, the size of its model, or the enthusiasm of its public posts.

It will be judged by whether it produces reproducible evidence about atmospheric representation learning, explains its failures honestly, and contributes knowledge that can be used responsibly in the long chain from weather observation to human protection.