I spent a while building a JEPA-style system on top of a truncated Qwen3 backbone, and I want to walk through what it actually is, not as a hype pitch, but as an explanation of the mechanism, what it's good at, and where it genuinely struggles. If you already know what next-token prediction is but have never seen a joint-embedding objective up close, this is written for you.

The idea in one paragraph

A transformer trained the usual way learns by predicting the next token: given everything so far, guess the exact next word. JEPA (Joint Embedding Predictive Architecture, a framing Yann LeCun has written about) learns something different: given part of an input, predict the representation of the missing part, not the missing part itself. No vocabulary, no softmax over tens of thousands of token options, no cross-entropy loss. Just two numeric vectors that need to end up close to each other.

That sounds like a small change. It isn't, and the rest of this post is about why.

Why predicting tokens and predicting meaning aren't the same thing

Here's the distinction that actually matters. If a model is trained to predict "bank" after "river," it can get very good at that by learning statistical patterns in text, which word tends to follow which, without ever needing anything you'd call an understanding of what a riverbank is. Token prediction is a proxy task. It correlates with understanding, but it isn't the same thing as understanding, and a model can get a long way on the proxy without the real thing.

JEPA sidesteps the proxy. Instead of asking "what's the exact next token," it asks "what does the representation of the missing piece look like," and lets the model find whatever internal structure makes that predictable. The hope is that structure ends up looking more like genuine semantic understanding, because there's no shortcut through surface-level word statistics to get there.

Two ways to learn from text

A transformer's loss function is cross-entropy over the vocabulary: did you guess the exact right word, out of tens of thousands of options. JEPA's loss is a distance in embedding space, usually mean-squared error or cosine distance between two vectors. That's the whole difference, mechanically. Everything else follows from it.

The architecture, piece by piece

A JEPA system needs three components, and understanding why each one exists is more useful than just naming them.

Pipeline

