axonlm-jepa

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

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}}
}

Links

• 📄 Full paper (PDF)
• 💻 Code + training configs + probing suite
• 🧠 Ground truth: Allen Mouse Brain Connectivity Atlas