A Small JEPA Word Embedding Model

After my prior blog post about SIGReg, I figured I’d train a small Joint-Embedding Predictive Architecture (JEPA) model to demonstrate it.

The paper “LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels” (Maes et al. 2026) suggested significantly reducing the complexity of JEPA models by removing stop-gradients, and the exponential-moving-average encoder. This was in the context of world models and planning.

In this case, I’m applying JEPA to the task of creating word embeddings. Prior methodologies include Word2vec (Mikolov et al. 2013) which uses a log-linear model and negative sampling, MLP next-word prediction (Bengio et al. 2000), applying CCA to small context windows (Dhillon et al. 2011), and training an autoencoder on small context windows (Shao et al. 2025).

We’ll be training a linear JEPA model with SIGReg on a small shakespeare dataset to show that it learns some informative embeddings. In other words, we’ll train an encoder that turns two words into word embeddings, and train a linear predictor that predicts the second word embedding from the first. It would be easy to extend this methodology to non-linear encoders and use larger contexts than single words.

Overview

First we’ll take our dataset and convert it into tokens. For example,

Then we create the dataset where we have context/target pairs. So in this case that would look like,

Then we create the JEPA model which uses the same embedding for the context and target , and then has predictor predict the target from the context. More formally,

Code

Preparing The Data

The dataset is a text file containing some of Shakespeare’s plays.

Code

To turn this into our dataset, we’ll convert everything to lower-case, split out punctuation, and then split on spaces to get our tokens. We’ll treat everything with frequency below 5 as an unknown token. The dataset consists of a single context word and the target word is simply the next word.

Code

SIGReg

We use the same SIGReg code as in my prior blog post. This is what makes the embedding space a bit Gaussian, avoiding dimensional collapse, as you’ll see later in Figure 1 which plots the first two embedding dimensions against each other.

Code

Making and Training The Embedding Model

Now we’re ready to train a model. The encoder in this case is just an Embedding module and the predictor is just a Linear module. We use MSE loss comparing the predicted next word embedding versus the actual next word embedding as the objective function.

Code

Visualize

Code

The visualization above shows only the first two dimensions of the embedding space. We can see several clusters including character names, royal titles, and tokens that follow apostrophes in words like ne’er, ’tis and o’er. This shows that the JEPA model is learning informative embeddings.

Further Embedding Investigation

We can also observe what words are closest in embedding space.

Code

It’s encouraging that “king” is near other professions, young is near other adjectives (including its opposite - old), and “romeo” is near other names.

Future Directions

As I mentioned earlier, it would be easy to extend this methodology to non-linear (e.g. MLP, CNN, RNN, Transformer) models and use larger contexts and targets than single words. It’s also possible to play around with other hyperparameters like SIGReg regularizer coefficient, embedding dimension, and hidden dimension and try larger datasets.

Compared to many prior methods for getting word embeddings, this does appear to be less complicated than things like negative sampling (e.g. word2vec).

There’s also recent work on approximations to SIGReg that are likely more computationally efficient with very little downside (Akbar 2026).

References

Akbar, Habibullah. 2026. Weak-SIGReg: Covariance Regularization for Stable Deep Learning. arxiv.org/abs/2603.05924.

Bengio, Yoshua, Réjean Ducharme, and Pascal Vincent. 2000. “A Neural Probabilistic Language Model.” In Advances in Neural Information Processing Systems, edited by T. Leen, T. Dietterich, and V. Tresp, vol. 13. MIT Press. proceedings.neurips.cc/paper_files/pape...206c2a01bf572b5940d7d9a8fa4c-Paper.pdf.

Dhillon, Paramveer, Dean P Foster, and Lyle Ungar. 2011. “Multi-View Learning of Word Embeddings via CCA.” In Advances in Neural Information Processing Systems, edited by J. Shawe-Taylor, R. Zemel, P. Bartlett, F. Pereira, and K. Weinberger, vol. 24. Curran Associates, Inc. proceedings.neurips.cc/paper_files/pape...761a28b734fe93831e3fb400ce87-Paper.pdf.

Maes, Lucas, Quentin Le Lidec, Damien Scieur, Yann LeCun, and Randall Balestriero. 2026. LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels. arxiv.org/abs/2603.19312.

Mikolov, Tomas, Kai Chen, Greg Corrado, and Jeffrey Dean. 2013. Efficient Estimation of Word Representations in Vector Space. arxiv.org/abs/1301.3781.

Shao, Chenze, Darren Li, Fandong Meng, and Jie Zhou. 2025. Continuous Autoregressive Language Models. arxiv.org/abs/2510.27688.