cs-labs/experiment-6-v1-forecast-model

Experiment 6 — Forecast DiT (Step 5)

Production Forecast DiT model for the Experiment 6 weather prediction pipeline.

Model Details

• Architecture: WeatherDiT, 560M params
• Task: 6-hour ahead atmospheric state forecast in VAE latent space
• Training data: 2005-2014 (10 years), RKSI 300x300 1km snapshots
• Hardware: 8x NVIDIA B200, bf16 mixed precision
• Optimizer: Muon + AdamW (dual), gradient_accumulation_steps=4
• Flow matching: shift=1.0 (rectified flow)
• Best epoch: 10/60 (training terminated at epoch 11)
• Best val T2 MAE: 0.7838 K
• Best val FM loss: 0.4537

Training Curve

Epoch	train_loss	val_fm_loss	val_t2_mae	val_t2_bias
1	1.966	1.171	1.996 K	—
3	1.284	1.072	1.455 K	—
5	0.726	0.942	1.097 K	+0.282
7	0.334	0.775	0.979 K	+0.003
8	0.273	0.662	0.927 K	-0.411
9	0.231	0.551	0.957 K	-0.542
10	0.196	0.454	0.784 K	-0.461
11	0.163	0.383	0.888 K	-0.797

Files

• best_step2_forecast_dit.pt — Best checkpoint (model weights only, epoch 10)

Usage

Load with the Experiment 6 inference pipeline:

import torch
from shared.models.dit.model import WeatherDiT

dit = WeatherDiT(...)  # see step_5_forecast_dit/master_config.py for arch params
dit.load_state_dict(torch.load("best_step2_forecast_dit.pt", map_location="cpu", weights_only=True))
dit.eval()

Dependencies

• Requires VAE from cs-labs/experiment-6-v1-vae-model
• Requires Interpolation DiT from cs-labs/experiment-6-v1-interpolation-model
• Requires Zarr dataset from cs-labs/experiment-6-v1-dataset-large-21-years