The context encoder sees the visible part of the input (the part that wasn't masked out) and turns it into a representation. In my setup this is a truncated Qwen3 backbone, using only the early layers rather than the full stack, because the full model's VRAM footprint doesn't fit a single consumer GPU for this kind of training.

The predictor is a small network that takes the context encoder's output and tries to guess what the representation of the masked part should look like. It's deliberately lightweight, a couple of linear layers with a nonlinearity between them, because its job isn't to re-derive meaning from scratch, just to extrapolate from what the context encoder already extracted.

The target encoder produces the actual representation of the masked part, the thing the predictor is trying to match. This is the piece that makes JEPA work at all, and it deserves its own section.

The target encoder trick: why it doesn't just collapse

Here's a problem you'd hit immediately if you tried the naive version of this idea. If you train a network to predict "a representation that matches some other network's output," and both networks are learning at the same time, the easiest way to drive the loss to zero is for both networks to collapse to outputting the same constant vector for every input. Loss goes to zero, and the model has learned nothing. This is a real, well-documented failure mode in joint-embedding self-supervised learning, not a hypothetical edge case.

The fix is the target encoder. It's a copy of the context encoder's weights, updated not by gradient descent but by an exponential moving average (EMA) of the context encoder's weights over time. Critically, no gradients ever flow into it directly: it's detached from the loss, sometimes called a stop-gradient.

What flows and what doesn't

This matters because the target now moves slowly and smoothly, shaped by training but never directly chasing the loss itself. The predictor has something real and slowly-evolving to aim at, instead of a target that can cheat its way to triviality. Removing the stop-gradient is not a minor tweak, it's the piece that prevents the whole system from collapsing.

Watching for collapse, because it's a real risk, not a footnote

Collapse doesn't always look dramatic. Loss can go down steadily while the model is quietly mapping every input to nearly the same point in representation space, and if you're only watching the loss curve, you won't notice until downstream performance is clearly broken.

The way to catch it is to track effective rank, roughly, how many independent directions the representations are actually using. If every input collapses toward the same vector, effective rank drops toward 1. If the model is using its full representational capacity, rank climbs toward the dimensionality of the embedding space. Checking this at multiple levels, token-level, pooled-sentence level, and at both the target and predictor outputs, gives you an early warning that loss alone won't.

I have effective-rank numbers from an earlier draft of this project, but I can't confirm they come from the final configuration, so I'm leaving them out rather than presenting them as fact.

The training loop

Training loop

Each step looks like this: take a batch of text, mask out some contiguous spans, run the visible tokens through the context encoder, run that output through the predictor, and compute the loss against what the target encoder produces for the actual (unmasked) input. Gradients flow backward through the predictor and the context encoder. The target encoder gets updated afterward, not through backprop, but by nudging its weights toward the context encoder's current weights via the EMA rule. Repeat for however many steps you've budgeted, checkpointing along the way.

The EMA momentum (how slowly the target moves) is a real hyperparameter with real consequences. Too low, and the target chases the context encoder too closely, reintroducing the collapse risk. Too high, and the target barely moves, so the predictor gets almost no useful signal to learn from. The values I've seen work in this range sit close to the target encoder updating very slowly relative to the context encoder, in the neighborhood of 0.996 to 0.999, but I'd treat that as a starting point to tune, not a universal constant.

Two tracks, one backbone

This repo actually contains two separate systems that happen to share the same base model.

Repo tracks

Track 1 is a decision model: given a state description, a question, and a list of candidate answers, it returns a calibrated probability for each candidate in a single forward pass, without generating any text. That's the system I wrote about in an earlier post, including the honest numbers on where it struggles (many-candidate questions, in particular).

Track 2 is the JEPA pretraining experiment described in this post, a representation-learning method, not a decision-making system. It's evaluated differently too, with linear probes and k-nearest-neighbor classification on the learned representations, rather than decision accuracy.

They share a backbone and some infrastructure, but they're answering different questions: Track 1 asks "can this model make good decisions," Track 2 asks "can this training method learn useful representations without labels."

Design choices and the actual tradeoffs behind them

Truncating the backbone. Using only the early layers of Qwen3 instead of the full 28 keeps the VRAM budget inside what a single consumer GPU can handle. The tradeoff is real: deeper layers typically encode more abstract semantic structure, and truncating trades some of that away for the ability to train at all on accessible hardware.

LoRA instead of full fine-tuning. Low-Rank Adaptation injects small trainable matrices into the attention layers instead of updating every parameter in the model. This is the difference between training a small fraction of the parameters versus all of them, and on an 8GB GPU, it's often the difference between a training run that fits in memory and one that doesn't.

Sequence packing instead of padding. Padding every short sequence out to a fixed length wastes a lot of compute on padding tokens that carry no signal. Packing concatenates multiple shorter texts together and splits at exact boundaries instead, which keeps more of each batch doing useful work, with a block-causal attention mask ensuring sequences packed together don't leak information across their boundary.

Where this genuinely breaks or gets uncertain

Multi-candidate scaling, on the decision side. As I covered in the earlier post, accuracy drops as the number of candidates grows, down into the high 30s percent by 8 candidates. That's a real limitation of the current scoring approach, not a minor footnote.

Order invariance is precise but not absolute. The scoring head is designed so that permuting candidates doesn't change their individual scores, and at the head level this holds extremely tightly. But once you re-tokenize a request with candidates in a genuinely different order, the backbone's attention computation involves floating-point operations in a different order too, and that introduces a small amount of drift end-to-end. It's small enough not to change which answer gets selected in testing, but it's not a mathematical guarantee, and I'd rather say that plainly than round it up to "perfectly invariant."

Representation collapse is a standing risk, not a solved problem. The monitoring (effective rank tracking) catches it, and the stop-gradient plus EMA target architecturally resists it, but this is a known, documented failure mode across this entire family of methods, not something unique to this implementation that's been fully eliminated.

A sketch of the training step, for intuition

This isn't a verbatim excerpt from the actual codebase, just a simplified sketch to make the mechanism concrete:

# illustrative only, not the real implementation context_repr = context_encoder(masked_input) predicted_target = predictor(context_repr) with torch.no_grad(): target_repr = target_encoder(full_input) # stop-gradient loss = mse_loss(predicted_target, target_repr) loss.backward() optimizer.step() # target encoder updates by EMA, not backprop for p_target, p_context in zip(target_encoder.parameters(), context_encoder.parameters()): p_target.data.mul_(momentum).add_(p_context.data, alpha=1 - momentum) 

No vocabulary, no softmax, no cross-entropy anywhere in this loop. That absence is the entire point.

Try the decision model

The decision model (Track 1) is the part of this repo that's furthest along, though I want to be precise about its current state rather than call it finished: security hardening (auth, rate limiting, safetensors-only checkpoints, pinned model revisions) is implemented and locally tested as of this writing, but hasn't yet been confirmed by CI on a clean push, and broader reliability work (multi-seed reruns, load testing) is still open.

git clone https://github.com/bharqav/mini-jepa.git cd mini-jepa make quickstart curl -X POST http://127.0.0.1:8000/v1/decide \ -H "Content-Type: application/json" \ -H "Authorization: Bearer <your JEV_API_TOKEN>" \ -d '{ "state": "Production is down.", "questions": [{ "id": "q1", "type": "choice", "question": "Severity?", "candidates": [ {"id": "p1", "text": "P1"}, {"id": "p2", "text": "P2"}, {"id": "p3", "text": "P3"} ] }] }' 

If you're running locally without setting JEV_API_TOKEN, you'll need JEV_DEV_MODE=1 instead, and the service will log a warning that auth is off. It won't serve requests unauthenticated otherwise.

What I'd actually take away from this

JEPA isn't a drop-in replacement for transformers, and I don't think it's useful to frame it as a horse race. It's a different bet about what the training signal should be: instead of getting good at guessing exact words, get good at building a representation space where prediction in that space is possible. Whether that bet pays off at larger scale is genuinely an open question, not something a single small-scale repo on an 8GB laptop GPU settles either way. What this project does show is the mechanism working end to end, with the collapse risk real and monitored rather than hand-waved away, and with the limitations (order-invariance precision, multi-candidate scaling) reported honestly instead of smoothed over.

The code, the full architecture, and what's actually verified so far are on GitHub:

GitHub logo bharqav / mini-jepa

Order-invariant decision models and a JEPA-style latent predictor on a truncated Qwen3 backbone

mini-jepa

Order-invariant decision models and a JEPA-style latent predictor on a truncated Qwen3 backbone

Zero-token inference · Near-exact candidate-order invariance · Calibrated probabilities · EMA-target self-supervision

CI CodeQL Python 3.11+ License: MIT Code style: ruff

mini-jepa turns a pretrained causal language model into something that answers questions instead of completing text. Give it a state, a question and a set of candidate answers; it returns a calibrated probability distribution, in a single forward pass, generating zero tokens.

It also contains a JEPA-style latent predictor that learns representations by predicting masked spans in a target encoder's embedding space, rather than by predicting tokens.

Naming: the distribution is mini-jepa; the importable package is jev_qwen (from jev_qwen.infer import DecisionEngine).

Why build this? Autoregressive generation is a poor fit for structured decision problems: you pay for tokens you did not ask for, you get invalid output that must be parsed and validated, and you cannot permute the options…