Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture
Paper • 2301.08243 • Published • 7
AxonLM-JEPA
A 124M-parameter GPT-2-style language model exploring whether JEPA-style self-supervised pretraining (Joint-Embedding Predictive Architecture, LeCun 2022) produces more usable representations than standard next-token pretraining — tested through a neuroscience-grounded probing methodology.
This repository hosts the model checkpoints for the AxonLM-JEPA paper. Full code, training configs, and the four-probe evaluation battery are on GitHub.
TL;DR
I pretrained a GPT-2-style encoder with a JEPA objective (span masking, EMA target encoder, VICReg loss) on PubMed abstracts, then fine-tuned it with a standard language-modeling head, and compared it against an identical model trained from scratch on the same data with no JEPA stage. Both were probed for whether they encode real mouse-brain anatomical connectivity information (from the Allen Mouse Brain Connectivity Atlas) in their internal activations, and whether that information is causally used during generation.
Result: at this scale (124M params, 117M training tokens, single consumer GPU), JEPA pretraining does not outperform the scratch-trained baseline on generation quality, causal knowledge activation, or retrieval accuracy, and roughly ties on linear representation-probing quality. Full details, four independent probing methods, and a documented self-correction of an early methodological artifact are in the paper.
Checkpoints in This Repo
| File | Description |
|---|---|
jepa_pretrain.pt |
JEPA-pretrained context encoder (50k steps), before any LM fine-tuning |
jepa_lm_finetuned.pt |
JEPA encoder + fine-tuned LM head (joint fine-tuning, 20k steps) |
baseline_lm.pt |
Identical architecture, trained from scratch on the same data (no JEPA stage) |
All three share the same architecture (see below) and tokenizer (GPT-2 BPE), so any can be loaded with the same code.
Model Architecture
- 124M parameters (12 layers, 12 heads, 768 hidden dim, 256-token context)
- GPT-2 BPE tokenizer (vocab size 50,304)
- Implemented from scratch following Karpathy's build-nanogpt, with Flash Attention and
torch.compile - JEPA components (context encoder, EMA target encoder, predictor) follow I-JEPA and VICReg
Training Data
PubMed biomedical abstracts, 117M training tokens. This is a small, single-domain, single-GPU-scale corpus by design — the study is explicitly framed around what is reachable in a compute-constrained, independent-research setting, well below Chinchilla-optimal token counts for this model size (2.5B tokens would be optimal).
How to Load
import torch
from model.encoder import EncoderConfig, TransformerEncoder # from the GitHub repo
from model.lm_head import AxonLMWithHead
config = EncoderConfig(
vocab_size=50304, block_size=256,
n_layer=12, n_head=12, n_embd=768,
dropout=0.0, bias=False,
)
encoder = TransformerEncoder(config)
model = AxonLMWithHead(encoder=encoder)
state = torch.load("jepa_lm_finetuned.pt", map_location="cpu", weights_only=False)
model.load_state_dict(state["model_state_dict"])
model.eval()
The model classes (TransformerEncoder, AxonLMWithHead, etc.) are defined in the GitHub repository — this is a research artifact meant to be paired with that code, not a drop-in transformers-library model.
Intended Use
This model is a research artifact for interpretability and self-supervised learning research, not a general-purpose text generation model. It was trained on a narrow biomedical corpus at small scale, and its generation quality (perplexity ≈ 15.7–20.3 depending on variant) reflects that. Its value is in what its internal activations reveal about representation learning, not in the fluency of its outputs.
Citation
@misc{salman2026axonlmjepa,
author = {Salman, Efekan},
title = {AxonLM-JEPA: A Controlled Study of Joint-Embedding Predictive
Pretraining for Sleeping Knowledge in Language Models},
year = {2026},
howpublished = {\url{https://github.com/EfekanSalman/axonlm-jepa}}
}