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.

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.

All three share the same architecture (see below) and tokenizer (GPT-2 BPE), so any can be loaded with the same code.

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).

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.

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.