JEPA Latent Prediction: exploratory 2hr result (1.2478 val_bpb, 3.850 val loss, 4×GPU, 11% training)

[aiejvn]

First application of JEPA (Joint Embedding Predictive Architecture) to the slowrun two-hour track.

JEPA adds a self-supervised auxiliary loss alongside CE: random contiguous token spans are masked, and the model must predict their latent representations (not the tokens themselves) from context. A frozen EMA copy of the encoder provides stable targets, and VICReg regularization prevents collapse.

Key architectural components:

• JEPAPredictor: residual ReLU² MLP mapping context encodings to masked-span predictions
• sample_block_spans(): geometric-distribution non-overlapping span sampler (~4 spans, mean length 16)
• vicreg_var_loss() / vicreg_cov_loss(): variance + covariance anti-collapse penalties
• EMA momentum schedule: 0.9 → 0.999 over training
• Dropout + AdamW weight decay for regularization

Architecture: 30L / 2048d / 16h, 1.717B params.

Result

This was a resource-constrained exploratory run: 4 GPUs (standard is 8) and the large model meant only 673 of 6100 steps (11%) completed in 2 hours.

Step	Val Loss	Val BPB	Train Time
200	4.4219	1.4334	35.9 min
400	3.9972	1.2957	70.7 min
600	3.8876	1.2601	105.4 min
673 (2h cap)	3.8495	1.2478	118.1 min

The loss curve is steep and still descending at cutoff — indicating the model had not yet saturated. With 8 GPUs and ~1400 steps reachable in 2h, or a smaller model tuned for this budget, JEPA should be competitive with the 3.144 record.

Changes

New file: two_hour/train_jepa_2h.py (~965 lines), ported from the unlimited-track unlimited/train_jepa.py with a 2-hour time cap and 2hr-track data format.

Key additions vs two_hour/train.py:

• JEPA predictor + EMA target encoder in training loop
• Span masking per micro-step
• VICReg variance + covariance loss terms
• --dropout, --weight-decay (AdamW for head optimizer)
• --jepa-lambda, --jepa-ema-start, --jepa-pred-dim, etc. for JEPA hyperparams
• 2-hour time cap (stops at step boundary when 7200s elapsed)

Test plan

• Run on standard 8×H100 node to get ~2× more steps within the 2h budget
• Evaluate longer runtime (unlimited track) to see where JEPA converges

[Nils-p-o]

This seems like a promising approach, but the ~double forward pass per step seems a bit suboptimal. the simplest solution that comes to my mind is to do the encoding pass in parallel for both JEPA and LM parts by stitching them together along the batch dim, then split off the jepa states at the end of the encoder. of course, i dont know if the memory allows for that, but that may be a speed up.

second thing that comes to mind, did you test different depths, or just went with the middle, because i believe that from some research and the speedrun community generally concludes that the most abstract representation is at ~2/3 of the model depth, so maybe its better to place the jepa latents there.

maybe you will appreciate this paper, which is somewhat related: https://arxiv.org/pdf/2511.05963