Masked language modelling (BERT-style) trains a model to reconstruct discrete tokens from masked positions: the prediction target is a one-hot token identity. I-JEPA, applied originally to vision, replaces that with a different objective: predict the representation of the masked region, not the pixel values or discrete class labels.
Text-JEPA adapts this to language. Instead of predicting which tokens were masked, the model predicts what the target encoder's representation of those tokens would look like. The model never decodes back to discrete tokens; everything happens in representation space.
Three components interact during training:
| Component | Architecture | Role |
|---|---|---|
| Context Encoder | 6-layer transformer (distilbert-base-uncased, 66M params) | Processes masked input; receives gradient updates |
| Target Encoder | EMA copy of context encoder | Produces stable prediction targets; updated via EMA only, no gradients |
| Predictor | 3-layer transformer bottleneck (768 → 384 → 768) | Maps context encoder output to predicted target representations |
The target encoder is never directly trained. Its weights are an exponential moving average of the context encoder's weights, which provides prediction targets that are stable enough to train against without gradient collapse. The predictor is intentionally shallow: it is meant to bridge context to target representation, not to reduplicate the full transformer's capacity.
The loss is cosine similarity between the predicted and target representations after L2 normalization (loss = 2 - 2 * (pred · tgt)), where 0 means perfect alignment and 2 means opposite directions. Minimising this directly trains the model to produce representations that structurally match what a more stable version of itself would output.
| Parameter | Value |
|---|---|
| Base model | distilbert-base-uncased |
| Corpus | WikiText-103 (103M tokens) |
| Training steps | 5,000 |
| Effective batch size | 256 (gradient accumulation) |
| Learning rate | 1.5e-4, cosine decay |
| Masking strategy | 4 contiguous spans per sequence, mean length 5 tokens |
Span masking rather than uniform random masking is a deliberate choice. Predicting a single missing token is close to trivial given surrounding context; predicting a contiguous span of five tokens forces the model to reason about something closer to phrase-level structure.
| Metric | DistilBERT baseline | Text-JEPA (5k steps) |
|---|---|---|
| Avg pairwise cosine similarity | 0.93 | 0.60 |
| SST-2 linear probe accuracy | ~81% | 75% |
| STS-B Pearson correlation | ~0.48 | 0.16 |
Embedding geometry improved substantially: avg cosine similarity dropped from 0.93 to 0.60, reducing anisotropy
Downstream task performance declined across both classification (SST-2) and semantic similarity (STS-B)
Between steps 2000-4000, geometry kept improving while SST-2 accuracy partially recovered; the two metrics decouple
The most interesting finding is not the performance numbers themselves, but the decoupling between them. Between steps 2000 and 4000, embedding geometry continued to improve: cosine similarity kept falling, the space became more isotropic, while SST-2 accuracy was moving in the opposite direction, partially recovering. Two quality dimensions that are often conflated were demonstrably moving independently.
BERT and its variants are famously anisotropic: all token embeddings cluster in a narrow cone of the embedding space, which limits the expressiveness of the representations and causes problems for distance-based downstream uses. The representation prediction objective does fix this: the cosine similarity drop from 0.93 to 0.60 is substantial. But the features that make a space isotropic are not the same as the features that make it useful for sentiment classification or semantic similarity.
What this suggests is that a better-structured embedding space is a necessary but not sufficient condition for better task performance. The JEPA objective improves geometry, but something is being reorganised in the process that the linear probes can't recover from at this training scale.
5,000 training steps is substantially shorter than what JEPA-style objectives typically require. Standard self-supervised pretraining uses hundreds of thousands of steps. At 5k, the model may simply be in an early reorganisation phase: the representation structure is being disrupted and rebuilt, but hasn't converged to something downstream-useful yet. Whether task utility would recover with longer training is an open question this experiment can't answer.
The initialization from pretrained distilbert-base-uncased weights also complicates attribution. Improvements to geometry could be due to the JEPA objective specifically, to general disruption of the pretrained anisotropic structure, or to some interaction between the two. A clean attribution would require training from random initialization, which is expensive at corpus scale.
Full source, training scripts, and evaluation code on GitHub →