1. Introduction

2. The semantics of language generation

3. The two systems, and a third

4. A generalized, constrained decision model

5. The Ideal Case

6. What is Jev?

7. Methodology

   1. A problem, and an open-source solution to it

   2. Architecture: the encoder was always the right tool

   3. Data: outcomes, not labels

   4. Training: what the objective actually asks for

   5. This is a calibration metric used as a loss

   6. RLCD: Reinforcement Learning for Calibrated Decisions

   7. RLCD against RLHF and RLVR

8. Hands-on Jev

9. Results

10. Thoughts

11. Conclusion

1. Introduction

Using a modern LLM to do a predictive task is harder than it should be, and the difficulty is structural rather than incidental. You want a label, a score, or a decision your code can branch on, and what you have is a system engineered to produce text for a human to read, so you coerce it; you ask for JSON in the prompt, constrain the decoding grammar, parse the result, and add a retry for when it comes back malformed anyway.

Three problems follow:

• The estimates are uncertain in a way you cannot inspect, because a confidence number produced by a text generator was never trained to correspond to how often the model is actually right.

• Hallucination in a decision setting means a label outside your schema or a confident answer the context does not support.

• And cost is architectural, since autoregressive models emit one token at a time, each conditioned on the last, so a decision that is logically a single choice gets paid for as a sequence and the latency cannot be bought away with more hardware.

What makes this worth dwelling on is that predictive tasks are not a niche corner of what we do. Routing a ticket, scoring a lead, flagging a transaction, judging whether a tool call is safe to run; the overwhelming majority of AI inside real software is making a judgment rather than writing something, and in almost every case we are using generation as a proxy for prediction. It works, which is why nobody has pushed on it, but proxies have a way of becoming invisible until somebody removes one.

That is what TypeSafe did in September with Jev, a model that does not generate text at all; you hand it context and typed questions defined in your code, and it returns choices, scores and calibrated probabilities directly, with nothing to parse. They call it the first System One model, after Kahneman’s fast intuitive mode of thought as against the slow deliberate System Two that chat models are built for. One caveat before we start, because it shapes what this article can be: there is no paper. TypeSafe spent two years in stealth and announced a new architecture, a new sampler and a new training algorithm while publishing essentially none of it, so instead of deriving somebody’s method from their description, we will ask what a model of this kind would have to look like built from first principles, and then hold that against what has actually been published.

(Usual caveat: the information in this article is upper bounded by my own grasp of the topic, my patience for typing for hours, and a writing medium that still refuses to support inline LaTeX. Please re-visit everything you read here, and since there is no paper to point you to this time, go read the docs and form your own view.)

2. The semantics of language generation

Before we can say what a decision model should look like, we need to be precise about what we are asking a language model to do when we use one for prediction, and the cleanest way in is through an analogy most readers will already have strong feelings about.

2.1 Strongly typed, loosely typed

In a strongly typed language, you declare what a function returns and the compiler checks it before anything runs; a mismatch is caught at build time, and by the time the program executes, an entire class of error has been eliminated. In a loosely typed one, the same mismatch is discovered at runtime, which is to say in production, by a user. The argument for the former has never really been about aesthetics, it is that the contract is verified before the cost of being wrong is incurred.

Now look at how we currently get structured output from an LLM. You describe the shape you want in a prompt, the model generates a string, you parse it, and you discover whether the contract held. That is runtime checking in the worst sense, because the thing being checked is not even deterministic; the same prompt can honour the schema on one call and invent a category on the next. Structured output and constrained decoding tighten this considerably, but they tighten it by validating the string as it is produced rather than by removing the string from the problem.

The same three-way choice, declared two ways. In the top row the contract is a sentence and the check happens after the fact; in the bottom row the contract is a type and the invalid answer has nowhere to come from.

The thing worth noticing is that a typed question defined in your code is checkable in a way a prompt never is. If the answer space is an enum with three members, the return type is one of three members, and that is true before the model runs, not after it has spoken. That is the actual difference between the two paradigms, and the rest of this section is about why it costs so much to get there through generation.

2.2 How much of a sentence is actually information?

Take a sentence that encodes a decision, something like “Based on the transaction history, this appears to be a duplicate charge, and therefore the refund policy supports approving it.” How much of that is the decision?

The decision is roughly two bits. Duplicate or not, approve or not. Everything else is scaffolding; based on, appears to be, and therefore, supports. Those words are doing real grammatical work, they make the sentence a sentence, but they carry almost no information about the answer, and crucially they are highly predictable given what came before. Once a model has committed to “and there-”, the continuation “-fore” is nearly certain.

Illustrative per-token surprisal for the sentence above. The information is concentrated in two positions; every other forward pass is spent emitting grammar.

We have met this idea before, in the Byte Latent Transformers article, where the whole patching mechanism was built on exactly this observation; sequences of low-entropy bytes get grouped together because they are cheap and predictable, while high-entropy points are where the model should spend its attention. BLT used entropy to allocate compute within the input. The same lens applied to the output says something uncomfortable.

If we write the information content of a generated sequence as the sum of per-token surprisals,

then for a decision expressed in prose, the overwhelming majority of that sum comes from a handful of positions and the rest contributes almost nothing. But an autoregressive model does not know that in advance, and more importantly it cannot act on it even if it did, because every token requires a forward pass regardless of how predictable it was. You pay uniformly for a signal that is distributed extremely unevenly. A decision worth two bits costs you forty forward passes, thirty-eight of which are spent emitting grammar.

And those passes are sequential, which is the part that no amount of hardware fixes. Each token is conditioned on the last, so the latency floor is the length of the output, not the size of your cluster.

2.3 Generation as constrained optimization

Let us make this precise, because the optimization framing is where the difference becomes structural rather than merely wasteful.

Unconstrained generation searches over all sequences the vocabulary can produce:

where V* is every finite string over the vocabulary, an unimaginably large space of which the overwhelming majority is nonsense. The model’s job is to concentrate probability mass on the thin sliver that is fluent and relevant, and it does this remarkably well, which is the entire achievement of language modelling.

Structured output narrows the feasible set to the strings that satisfy a grammar:

This is a genuine improvement and it is why constrained decoding works. But read the mechanism carefully; at each step you compute a distribution over the full vocabulary and then mask out the tokens that would violate the grammar. The search space has been restricted, the computation has not. You are still taking |y| sequential passes, still evaluating all of V at each one, and still doing it to arrive at an answer that was, all along, one of three options.

Now write what we actually wanted:

Three nested feasible sets. Constrained decoding moves you from the outer box to the middle one; a decision model operates natively in the inner one.

The feasible set is no longer a subset of string space, it is a finite answer set defined in your code; A might have two members or two hundred, and Jev caps it at 255. And this changes the problem qualitatively rather than quantitatively, because when the answer set is small and enumerable you do not need to search it one token at a time. You can evaluate the whole distribution over A in a single pass, which is precisely what removes both the sequential latency and the uniform cost on uninformative tokens.

That is the gap this article is about. Constrained decoding projects a generation problem onto the feasible set after the fact. What we are describing optimizes over the feasible set natively, and a model built that way is a different object from a language model wearing a grammar.

But having established the shape of the difference, we are left with two questions that the mathematics does not answer. The first is practical: which real problems actually have this shape, because |A| being small and known is a property of the task rather than something you can impose on an arbitrary one, and it would be useful to have a way of recognizing those tasks on sight. The second is about trust: if the model returns a distribution over A rather than a sentence, we are going to make decisions based on those numbers, and we need to know what the numbers mean before we build anything on top of them. The next section gives us vocabulary for both.

3. The two systems, and a third

Two things need naming before we can look at the model itself, and they answer the two questions we just left open. The first is a frame for recognizing which problems belong in this category at all, borrowed from psychology and used here somewhat loosely. The second is the property that decides whether the output can be trusted, which is considerably more technical and which I think is the single most important idea in this article.

3.1 System One and System Two

The name comes from Kahneman, who split human cognition into two modes in Thinking, Fast and Slow.

System 1 is fast, automatic and effortless; it is what recognizes a face, reads a word without deciding to, senses that a room feels tense, or knows a sentence is ungrammatical before you could say why. System 2 is slow, deliberate and expensive; it is what multiplies seventeen by twenty-four, compares two mortgage offers, or follows an argument through four steps to check whether it holds.

The useful part of the distinction is not speed, though, it is structure, and this is where the section connects back to what we built in Section 2. System 1 is not a faster System 2 working on the same problem; it is a different operation on a differently shaped problem, a match against an answer space that is already known rather than a search through a space that has to be constructed. Which is exactly the distinction we drew in optimization terms, because searching over V* for a fluent sequence and selecting from a known A are not the same problem at different speeds, they are different problems. A System 1 task is one whose feasible set is small, enumerable and fixed in advance. That is the whole of it, and it is why the two framings are really one framing; Kahneman’s vocabulary describes in cognitive terms what the nested-sets diagram describes in mathematical ones.

The same distinction twice. Kahneman’s vocabulary describes in cognitive terms what the nested-sets diagram describes in mathematical ones.

Today’s chat models are System 2 machines by design and correctly so. They are built to work through novel problems, construct chains of reasoning, and produce answers nobody enumerated beforehand, with a human reading the output. That is a genuine achievement and the right tool for an enormous number of tasks.

But it is the wrong tool for ticket routing, and now we can say why precisely rather than observing that it feels slow. Ticket routing is a System 1 problem being solved by a System 2 machine. The answer space is three teams, known in advance, enumerable; there is nothing to search and nothing to construct, and what is needed is a match against candidates we already hold. Instead we deploy a machine that builds a sentence one token at a time in order to tell us which of the three it picked.

I should flag that the metaphor is marketing as much as it is science. Kahneman’s dual-process account is a description of human cognition still under dispute among psychologists rather than an architectural specification, and nothing about Jev literally implements System 1. What the name usefully captures is the shape of the problem being targeted; bounded, enumerable, pattern-matched, answers known ahead of time. Read as a claim about the task rather than the mechanism, it holds up well.

3.2 A hypothesis: System Three

Here is where I want to speculate slightly, because the two-system framing has an obvious hole once you start building with it.

If some of your decisions are System 1 and some are System 2, something has to decide which is which. That dispatcher is not doing either kind of work; it is not pattern-matching an answer and it is not reasoning through a novel problem, it is deciding what kind of problem it is looking at and routing accordingly. Call it System Three, and note that this is not quite as invented as it sounds, since Stanovich’s tripartite model of cognition already splits Kahneman’s System 2 into an algorithmic mind that executes and a reflective mind that decides what is worth executing.

The dispatcher nobody names. It is not doing either kind of work; it is deciding which kind of work the request needs, and that judgment is itself bounded and enumerable.

The reason this matters practically is that every agent harness already has a System Three and nobody calls it that. Something decides whether a user request needs the fast model or the expensive one. Something decides whether a tool call is risky enough to interrupt and ask about. Something decides whether the current answer is good enough to return or needs another loop. Those are all metacognitive decisions, made constantly, and until recently they were made either by hand-written heuristics or by burning a full reasoning-model call to ask.

And the interesting thing, which we will see concretely in the hands-on section, is that a System One model turns out to be a very good System Three. Deciding which model should handle this request is itself a bounded, enumerable, pattern-matched judgment; the answer space is your list of models, the state is the request, and the decision wants to be fast because it sits on the critical path of every single call. So you end up with a System One model whose job is to decide when System Two is needed, which is a slightly recursive arrangement but a sensible one.

I offer this as a frame rather than a claim. TypeSafe does not use the term and nothing in their material depends on it. But I find it clarifies why this category of model is interesting beyond cost; it is not only that some decisions get cheaper, it is that the orchestration of expensive reasoning becomes something you can afford to do well.

3.3 Calibration, which is the load-bearing word

Now the property everything rests on, and the one most likely to be skimmed because it sounds like a technicality.

Any classifier can emit a number between zero and one; a softmax guarantees they sum to one and tells you nothing whatsoever about whether they mean anything. Calibration is the much stronger property that the numbers correspond to observed frequencies, so that among all the cases where the model says 0.8 it is right about 80% of the time. If it is right 95% of the time the model is underconfident, and if it is right 60% of the time it is overconfident, and in both cases the number is lying to you even though the ranking may be perfectly good.

We went through this properly in the uncertainty and calibration primer, including why modern networks tend to be badly overconfident and what a reliability diagram actually shows, so I will not rebuild the machinery here.

What I want to add is why it matters more in this setting than in the ones we discussed there. In a conventional classifier the probability is usually decoration; you take the argmax and move on, and whether 0.8 really meant 80% rarely changes what the system does. In a decision model the probability is the control surface. You are going to write code that acts automatically above 0.9, escalates to a human below 0.6, and calls a reasoning model in between, and every one of those thresholds is meaningless unless the number means what it says. An uncalibrated model with excellent accuracy will route confidently wrong cases straight past your human review, because it was confidently wrong about precisely the cases the threshold existed to catch.

A reliability diagram. The shaded area is the distance between what the model claims and what it delivers, and the threshold is where your code stops asking anyone.

This is also why TypeSafe’s training algorithm is called Reinforcement Learning for Calibrated Decisions rather than something about accuracy. The claim is not merely that the answers are good, it is that the uncertainty attached to them is honest, and those are separate properties that can be optimized separately.

One caveat the docs are admirably direct about, and which is easy to misread. Calibration is measured across groups of predictions, never within one. A well-calibrated 0.8 tells you that cases like this one are right about eighty percent of the time; it does not tell you that this particular answer is eighty percent correct, because this particular answer is either right or wrong and there is no eighty about it. The distinction sounds pedantic until somebody builds a system that treats a single confidence score as a per-case guarantee, which is exactly the mistake the property invites.

So we now have both pieces. A System One task is one whose answer space is small, enumerable and fixed in advance, which is Section 2’s feasible set stated in cognitive rather than mathematical terms. And calibration is what makes the returned distribution usable as a control surface rather than decoration, which is what we need if code is going to act on it without a human reading the output first.

Which raises the obvious question, and it is the one the next section is about. If a model is going to serve every task of this shape rather than being trained for one of them, what exactly is that model? It cannot be a classifier in the usual sense, because a classifier’s labels are baked in at training time and we have just said the answer space arrives at call time. So what is it?

4. A generalized, constrained decision model

4.1 The generalized classifier

The answer begins with an observation that is obvious in hindsight and that I had never quite put together until I started reading about this.

Imagine you have trained three models; one that detects spam, one that flags fraudulent transactions, and one that tags named entities in text. These are different tasks in different domains with different label sets, and in the normal course of things they are three codebases, three training runs, three deployments and three sets of eval infrastructure.

Now ask what they actually have in common, because the answer is not what you might first assume. It is not the architecture; one might be a fine-tuned BERT, another gradient-boosted trees over engineered features, another a sequence tagger with a CRF on top. It is not the data either, in the sense of subject matter, since an email and a card transaction and a sentence of news copy share nothing at all.

What they share is the shape of the problem. Each takes some context and returns a member of a known, finite set. Spam or not. Fraudulent or not. Person, organization, location, or none of them. The domains differ wildly and the structure is identical, and the thing we have been doing for fifteen years is building a new model for each domain while rebuilding that identical structure every time.

Three unrelated domains, three unrelated architectures, and an identical problem shape underneath all of them.

Which suggests an inversion. If the structure is what is constant, then instead of fixing the data and varying the model, fix the model and vary the structure. One model that understands natural language well enough to make a judgment about anything, into which you inject the answer space at call time rather than at training time. The spam classifier becomes a question with two options, the fraud detector becomes a question with two options and some transaction records as context, the entity tagger becomes a question with four options asked once per span.

The inversion. Rather than training a model per label set, the label set arrives with the request.

That is what Jev is, and it is worth saying plainly before we complicate it: a generalized classifier whose label set you supply at inference time. Everything else in this section is about why that is harder than it sounds and what you get when it works.

4.2 Why not just use an LLM for this?

The obvious objection is that we already have models which understand language well enough to make these judgments, and we have been using them for exactly this. So what is actually wrong?

It helps to think about it in terms of entropy, since that is the language we used in Section 2 and it makes the mismatch precise.

An autoregressive language model places a distribution over sequences, and the entropy of that distribution is enormous. At each step it is choosing among a vocabulary of perhaps a hundred thousand tokens, and over a sequence of length n the space has size |V|ⁿ. The model’s entire job is to concentrate probability onto the vanishingly small fraction of that space which is fluent and relevant, and it is extraordinarily good at this, which is the whole achievement of language modelling.

A decision has an entropy of at most log|A|. For a three-way routing choice that is about 1.6 bits. For a yes-or-no question it is one bit.

H(decision) ≤ log₂ |A| versus H(sequence) ~ n · log₂ |V|

So when we use an LLM for a decision, we take a machine built to navigate an astronomically large space and point it at a problem whose answer space has three members. The model does not know the space is small. It is still computing a full distribution over its vocabulary at every step, still spending its capacity on the question of what word plausibly comes next, and the fact that only 1.6 bits of the output actually matter is information available to us and not to it.

There is a second mismatch that is subtler and, I think, more damaging. The probability of a token is not the probability of a decision. When a model emits the token “billing” with probability 0.87, that number is the likelihood of that token appearing at that position in text, conditioned on everything before it. It is not the model’s belief that billing is the right team. Those two quantities are correlated, often quite strongly, and they are not the same thing, and nothing in next-token-prediction training ever asks them to be. This is why extracting confidence from an LLM’s logits gives you something that behaves roughly like confidence while failing exactly the calibration test from Section 3.

So the answer to “why not just use an LLM” is that you can, and people do, and it works adequately. But you are paying for an enormous search you do not need, and the uncertainty you get back is a byproduct of a different objective rather than an estimate of what you actually asked.

4.3 Where the tokens go

The cost consequence is worth seeing laid out, because it is not uniform across the three parts of an LLM call.

A chat-model request has an input, possibly some reasoning, and an output. The input is usually the largest in raw token count, but it is also the cheapest per token, it is processed in parallel rather than sequentially, and it is frequently cached. Reasoning tokens and output tokens are where the cost and latency actually live, and they are priced accordingly; on current frontier models output tokens cost several times what input tokens do, precisely because each one requires its own forward pass that cannot be batched with the others.

Now look at what happens to that profile when the output is a choice rather than a sentence.

Illustrative rather than measured. The input does not shrink and should not, since the context is what grounds the decision; what collapses is the expensive end.

The input does not shrink, and should not, because the context is what grounds the decision and you want to give it everything relevant. Reasoning may or may not be present depending on whether the task needs it. But the output collapses from a sentence to a value, and since the output is the expensive end, the economics change shape rather than merely scaling down.

I want to be careful here, because there is a sloppy version of this claim going around that says output cost goes to zero. It does not. A choice still has to be produced, there is still computation behind it, and TypeSafe still charges for tokens. What is true is that the output becomes insignificant relative to the input, which is a different and more defensible statement. We will return to this in the Thoughts, because the gap between “free” and “negligible” matters more than it sounds.

4.4 What this unlocks, practically

Abstractly, cheap calibrated decisions sound like an infrastructure improvement. Concretely they change what is worth building, and the clearest examples are in the places where volume has always been the obstacle.

Take customer interaction at scale. A sales team with fifty thousand inbound leads cannot afford to have a frontier model read each one and judge intent, fit and urgency, so they use keyword rules and a scoring heuristic that everybody knows is crude and nobody can justify improving because the model calls would cost more than the leads are worth. Drop the per-decision cost by two orders of magnitude and the calculation inverts; now you can ask six questions of every lead, and more importantly you can ask them again whenever the context changes rather than scoring once at capture.

The same pattern appears anywhere a judgment currently happens either rarely or crudely because it was too expensive to do properly. Content moderation on every message rather than on a sample. Support ticket triage on arrival rather than after a queue delay. Marketing segmentation recomputed per interaction rather than nightly. Routing every agent request to an appropriately sized model rather than sending everything to the largest one.

The band in the middle is the article’s practical claim: a set of judgments that were always desirable and never affordable.

And this is where the naming becomes relevant, because Jev is named after William Stanley Jevons, who observed that making a resource cheaper to use does not reduce consumption of it but increases it. If TypeSafe is right about the cost curve, the interesting consequence is not that your existing classification bill falls by 99%. It is that decisions become cheap enough to put in a hundred places where nobody would currently consider putting one, and the aggregate load goes up. That is a claim about the shape of what gets built, not about anyone’s invoice, and I will come back to it at the end.

5. The Ideal Case

That said and given our understanding of pre-requisits; imagine if we could specify the model we want, here is what we would ask for.

1. It reads like a language model. Messy human text, a ticket with typos, a log excerpt, three transaction records and a paragraph of policy, understood the way a competent colleague would understand it on first reading.

2. It answers like a function. I declare the answer space in my own code, and that declaration is a contract; if I say the answer is one of three things, a fourth has nowhere to come from. No parsing, no validation, no retry.

3. The probability means something. Not a softmax, not a confidence the model generated because I asked for one. A number I can threshold; when it says 0.9 across a thousand cases, roughly nine hundred are right.

4. Many questions, one price. Real decisions decompose into several independent judgments about the same context. Asking five should cost about what asking one costs, because nothing about the second depends on the first.

5. Latency set by the model, not the answer. An autoregressive model’s floor is how long the output is. A model evaluating a known set in one pass has a floor set by its own size. The second is low enough to sit inside a request a user is waiting on.

6. Pay for reading, not for replying. Charge me properly for the context, which is the whole point, and nominally for an answer carrying two bits. Otherwise the economics steer me towards asking fewer questions than I should.

7. Grounded in what I gave it. If the context does not support a conclusion, that should show up in the number rather than being filled in from training. A judgment drawing on knowledge I cannot see is one I cannot audit.

The tension

Read that list back and there is a contradiction sitting in it.

One requirement and six, pulling against each other. Everything in Sections 6 and 7 is about how that gets resolved.

Point 1 wants deep language understanding, and we know exactly one way to get that: pretrain something enormous on text. Points 2 through 7 want behaviour that pretraining does not give you and in places works against. The thing that understands language is an autoregressive generator, and an autoregressive generator is slow, untyped, and produces token probabilities rather than decision probabilities.

So the specification is for something that is a language model on the way in and a classifier on the way out. Not an LLM with a grammar bolted on, which keeps all the sequential cost. Not a classifier with better features, which has its labels welded in and no real understanding of what it reads.

What we are giving up

No explanation. A model returning a value has nowhere to put a reason. When an answer looks wrong you have the context, the question and the number, and nothing else.

No surprises. The answer space is fixed before the model runs. If the right answer is not in the set, you get the nearest one with whatever confidence it has, not a warning that the set was wrong.

Neither is fixed by a better model. They are the price of the shape, and the shape is what we came for.

6. What is Jev?

6.1 The answer to the tension

Section 5 ended on a contradiction: the specification needs a language model on the way in and a classifier on the way out, and those are two different machines.

Jev’s answer, as far as anyone outside TypeSafe can tell, is to keep the first and discard the second half of what usually comes with it. It understands language the way a pretrained model does and it does not generate. You send it a state, the context the judgment is about, together with one or more questions defined in your code, and it returns typed values with probabilities. No prose, no code, no rationale, nothing to parse.

Two inputs, one pass, a typed answer. Nothing in the path emits a token.

The crucial word is non-autoregressive. Jev does not emit tokens one at a time conditioned on the last, which is what collapses the latency floor from the length of the answer to the size of the model. That single architectural fact is doing most of the work behind points 2, 5 and 6 on our list.

TypeSafe calls this a System One model and Jev is the first public one. How it is built, they have not said.

6.2 Three primitives

The answer space is declared through one of three question types.

Choice returns a distribution over your options; Score returns a continuous value between ordered levels; Noul returns the probability that a statement is true.

Choice selects one option from a set, returning a probability for every option plus a confidence. Up to 255 members.

Score rates the state against ordered levels, and because the levels have a direction the model returns a continuous value between them; 1.4 on a zero-to-two frustration scale rather than a forced bucket.

Noul answers yes or no and returns the probability that the statement is true. TypeSafe’s own coinage, and the one you reach for most.

The design philosophy is decomposition. Rather than one large question requiring the model to weigh several things at once, you ask several small ones and combine them in code: was a refund requested, does the evidence suggest a duplicate charge, does the policy permit it. Judgment in the model, policy in your repository.

The model is asked what it can judge; the rules about what to do with those judgments stay in code you can read and change.

And because the questions are independent they are evaluated in parallel in a single pass, so a fourth question barely moves the latency. That is point 4 satisfied directly.

6.3 Where it meets the spec, and where it does not

Three of the seven are vendor claims rather than verified properties, and the third is the one I would most want independent evidence for, since calibration is what the whole control-surface argument rests on.

6.4 The numbers, and how to read them

TypeSafe reports 70–500 ms end to end, 40–200× faster than frontier models on this class of query, and headline figures of 193.6× faster and 444.6× cheaper. Context is 64k tokens per request, with state plus the longest question fitting inside 32k.

Treat these as vendor claims. They are comparisons the vendor chose against baselines the vendor selected, on a task shape the vendor defined. That does not make them wrong; it makes them unverified. The independent evaluations in Section 8 are more useful precisely because nobody involved had a stake in the result.

6.5 What it is not

It does not write. No prose, no code, no explanations. A complement to a generative model, not a replacement.

It is not a drop-in LLM replacement. It handles the classification work we currently route to chat models, and nothing else.

It is hosted and proprietary. Not fine-tuned on your data, not self-hostable, reachable through one endpoint. A different kind of dependency from a classifier you trained yourself, and worth noticing before cheap decisions make it tempting to put one in a hundred places.

And the method is undisclosed. Two years in stealth produced a new architecture, a new sampler and a new training algorithm, none of which has been published. The next section is about what can be inferred anyway.

7. Methodology

7.1 A problem, and an open-source solution to it

Everything so far has described what Jev does. This section is about how, and it runs into the wall I flagged in the introduction: TypeSafe has published nothing. Two years in stealth, a new architecture, a new sampler and a new training algorithm, all undisclosed.

Fortunately the category did not stay closed for long. Laya, from Convai Innovations, is an open-weight System One model released under Apache 2.0 with the same three primitives, the same single-pass evaluation, and the same training approach. It is 421M parameters, runs locally in about a gigabyte, and the weights, export scripts and training configuration are all public.

So for the rest of this section I will describe Laya and treat it as evidence about the category. The two models are not the same and I will be clear about which claims come from where, but the architectural questions, how do you get a language model to answer like a classifier and how do you train a probability to be honest, have public answers now, and those answers are almost certainly close to what TypeSafe is doing.

7.2 Architecture: the encoder was always the right tool

The answer to Section 5’s tension turns out to be slightly embarrassing in hindsight, which is that we already had the architecture and stopped using it.

Laya is a bidirectional encoder with a decision head on top. The English checkpoint is ModernBERT-large; the multilingual one is mmBERT-based, covering 100+ languages. That is the BERT lineage, not the GPT lineage, and the distinction is exactly what matters here.

An autoregressive decoder reads left to right because it has to; it is predicting the next token and must not see the future. An encoder has no such constraint. Every token attends to every other token in both directions, the whole input is processed at once, and what comes out is a representation of the entire state rather than a prediction about what follows it. For understanding a thing, that is strictly better. For generating text, it is useless.

The same sentence under two attention patterns. The triangle is why generation is sequential; the full square is why a decision is not.

And it is not only Laya. GLiNER2.5-Decide, an open 340M model from Fastino, uses DeBERTa-v3-large as its encoder and exposes the same call-time label sets in a single forward pass. Two independent teams reaching independently for the same architecture family is stronger evidence for the argument above than either one alone. The encoder was the right tool, and the only reason it stopped being obvious is that the field spent five years optimizing for generation.

Which is why the field moved on. Once generation became the goal the decoder won, and encoders became what you reached for when you had a specific classification task and some labelled data.

But we are not generating. We want a representation of the state good enough to answer questions about it, and a bidirectional encoder is the right machine for that and always was. The novelty is not the encoder; it is what sits on top.

The decision head takes the encoder’s representation together with the question’s answer space and produces a distribution over the options. One forward pass, no decode loop, and because the questions are independent given the representation, all of them are answered from the same pass. That is where the latency goes: Laya reports 32.8 ms p50 for a routed question against Jev’s 236–276 ms, and both are a different order of magnitude from a chat model’s one to three seconds.

7.3 Data: outcomes, not labels

Before the training objective, a word about what it needs to consume, because this is the part nobody discusses and it is harder than it looks.

The encoder pretraining is unremarkable and public. ModernBERT and mmBERT are trained the way encoders have always been trained, on large text corpora with masked-language-modelling objectives, and Laya inherits that lineage rather than inventing anything.

The decision stage is where it gets interesting. A proper scoring rule needs outcomes, not labels, and those are different objects.

A labelled dataset says this ticket belongs to billing. An outcome dataset says this ticket was routed to billing, and that turned out to be correct. The first is an annotator’s judgment; the second is what happened. For training a classifier the distinction barely matters, but for training a calibrated classifier it is the whole thing, because calibration is a claim about frequencies in the world and you can only learn it from the world’s frequencies.

Which creates a sourcing problem. Annotated corpora are abundant and outcome-resolved decision logs are not; they live inside companies, they are domain-specific, and they are rarely collected with this use in mind. How either lab assembled enough of them is unstated. Laya ships three checkpoints, English, multilingual, and one tuned for enterprise decisions like customer-service routing, which at least tells you the shape of what they targeted, and Convai’s evaluation uses a 2,000-decision benchmark of their own construction.

For Jev we know nothing at all. No data description, no mixture, no scale, and no indication of whether the outcomes were collected, simulated, or derived from labels by some procedure they have not described.

I flag this not as a complaint but because it bears on the zero-shot question in 7.8. If calibration is learned from outcome distributions, then how well it transfers to your distribution depends on how much your decisions resemble the ones it was trained on, and neither lab has given you the information to judge that.

7.4 Training: what the objective actually asks for

Start with what ordinary supervised learning asks for.

Train a classifier the standard way and your objective is cross-entropy against a hard label; the target is 1 for the correct class and 0 for everything else. Read that literally and notice what it is telling the model: you should have said one hundred percent. Every training example, without exception, pushes towards maximal certainty, and the only thing restraining it is that it cannot be certain about everything at once.

So the model learns to be right, and it learns that confidence is rewarded. What it never learns is when to be unsure, because nothing in the objective ever asks. There is no example whose label says “this one is genuinely a coin flip, please output 0.5”. The gradient points at 1.0 every time.

That is the root of the overconfidence we documented in the calibration primer, and it is why the usual remedy is post-hoc, temperature scaling, fitting one parameter on a validation set to squash the probabilities after training is already done. A patch applied to a model optimized for the wrong thing.

Now the alternative. A scoring rule assigns a reward to a predicted probability given the outcome that occurred. It is proper if the model maximizes expected reward by reporting its true belief, and strictly proper if that is the unique maximum:

Brier = (p − y)² Log-loss = − [ y log p + (1−y) log(1−p) ]

The difference from cross-entropy is not the formula, log-loss is cross-entropy, it is what you evaluate it against. Supervised training scores you against a hard label on one example. A scoring-rule objective scores you against observed outcomes across many predictions, which is the setting where honesty becomes optimal.

A worked example. Suppose the model genuinely believes an event has probability 0.7 and is deciding what to report. Over a hundred such cases the event happens about seventy times. Reporting 0.9 costs an expected 0.250, reporting 0.5 also costs 0.250, and reporting the honest 0.7 costs 0.210.

Overclaiming and hedging are both punished. The minimum sits exactly at the model’s true belief and nowhere else.

In general, the expected penalty for reporting q when the truth is p works out to

E[Brier] = p(1−p) + (p − q)²

the first term is irreducible, the second is the penalty for lying

a parabola in q with its floor at q = p. The first term is the world’s own uncertainty, which no report can remove; the second is what you pay for misrepresenting your belief, in either direction.

So the objective rewards honesty rather than confidence. The model cannot improve its score by sounding more certain, which is precisely the lever cross-entropy hands it.

7.5 This is a calibration metric used as a loss

Worth naming plainly, because if you have read anything on forecast evaluation you will already have recognized it.

A strictly proper scoring rule is the classical tool for measuring calibration, not a new invention. The Brier score has a well-known decomposition:

Brier = reliability − resolution + uncertainty

Reliability is exactly the gap in the reliability diagram from 3.3, the distance between what you claimed and what happened. Resolution is how much your predictions vary from the base rate, which is to say whether you are telling anyone anything. Uncertainty is a property of the data you cannot touch.

So the objective is not a novel mechanism. It takes a metric meteorologists have used for sixty years and moves it from the evaluation step to the training step, which is a real contribution but a different kind from the one the name implies.

And the decomposition surfaces a failure mode that matters. Optimizing calibration alone admits a degenerate solution: a model that ignores the input entirely and reports the base rate on every example is perfectly calibrated. If 12% of tickets are urgent, answer 0.12 always, and your reliability is flawless. It is also useless, because the resolution term is zero and you have learned nothing about any particular ticket.

A strictly proper rule does punish this, since the resolution term enters with a negative sign and the degenerate model forfeits all of it. But it means calibration and accuracy are both being optimized and they are in tension, and a model can trade some of one for the other. Which is the honest reason requirement three in the 6.3 table reads “claimed” rather than “met”.

7.6 RLCD: Reinforcement Learning for Calibrated Decisions

We now have both halves of the idea, so let us put the name to them.

Reinforcement Learning for Calibrated Decisions is TypeSafe’s training algorithm and, as far as the open-source work indicates, the same approach Laya uses. Neither lab has published the method, so what follows is the shape of it rather than the details, assembled from what both say about it and from what the mathematics requires.

The loop. Write the model as a policy π_θ that takes a state x and a question with answer space A, and emits a distribution over the options:

π_θ( · | x, A ) = ( p₁, p₂, …, p_|A| ), Σ_k p_k = 1

The world then resolves the question, giving an outcome y ∈ A. A strictly proper scoring rule S turns that resolution into a reward, and the objective is the expectation of it:

J(θ) = E over (x, A, y) of S( π_θ( · | x, A ), y )

A policy, an environment, a reward and an update. What is unusual is that the action is a probability and the reward is a scoring rule.

That is reinforcement learning in a fairly literal sense; a policy produces an output, the environment returns a reward, parameters move to increase it. What is unusual is that the action is a probability rather than a token or a move, and the reward is a scoring rule rather than a preference or a check.

Why it has to be RL rather than supervised fitting.

This is the part worth slowing down on, because at first glance you could imagine skipping the machinery; just collect outcomes, turn them into soft labels, and do ordinary supervised learning. It does not work, and the reason is instructive.

Imagine you are training a weather forecaster. On Tuesday she says there is a 70% chance of rain. It rains. Was she right?

You cannot answer that. Not “you do not have enough information yet”, the question genuinely has no answer from one day. A 70% forecast is not falsified by rain and not confirmed by it either, because the forecast was never a claim about Tuesday. It was a claim about the class of days that look like Tuesday, and the only way to check it is to collect every day she said 70% and count how many were wet.

So there is no per-example target. The supervised framing wants to write a loss against some target t(x), and the target does not exist. All you observe is y ∈ {0, 1}, and regressing probabilities onto hard outcomes is exactly the cross-entropy setup from 7.4 that drives everything towards 1.0.

The property we actually want is a statement about a group:

E[ y | p_θ(x) = q ] = q for every q in [0, 1]

Read it carefully: among all the inputs on which the model says q, the outcome rate should be q. That is a constraint on the distribution of predictions the model produces, not on any individual prediction, and it depends on the policy’s own behaviour. Change the model and you change which inputs land in the q bucket, which changes what the constraint is evaluated over.

There is a Bayesian way to see the same thing, and it makes the “one example tells you nothing” point concrete. Suppose you want to check whether the forecaster’s 70% claim is honest. Put a prior on the true rain-rate among 70%-days, say a Beta, and update it with each resolved day:

θ ~ Beta(α, β), θ | data ~ Beta( α + r, β + (n − r) )

where r is how many of those n days were wet. The posterior mean is (α + r) / (α + β + n), and the spread shrinks like 1/√n.

After one day the posterior is still almost the prior. The evidence about a probability accumulates across cases and cannot be extracted from one.

After one rainy day the posterior has barely moved; it is consistent with the true rate being anywhere from 0.3 to 0.95. After a few hundred days it has collapsed onto a narrow interval, and now you can say whether 0.7 was honest. And notice that this is what the model is implicitly being asked to do. A calibrated predictor is one whose outputs would survive exactly this test, bucket by bucket, which is why the training signal has to be evaluated over a population of its own predictions rather than per example.

That is precisely the structure reinforcement learning exists to handle: an objective defined over the distribution your policy induces rather than over fixed targets per input. You cannot get here by relabelling your training set, because the thing you would need to label is a property the model has not yet exhibited.

Another way to see it. Supervised learning asks “what is the answer to this?” and marks you against a key. RLCD asks “how often are you right when you talk like that?” and marks you against your own track record. The second question is incoherent for one example and perfectly well-posed for ten thousand.

What the three letters are doing. Reinforcement, because the signal comes from consequences rather than a supplied answer. Calibrated, because the reward is chosen to make honesty optimal rather than accuracy. Decisions, because the action is a typed choice over a known answer space rather than free-form output.

And the thing to keep hold of. Everything in 7.4 and 7.5 was reward design, not algorithm. RLCD’s contribution is the choice of what to reward, and that choice is a scoring rule from the 1950s. The RL wrapper is how you optimize something defined over a distribution of predictions; the proper scoring rule is why what you converge to is honest. Two separable ideas, and the second is doing the work.

7.7 RLCD against RLHF and RLVR

With RLCD described, it is worth placing it against the two reinforcement-learning approaches you will already know. All three are distinguished by where the reward comes from.

RLHF gets its signal from human preference. A person compares two outputs, a reward model learns to imitate that judgment, and the policy optimizes against the reward model. It works for things only a human can assess, tone, helpfulness, whether an answer is actually useful, and it inherits the annotators’ biases plus the reward model’s willingness to be gamed.

RLVR gets its signal from verification. The answer is checked mechanically; the test passes, the proof checks, the arithmetic is right. Far more reliable than a reward model, and confined to domains where correctness is machine-checkable. We looked at this family in SRL and at reward during pretraining in RPT.

RLCD gets its signal from outcomes, scored by a proper rule. Not whether a human liked the answer or whether a checker accepted it, but how well the stated probability matched what actually happened, across many predictions.

What each family can be asked to optimize, and what it costs you to ask.

The clarifying question is what each can ask for. RLHF can optimize for qualities nobody can define formally. RLVR can optimize for correctness where correctness is decidable. RLCD can optimize for honest uncertainty, which neither of the others can express, because neither has a notion of a well-calibrated answer, only of a good one and a right one.

8. Hands-on with first-order system

Which is where the theory runs out. We have the architecture, the objective, and the training loop. What none of it tells you is whether the thing is pleasant to use, and a decision model is not something you chat with and evaluate by reading; it is something your code calls a hundred times a second.

A note on how this section changed while I was writing it. I started out writing against Jev’s SDK and found myself guessing at method names, because the hosted client is thinly documented and I could not run the thing I was describing. Then, looking for something I could actually execute, I found GLiNER2.5-Decide, a 340M open model from Fastino, Apache 2.0, that does the same job locally. It also, on its own benchmark, beats both of the models this article has been about.

8.1 State, and how little of it there is

Let us see what state is in the open models variant (it is simpler than Jev’s version): just the text.

text = “”“From: compliance@group.example
Subject: Protocol update — action required today
Please confirm the new retention rule is applied before Friday’s audit.”“”

Jev additionally accepts JSON objects and arrays, which genuinely matters when your decision depends on structured records alongside prose; the crash-narrative study and the refund example both need that. GLiNER2.5-Decide takes a string, so structure has to be serialized by you. A real difference, and a point in the hosted model’s favour.

The questions are where both converge.

8.2 The smallest possible call

# pip install gliner2
from gliner2 import AutoExtractor

model = AutoExtractor.from_pretrained("fastino/GLiNER2.5-Decide")

model.classify_text(
    “Your mailbox is almost full. Click here in the next hour “
    “or we will delete every message.”,
    {”label”: [”spam”, “ham”]},
)
{”label”: “spam”}

Four lines, no API key, no network call after the weights are cached, and roughly a gigabyte of model sitting in memory. Three things worth noticing.

The label set is the second argument. Not a prompt, not a schema file, not a fine-tuning run; a Python dict, at call time. That is Section 4’s inversion in its most literal form, since the same loaded model answers a spam question now and a routing question on the next line.

There is no prompt anywhere. No instructions, no template, no “you are a helpful classifier”. The model was trained to take a label set directly, so the only English in the call is the labels themselves.

The output is the decision. A dict your code reads. No parsing, no retry, no JSON mode.

8.3 Several decisions at once

The decomposition argument from 6.2, made concrete. One text, three heads, one forward pass:

model.classify_text(
    text,
    {
        “intent”:  [”fyi”, “request”, “approval”, “complaint”,
                    “newsletter”, “security_alert”],
        “urgency”: [”low”, “normal”, “high”, “critical”],
        “route”:   [”support”, “billing”, “legal”, “security”,
                    “finance”, “archive”],
    },
)
{”intent”: “request”, “urgency”: “high”, “route”: “legal”}

Three questions, one pass over the encoder. This is the parallel-evaluation property from 6.2 arriving as something you can time yourself; add a fourth head and the latency barely moves, because the expensive part was reading the text and that happened once.

And it is the same design discipline. Three small factual judgments, combined by your code, rather than one large question that buries the policy inside the model.

8.4 Confidence, which you have to ask for

Here is the parameter the model card does not mention and which the whole of 3.3 depends on:

model.classify_text(
    text,
    {”route”: [”support”, “billing”, “legal”, “security”, “finance”, “archive”]},
    include_confidence=True,
)

It defaults to False, so by default you get a bare label and no probability, which is exactly the thing Section 5’s requirement three said not to accept. Turn it on and you have a number to threshold:

r = model.classify_text(text, {”route”: [...]}, include_confidence=True)
if r[”route”][”confidence”] > 0.90:
    assign(r[”route”][”label”])               # act alone
elif r[”route”][”confidence”] > 0.60:
    assign(r[”route”][”label”], review=True)  # act, but watch
else:
    escalate_to_reasoning_model()             # ask something expensive

The three-tier pattern from 3.3, now with a parameter that actually returns the number it depends on.

I should flag that I verified include_confidence in the package source but could not run the model to see the exact response shape, so treat the accessor above as the shape to check rather than the shape to trust. The parameter is real; the key names may differ.

And the Section 9 caveat applies with full force here. A confidence you have not validated on your own data is a number, not a threshold; measure the reliability before you branch on it.

8.5 Multi-label, thresholds, and ordinal scales

Three variations that cover most of what you will actually need.

Multi-label, when several answers apply at once:

model.classify_text(
    “Battery dies before lunch, but the keyboard and the screen “
    “are the best I have used.”,
    {”aspects”: {
        “labels”: [”battery”, “keyboard”, “screen”, “camera”, “price”, “support”],
        “multi_label”: True,
        “cls_threshold”: 0.4,
    }},
)
# {”aspects”: [”battery”, “keyboard”, “screen”]}

Ordinal scales, which are Jev’s Score primitive done the cheap way; pass the levels as ordinary strings:

model.classify_text(
    “Payroll file has to be corrected before the 5pm cutoff “
    “or the whole company is paid late.”,
    {”urgency”: [”0”, “1”, “2”, “3”, “4”, “5”]},
)

# {”urgency”: “5”}

Worth being precise about the difference. Jev’s Score understands the levels as ordered and returns a continuous value between them, 1.4 on a zero-to-two scale. This returns the nearest level as a string. For most SLA systems that is fine; if you need the interpolation, it is a real gap.

Labels with descriptions, for when the name alone is ambiguous:

model.classify_text(
    “Please reset the card PIN. The new one never arrived “
    “and the old one is locked.”,
    {”intent”: {”labels”: {
        “card_pin_change”: “The customer wants a new PIN or the current PIN replaced”,
        “card_lost”:       “The physical card is missing”,
        “balance_inquiry”: “The customer wants the current balance”,
    }}},
)

This is where the prompt quietly comes back. A description is a short instruction, and writing good ones is the same skill as writing good instruction strings for a Noul. The typed interface removes the parsing problem completely; it does not remove the specification problem.

8.6 A yes/no question over a passage

The closest thing to Jev’s Noul, and the primitive you reach for most:

model.classify_text(
    “The treaty was signed in Paris in 1992. It entered into force “
    “the following year, after the last signatory ratified it.”,
    {”answer”: {
        “labels”: [”yes”, “no”],
        “prompt”: “Did the treaty enter into force in 1992?”,
    }},
)
# {”answer”: “no”}

The text is the state, the prompt is the question, and the answer is one bit. No chain of thought, no extracted sentence, no explanation. This is the grounding-check shape, is this claim supported by this passage, which is one of the most common things people currently burn a chat-model call on.

8.7 All together

The refund case from 6.2, end to end, running locally (no API key required).

# pip install gliner2

from gliner2 import AutoExtractor

model = AutoExtractor.from_pretrained("fastino/GLiNER2.5-Decide")

def handle_ticket(message: str, transactions: list, policy: str) -> str: 

    # ---- 1. the state is a string, so serialize the structure yourself ----
    lines = [f"- {t['id']}: {t['amount']} on {t['ts']}" for t in transactions]
    state = (
        f"{message}\n\n"
        f"Recent transactions:\n" + "\n".join(lines) + "\n\n"
        f"Policy: {policy}"
    )

    # ---- 2. several independent judgments in one forward pass ------------
    r = model.classify_text(
        state,
        {
            "intent":      ["refund_request", "order_status", "cancel_subscription",
                            "login_problem", "bug_report", "speak_to_human", "other"],
            "duplicate":   ["yes", "no"],
            "urgency":     ["0", "1", "2", "3", "4", "5"],
            "needs_human": ["yes", "no"],
            "queue":       ["billing", "technical", "account"],
        },
        include_confidence=True,
    )

    # ---- 3. every business rule lives here, none of it in the model ------
    asked     = r["intent"]["label"] == "refund_request"
    duplicate = r["duplicate"]["label"] == "yes"
    urgency   = int(r["urgency"]["label"])
    queue     = r["queue"]

    if asked and duplicate and queue["confidence"] > 0.80:
      return {"action": "approve_refund"}
    elif urgency >= 4 or r["needs_human"]["label"] == "yes":
      return {"action": "escalate", "priority": "high"}
    else:
      return {"action": "queue_for_review", "suggested": queue["label"]}

    # ---- 4. log probabilities so you can audit calibration later ---------
    log_decision(outcome=outcome, raw=r)
    return outcome

if __name__ == "__main__":
    result = handle_ticket(
        message="I was charged twice for the same order and nobody has replied "
                "in three days. Please refund the duplicate.",
        transactions=[
            {"id": "tx_8821", "amount": "49.00", "ts": "2026-09-14T10:02Z"},
            {"id": "tx_8822", "amount": "49.00", "ts": "2026-09-14T10:03Z"},
        ],
        policy="Duplicate charges within 24h are refundable without review.",
    )
    print(result)
    # output : {'action': 'queue_for_review', 'suggested': 'billing'}

Four things this file demonstrates, and they are the same four regardless of which model sits behind it.

Step one is where the open model costs you something. The transactions have to be flattened into prose because the state is a string, and that flattening is a parsing problem you created. Jev’s JSON state avoids it. This is the clearest practical argument for the hosted model.

Step two is one forward pass. Five heads, one read of the text. Writing it as five calls would be five times slower for nothing.

Step three contains every business rule in the system, and none of them are in the model. Thresholds, the precedence of escalation over approval, the decision to treat low routing confidence as grounds for review; all readable, testable, and changeable by someone who has never heard of GLiNER.

Step four is the one people skip. Logging the probabilities alongside what actually happened is what lets you build the reliability diagram from 3.3 three months from now, and Section 9 will show that you will need it.

8.8 What using it actually feels like

Three observations from writing the above.

Running locally changes the calculus more than I expected. No API key, no rate limit, no per-call cost, no network hop. The Section 4 argument about cheap decisions unlocking new volume is one thing when a decision costs $0.000004 and another thing entirely when it costs nothing and runs on the machine already doing the work.

The decomposition discipline is the hard part, and it is a skill rather than a syntax. The temptation is to ask one big question, something like a single “decision” head with approve, deny and escalate as its labels, and the interface will happily let you. That puts the policy back inside the model and you lose the thing you came for.

And the label names are doing more work than they look like they are. refund_request versus wants_money_back is a prompt-engineering decision wearing a very small hat. The model reads those strings. You are still choosing words carefully, just fewer of them.

9. Results

9.1 Independent evaluation exists, which is unusual this fast

Two months after release there are at least six third-party evaluations of Jev on arXiv, which is more scrutiny than most closed models get in a year. Since the vendor’s figures are unverifiable, these are what this section rests on.

The broadest is a benchmark study from Bonn covering 37 datasets and 346,009 requests, run zero-shot on full evaluation splits with one frozen template per dataset, for under US$10 in total, which is itself a data point about the cost claim.

Where it is strong. Binary sentiment at 96.5% on IMDB and 96.4% on SST-2, language identification at 99.6%, commonsense at 95.5% on HellaSwag, science questions at 98.8% on ARC, and 86.7% on Belebele across 122 languages. Intent routing reaches 89.5% on CLINC150 with 151 options in a single Choice.

Against open-weight models scored on identical requests through their exact option probabilities, Jev beats Qwen3.8-27B on 27 of 37 datasets and Gemma-4-E4B on all 37. Worth noting what that comparison is: a hosted commercial model against a 27.8B and an 8B open model, and the margins are real but not enormous.

There is a third data point, and it is awkward for both models this article has been about. GLiNER2.5-Decide, an open 340M encoder, reports 60.2% exact-match accuracy on Fastino’s fast-decisions benchmark, 17 domains with 300 held-out examples each and the same text and candidate labels for every model, against 57.6% for JevK5 and 46.6% for Laya Router. A 340M model beating a hosted commercial one on a public benchmark is a real result, with the usual caveat that the benchmark belongs to the people who won it. Still, the direction is consistent with the Bonn study’s finding that Jev’s advantage over well-used open models is in interface and consistency rather than raw accuracy.

Where it fails. Emotion at 58.5% on noisy hashtag-derived labels, SST-5 five-way sentiment at 57.9%, GoEmotions macro-F₁ of 0.243, and worst of all AGB-DE at F₁ 0.204, deciding whether a German consumer-terms clause is void under German law, which the authors rightly describe as requiring domain knowledge rather than a snap judgment. Prompt-injection detection has perfect precision and 50% recall.

Thirty-one of the 37 datasets, sorted. The right-hand tail is where a snap judgment is the wrong tool.

9.2 The calibration picture is more complicated than the marketing

Here is the finding that matters most, and it cuts both ways.

Pooled, the calibration is genuinely good. Across 22 Choice datasets and 279,925 answers, expected calibration error is 0.028. Averaged per dataset, Jev’s ECE of 0.074 is level with Qwen’s 0.075 and far better than Gemma’s 0.184.

Per dataset, it ranges by two orders of magnitude. From 0.003 on Language ID and 0.005 on ARC, up to 0.279 on Emotion, 0.264 on AGB-DE and 0.236 on Prompt Injections. The aggregate number is an average over a spread that wide, and your task sits somewhere in it without you knowing where.

Every bar is the same model under the same evaluation harness. The dashed line is the number that gets quoted.

And Noul is weaker than Choice. Binary probabilities rank well but sit badly relative to a fixed 0.5 threshold; tuning the threshold on training data raises micro-F₁ on UNFAIR-ToS from 0.50 to 0.75. The same pattern appears in moderation: AUROC 0.989 on ToxicChat toxicity and 0.982 on prompt injection, with mediocre decisions at 0.5. The model knows which cases are more likely true; it does not know where to put the line.

9.3 Three domain studies agree

Crash narratives. A 27-question schema over 499,500 Texas police narratives, 195,857 coded by jev-1.13.0 at 3,673 mean input tokens and 0.20 s median latency. Against 2,416 blinded human judgments: F₁ 0.908, with precision 0.902 and recall 0.915. One frontier model gains 0.059; the other is statistically indistinguishable. Raw probabilities overstate prevalence; recalibrated on those judgments, out-of-fold calibration error falls to 0.0069 with slope 0.97. The authors’ conclusion is that calibration varies by model rather than by paradigm, so each model must be audited.

Radiology. Jev judging factual agreement between generated and physician-written reports. Raw ECE 0.1235, Brier 0.1269, AUROC 0.8858. After isotonic calibration on a held-out half, ECE falls to 0.0077, a sixteen-fold improvement. On controlled errors it reaches AUROC 0.977 for detecting a reversed finding. The recommendation is explicit: recalibrate per task before thresholding.

Medicine. The least flattering, and preliminary. ECE from 0.063 on MetaMedQA to 0.141 on PubMedQA. At a 0.9 threshold, accuracy is 75.9% at 17.3% coverage, meaning roughly one in four of its most confident answers was wrong.

A human-activity-recognition study adds the blunt version: fast and inexpensive to query, but its probabilities are not reliably calibrated for recognition.

9.4 The pattern

Six studies, one consistent finding. Calibration out of the box is decent on average and unreliable per task, and post-hoc recalibration fixes it dramatically. Every study that tried it reports a large gain from modest labelled data.

A few thousand resolved outcomes from your own distribution, and an isotonic fit, closes most of the gap.

Practitioners report the same from the other direction; this walkthrough makes the point that the outputs are not always calibrated and a confidently wrong answer is entirely possible. The marketing line that the model “never hallucinates” is true only in the narrow sense that it cannot emit a category you did not define. A wrong answer at 0.95 confidence is the same problem wearing a type signature.

Which lands awkwardly against Section 7. The argument for RLCD was that calibration should be trained in rather than patched on, that temperature scaling is a remedy for a model optimized for the wrong objective. The empirical finding is that you should fit an isotonic map on your own data before trusting a threshold. That does not make RLCD pointless, since starting at 0.1235 beats starting at 0.4, but it qualifies the claim considerably. Calibration is a property of a model on a distribution, and you are not running it on the distribution it was trained for.

9.5 The caveat the benchmarks hide

Now the thing most coverage skipped, and the most important practical fact here.

Laya’s shipped checkpoints work zero-shot, and on Convai’s own 2,000-decision benchmark they score 0.362. Fine-tuned on the target domain, the same model scores 0.766. Roughly double.

Read that against Section 4’s framing. The promise was one model serving every label set without retraining; the measurement says that out of the box, on a broad benchmark, the open model is mediocre, and most of its headline accuracy comes from domain adaptation. Which is a classifier with extra steps.

Two things soften it. The number is Laya’s, not Jev’s, and Jev is considerably larger and hosted, so its zero-shot gap may be much smaller, though TypeSafe explicitly does not fine-tune on customer data, which means whatever zero-shot performance you get is what you have, permanently. And a single aggregate over 2,000 heterogeneous decisions says little about any particular task; a three-way routing problem with well-separated options is not a fine-grained multilingual rubric.

It also connects straight back to 7.3 in the methodology. If calibration is learned from outcome distributions, zero-shot transfer is a question about how much your decisions resemble the training ones, and nobody has published enough for you to answer that in advance.

The honest summary: the category’s zero-shot generality is the least-evidenced of its claims, and it is the first thing I would test on my own data rather than take on trust.

9.6 What we still do not know about Jev

Laya tells us what the category looks like. It does not tell us what TypeSafe built.

Unknown: the architecture beyond non-autoregressive, the model size, the training data, the sampler they describe as novel, the specifics of their RLCD implementation, and any reliability data for the calibration claim. The 193.6× and 444.6× figures are vendor comparisons against vendor-chosen baselines.

What we can say is that an open 421M encoder reaches roughly 33 ms with competitive calibration using published techniques, which makes TypeSafe’s claims plausible in kind if not in magnitude. That is weaker than reading their paper, and considerably better than what we had two months ago.

10. Thoughts

1. Jevons was the right name, and more right than intended.

Jevons observed that making a resource cheaper increases total consumption. If a decision costs 400× less, the consequence is not that your bill falls 99.75%, it is that you start making decisions where nobody currently makes one. The benchmark study is the proof: 346,009 decisions for under ten dollars.

Which is fine until the second-order version. A system making a million calibrated judgments a minute is one whose behaviour is set by a million small thresholds that nobody reviewed individually. The cost curve is the easy part; the governance of a hundred thousand confidence thresholds is the part nobody has thought about.

2. The speedup numbers are measured against the most expensive possible baseline.

“40–200× faster” and “193.6× faster” are comparisons against frontier models, the Opus 5.5 and GPT-5.6 class, large models with reasoning paths behind a hosted API. Of course a 400 ms typed decision beats a multi-second reasoning call by two orders of magnitude; that baseline was never the right tool.

Approximate, and the point is the ratio rather than any one number. Against the right baseline the gap is real but ordinary.

Compare against what a sensible engineer would actually reach for, GPT-OSS, or Qwen or Gemma at a few billion parameters doing constrained decoding on a short classification, and the gap compresses to roughly 2–3×. The Bonn study is instructive here: it scores Qwen3.8-27B and Gemma-4-E4B on identical requests via a single forward pass over option tokens, which is the right way to use an LLM for this, and Jev’s advantage is in accuracy and interface rather than orders of magnitude of latency.

The sharpest version comes from the open side. Laya is a 421M encoder running locally at 32.8 ms, against Jev’s 236–276 ms. The open model is 7.8× faster than the thing whose headline figure is 193.6×, which tells you most of that number was about the baseline rather than the method.

The architecture argument still holds; non-autoregressive single-pass evaluation genuinely removes the sequential floor. It is the magnitude that is a marketing artifact.

3. Output tokens are priced at zero, which is not the same as costing nothing.

TypeSafe charges US$0.042 per million input tokens and nothing for output. So in billing terms the output genuinely is free.

But that is a pricing decision rather than a statement about physics. Producing the answer still costs compute; it is bundled into the input price, which TypeSafe can afford to do precisely because the output is two bits rather than a paragraph. And if reasoning is involved it is charged normally, so “free output” describes the choice, not the deliberation that might precede it.

The interesting consequence is the incentive it creates. With input dominant and output free, the economics push you to send more context, which is the opposite of the instinct everyone developed optimizing chat prompts. The crash-narrative study runs 3,673 input tokens per decision without comment.

4. Grounding reduces hallucination and relocates the failure.

The model cannot invent a category outside your answer space, which removes a whole class of failure. That is real.

But it does not remove being wrong, it concentrates it. One in four of the medicine benchmark’s most confident answers was incorrect; the model was not hallucinating, it was selecting the wrong option from a valid set with high confidence. And a bare value with no rationale leaves nothing to inspect. A chat model that reasons its way to a wrong answer at least leaves a trail.

So the choice becomes a single point of failure in a way prose is not. The mitigation is the decomposition discipline from 8.6, several small questions whose answers can disagree, rather than one large one that cannot.

5. Is this a step toward AGI?

No, and the framing is wrong in an interesting way. This is not a weaker general intelligence, it is a parallel branch on a simpler assumption: that most decisions do not require general intelligence. Routing a ticket needs no world model.

Where it touches the deeper argument is on representations. LeCun’s position, which we went through in the LeJEPA article, is that intelligence is largely about representations good enough to support prediction, and a System One model is precisely a bet that a good representation plus a bounded answer space is all most tasks need.

6. The thing I would actually watch.

Not speed, which will be matched. Not cost, which will fall. Whether calibrated uncertainty becomes a standard thing to expect from a model.

We have spent three years building on outputs whose confidence means nothing, compensating with evals and guardrails and human review. If the number becomes trustworthy enough to branch on, that changes what software can safely delegate far more than another order of magnitude on latency would.

The evidence says we are partway: calibrated on average, unreliable per task, fixable with modest labelled data. A real capability with a real caveat.

11. Conclusion

We talked about Jev, and specifically about the gap it is trying to close; most of what we use AI for in production is a judgment rather than a piece of writing, and we have been making those judgments with a text generator because that is what was available.

The route out was noticing that this is not a rough edge but a category error. A sentence encoding a decision is mostly grammar, so you pay forty sequential forward passes for two bits; constrained decoding narrows the search space without touching the computation; and the probability you get back is the likelihood of a token appearing at a position, which is not the model’s belief about the answer and was never trained to be.

That turned the question from how do we parse the output into what shape should a decision model be, and the answer has a contradiction in it. It has to read like a language model and answer like a function, and the thing that reads like a language model is an autoregressive generator, which is slow, untyped, and gives you token probabilities. The resolution is slightly embarrassing in hindsight: we already had the architecture. A bidirectional encoder was always the right machine for understanding a fixed input, and we stopped reaching for it when generation became the goal.

Calibration was the harder half. Cross-entropy against hard labels tells a model it should have said one hundred percent, every time, which is why modern networks are confidently wrong. A strictly proper scoring rule asks a different question, how often are you right when you talk like that, and has a unique optimum at the model’s honest belief. That is a metric from the 1950s moved from the evaluation step to the training step, wrapped in reinforcement learning because calibration is a property of a distribution of predictions rather than of any single one.

The empirical case is more mixed than the marketing. Thirty-seven datasets, 346,009 decisions for under ten dollars, beating a 27B open model on most of them. Speed real but smaller than advertised, since 193.6× was measured against a frontier model and an open 421M encoder runs 7.8× faster than Jev itself. And calibration pooled at 0.028 while ranging from 0.003 to 0.279 across tasks, which is decent on average and unreliable where it matters.

What strikes me most is how little here is new. The encoder is BERT’s architecture, the Brier score is from 1950, proper scoring rules have been the standard tool for evaluating forecasts for sixty years. The contribution is noticing that a decision is not a short piece of text, and everything else follows from taking that seriously.

Whether Jev specifically becomes standard I have no idea. The durable part is the claim underneath it: that a model can tell you how much to trust it, in a number you can act on. We have spent three years building on outputs whose confidence means nothing, and if that changes, it changes what software can safely delegate more than any amount of latency ever would.

This marks the end of our walk through System One models, and with this we conclude.

References:

TypeSafe docs: https://docs.typesafe.ai

Laya, open weights: https://huggingface.co/convaiinnovations/laya

Benchmark study: https://arxiv.org/abs/2609.37647

Crash narratives at scale: https://arxiv.org/abs/2609.24052

Radiology factuality: https://arxiv.org/abs/2609.27607

That’s all for today. Follow me on LinkedIn and Substack for more such posts, till then happy Learning. Bye 👋

1. Introduction

2. What a JEPA actually is

   1. What is a world model?

   2. Why predict in representation space?

3. Connecting probability with optimization

4. The heuristics zoo, and what “under-specified” really means

   1. Two perspective of SSL tasks

   2. Various methods (Negative samples, Stop-gradients, Whitening, etc.)

5. The Ideal Case

6. Methodology

   1. Finding Q: why the isotropic Gaussian?

   2. Finding D: measuring the distance to a Gaussian

7. SIGReg in code

8. Results & Outcomes

9. Thoughts

10. Conclusion

Introduction

I have been a fan of self-supervised learning for a very long time, and at this point I don’t think that qualifies as a niche taste anymore. If you look underneath almost anything that has worked in deep learning recently, you will find some version of SSL doing the actual work, and it tends to fall into one of two camps. Next-token prediction in LLMs and most of what we do with time series are learning a temporal world; what follows what, and how state evolves along an axis we cannot run backwards. The DINO family, masked image modelling, contrastive image-text pairs are learning a spatial world; what belongs with what, and how structure holds together across a scene that is all present at once. The supervision never really went away in either case, it just stopped coming from a human and started coming from the structure of the data itself. We had been writing multiple articles on SSL and respective methodologies; in case if you want to brush up some more details:

The interesting question then is no longer whether SSL works, we have more than enough evidence for that. It is what exactly we are asking a model to do when we train it this way, and among all the SSL families, JEPA has committed to the sharpest answer by deriving the model from the combination of spatial and temporal which is a video. A video is temporal structure laid over spatial structure, and if your input space is video, then the thing you are actually trying to learn is neither the layout of a frame nor the ordering of frames, but the dynamics that carry one into the next. The way JEPA goes after this is by predicting in representation space rather than in pixel space, which frees the model to discard whatever it cannot predict instead of wasting capacity reconstructing it. That single shift does most of the heavy lifting here, and it is worth unpacking properly, which we will do in the next section.

There is a framing of all this that I keep coming back to, and it comes from Yann’s lecture at NYU. SSL, in his telling, is not really a learning trick at all; it is an optimization problem with constraints, and the constraints are where the entire inductive bias of the method lives. The prediction task on its own is trivially solvable, because a model that maps every input to the same constant vector predicts perfectly and learns absolutely nothing in the process. Hence the constraint is not some detail bolted onto the objective for stability, it is closer to being the objective itself, and everything we discuss in this article follows from taking that seriously.

Which is exactly where things get uncomfortable. If the constraint is carrying all of the inductive bias, then we should have some principled theory of which constraint to use, and we simply don’t. What we have instead is an accumulated pile of mechanisms that were found empirically and kept because they happened to work; negative samples (SimCLR, MoCo), teacher–student networks with hand-tuned EMA schedules (BYOL, DINO), stop-gradients (SimSiam), feature whitening (Barlow Twins, W-MSE), and explicit variance-covariance penalties (VICReg). Every one of these prevents the model from cheating, and not one of them was derived from a statement about what the embeddings ought to look like, largely because no such statement existed to derive from.

LeJEPA is an attempt to finally write that statement down. Balestriero and LeCun ask which distribution the embeddings should follow in order to be useful for downstream tasks that nobody has seen yet, prove that there is exactly one such distribution, construct a penalty term that enforces it, and then delete the rest of the machinery; no stop-gradient, no teacher–student, no schedulers, one hyperparameter. The route they take to get there is genuinely lovely, involving an isotropy argument that falls straight out of plain bias-variance, a theorem that lets you read a high-dimensional distribution entirely off its one-dimensional shadows, and a statistical hypothesis test repurposed as a loss function.

As always, we will build this from the why forward, so that by the end the objective stops looking like a proposal someone made and starts looking like the only thing that could have been written.

(Usual caveat: the information in this article is upper bounded by my own grasp of the topic, my patience for typing for hours, and a writing medium that still refuses to support inline LaTeX. Please re-visit everything you read here, and go read the paper.)

2. What a JEPA actually is?

Let us start by understanding the claim which JEPA starts with, i.e. it being a world model.

2.1 What is a world model?

The phrase gets used loosely enough to have lost most of its meaning, so let me pin down the version that matters here.

A world model is an internal representation of how things work that is good enough to run forward. Not a recording of what happened, a mechanism for anticipating what happens next. You have one. It is why you flinch before the ball hits you, why a wobbling glass makes you reach out, and why you can read half a sentence and know roughly where it is going. None of that involves retrieving a stored instance; it involves running a model of the situation slightly ahead of the situation itself.

Three properties matter, and each one turns out to be a design constraint.

1. It predicts, so it must be usable in the forward direction. A representation that describes the present beautifully but supports no inference about the next moment is a description, not a model. The test is always whether you can push it forward and get something useful.

2. It is abstract, so it discards. The version of the falling cup that lets you predict the outcome contains gravity, brittleness and height. It does not contain the reflection on the ceramic. Detail that does not affect what happens next is not merely unnecessary, it is actively in the way; a model cluttered with it is harder to run and easier to confuse.

3. It is reusable, so it is not built per task. The same intuitions about objects and support and momentum serve you when catching a ball, stacking dishes and judging whether a shelf will hold. You did not learn a separate model for each. That reuse is exactly what “foundation model” is meant to name; one representation, many tasks, none of them known when the representation was learned.

It is worth seeing where this sits in the larger picture, because LeCun’s proposal for autonomous machine intelligence is considerably bigger than a world model alone.

LeCun’s proposed architecture for autonomous machine intelligence. The world model is one module among several.

Reading the diagram: perception turns percepts into a representation of the current state. The world model takes that state, plus an imagined action, and predicts the state that would follow, which is what lets the system evaluate a plan without executing it. The actor proposes actions, the cost and critic modules score the predicted outcomes, short-term memory holds the trajectory being considered, and the configurator adjusts the other modules for whatever task is at hand.

The point I want to draw from this is a limiting one. JEPA is a proposal for how to train the green box, and nothing else. There is no actor here, no cost module, no planning loop; the architecture around it is a research programme rather than something anyone has built. So when we talk about world models in this article we mean the representation-learning half specifically, and questions about what you would do with one, planning, control, action-conditioning, are outside what LeJEPA addresses.

Which is why the paper’s framing matters. If the world model is the piece everything else depends on, then getting its representation right is not one optimization problem among many, it is the one that determines whether the rest of the architecture has anything to work with.

The same scene under three objectives. Only the third keeps what determines the outcome and lets the rest go.

Why prediction is the natural way to learn one

Here is the part that makes this more than philosophy. If you want a model with those properties and you have no labels, prediction is the only training signal available that produces them for free.

To predict what comes next you must represent what determines what comes next. There is no shortcut; a model that ignores the causal structure cannot anticipate it. So the prediction task does not merely test the world model, it forces one into existence, and the pressure to be abstract comes along automatically, because carrying unpredictable detail costs capacity and buys nothing.

This is the argument for self-supervision at its strongest, and it explains why the field keeps returning to it across modalities that otherwise share nothing. Next-token prediction, next-frame prediction and masked-patch prediction are the same bet: that predicting is how you are forced to understand.

Where JEPA fits

JEPA is a specific and, I think, correct answer to the question of where the prediction should happen.

Everything above says predict. Nothing above says predict in pixel space, and that turns out to be the crucial degree of freedom. Reconstruct the next frame and you are graded on the reflection on the ceramic, which forces exactly the detail-hoarding that the second property said to avoid. Predict in representation space and the model is free to discard whatever it cannot anticipate, which is what the same property demanded.

So JEPA’s claim is narrow but load-bearing: predict, but predict abstractly, and let the abstraction be learned rather than specified. The rest of this section is what that looks like when written down.

2.2 The setup

The architecture itself is deliberately plain, and it helps to see it as four steps before we complicate it:

1. Take two views of the same underlying thing.

2. Push both through an encoder Enc(·), which lands them in some embedding space R^K.

3. Ask a small predictor Pred(·) to produce one of those embeddings from the other, where one is used at all.

4. Measure how far it missed, and backpropagate.

Everything interesting is hiding in that first step, because what counts as “two views of the same thing” is entirely up to you. If your data is video then time gives you views for free, and view one is frame t while view two is frame t+1, which means you are effectively asking the model where the world is going next. If your data is images then you have to manufacture views instead of finding them, so you take two random crops of the same photo, or a masked patch and the context around it, or the same image under two different colour jitters, and you ask the model what belongs with what. The machinery downstream is completely identical in both cases and only the view-generator changes, which is why I will mostly just say “two views” from here on and stay agnostic about where they came from.

2.3 Why predict in representation space?

We asserted this in 2.1 and it is worth earning properly, because the cleanest way to see it is to write the two objectives next to each other.

Reconstruction, which is what autoencoders and masked image models do, predicts back into the input space:

L_recon(θ) = E [ ‖ Dec(Enc(x̃)) − x ‖² ], x ∈ R^D

whereas JEPA predicts into the embedding space instead:

L_pred(θ) = E [ ‖ Pred(Enc(x_v)) − Enc(x_v′) ‖² ], Enc(x) ∈ R^K

The same two encoders, two different places to put the target. Moving it from pixel space to representation space is the entire architectural claim.

The difference between these two lines looks almost cosmetic, and it is anything but, because it changes three things at once.

The first is that unpredictable detail becomes free to throw away, which is the world-model property from 2.1 arriving as a loss function. Under reconstruction every pixel of the missing patch contributes to your loss, so the model is actively penalized for failing to predict things that are not predictable in principle; the position of each leaf on a tree, the grain of the sensor noise, the exact texture of gravel. A real fraction of model capacity ends up going into hallucinating plausible detail, which is a genuine skill but not the one you were trying to teach.

Suppose I nudge a cup off the edge of a table and ask what happens next; you will say without hesitation that it falls, hits the floor and probably breaks. What you will not do is give me the coordinates of each ceramic shard, and if I insisted you would rightly say the question is unreasonable, not because your understanding is poor but because that information does not yet exist in any form your prediction could access. Your grasp of the dynamics is excellent and your ability to reproduce the pixels is nil. Reconstruction grades you on the second; JEPA grades you on the first.

The second consequence is subtler and matters more than it first appears: the target is no longer fixed. In reconstruction, x is handed to you by the dataset and cannot move no matter what the optimizer does, whereas in JEPA the target Enc(x_v′) is itself a function of θ, so both sides of the loss shift every time you take a step.

Which leads directly to the third, namely that this is exactly where the trouble begins. A fixed target makes collapse structurally impossible, because there is no way to cheat your way into reproducing a photograph you have never seen. A learnable target makes collapse the single easiest thing the optimizer can do, and it will find that shortcut unless you stop it. Hence JEPA needs anti-collapse machinery for a structural reason that autoencoders never had, and that is not a flaw so much as the price of the freedom we bought in the first consequence.

One aside about the predictor network, since it will come back later. Its stated job is to absorb whatever systematic difference exists between the two views before the comparison, and that is well-motivated when the views really are asymmetric in information content, for instance when you condition the prediction on an observed robot action so that there is an actual transformation to be learnt. It stands on much shakier ground when the views are symmetric, in which case it looks less like a modelling choice and more like a device for breaking an inconvenient symmetry. Hold that thought; it turns out LeJEPA does not need one at all.

2.4 The definition, and the crack running through it

The paper states the JEPA criterion compactly, and the shape of the statement matters as much as its content:

Enc(x_{t+1}) is predictable from Enc(x_t), for all n and t, and Enc(·) is not degenerate.

It is worth reading twice, because it is two clauses joined by an “and”, and those clauses are not remotely the same kind of object.

The first is a loss function, specifically the L_pred we wrote above; you can put it on paper, differentiate it, and hand it to an optimizer with no ambiguity about what you are asking for. The second is an English sentence. “Not degenerate” is a property everyone agrees they want and nobody has specified, which means it has no gradient, no closed form, and nothing you can minimize, so in practice it gets silently translated into whichever mechanism happened to fix the last failure someone observed. That crack between a clause you can compute and a clause you can only gesture at is, in a very real sense, the entire subject of this article.

2.5 Writing the whole thing as one objective

Let us now push both clauses into a single equation, because doing that makes the missing piece impossible to ignore.

The general form of constrained optimization is something we built up from scratch in the regularization article, and the shape of it is that you want to minimize an objective while respecting a constraint,

min_θ f(θ) subject to g(θ) ≤ 0

which you then fold into a single expression by attaching the constraint to the objective through a multiplier:

L(θ) = f(θ) + α · g(θ)

[EMBED: Regularization : What? Why? and How? (Part -1)]

Mapping JEPA onto that gives us the skeleton that the rest of this article exists to fill in, with the predictive term playing the role of the objective and the anti-degeneracy requirement playing the role of the constraint:

L_JEPA(θ ; X) = L_pred(θ ; X) + α · L_constraint(θ)

where L_pred is f(θ), i.e. clause one, and L_constraint is g(θ), i.e. clause two.

Three things fall out of this, and together they set up everything that follows.

The first is that f is settled while g is not, since every JEPA variant in the literature agrees almost exactly on the predictive term, and every disagreement between them, without exception, turns out to be a disagreement about g. The second is that every existing method is really an unwritten g in disguise, because stop-gradients, EMA teachers, whitening layers, negative samples and VICReg’s variance-covariance penalties are all mechanisms that behave like constraints without ever having been stated as one, and several of them are not even in the loss at all but are instead surgery performed directly on the gradient flow. The third is that α cannot exist until g does, which is worth noticing because a single principled trade-off parameter is only available to you once the thing being traded off has an actual form; if your constraint is a collection of architectural interventions then there is nothing to put a coefficient on, and you are left tuning schedules instead.

All of which lets us state the paper’s contribution in one line, before we have looked at a single piece of its machinery: LeJEPA is a specific, derived choice of g, and everything from Section 6 onwards is the derivation of it.

2.6 What “degenerate” actually looks like

Since g exists to prevent degeneracy, we should be concrete about what it is preventing, and there are two versions of the failure with the second being considerably more dangerous than the first.

Complete collapse is the blunt version, where the encoder simply maps every input to the same constant vector:

Enc(x) = c ∀x ⟹ L_pred(θ) = 0

Both embeddings are then identical, the predictor happily learns the identity function, the loss drops to zero, and the model has learnt precisely nothing. This is the failure mode we flagged in the introduction and it is the easy one to deal with, because it is loud in the loss curve and obvious the moment you look at the embeddings.

Dimensional collapse is the version that actually ruins your week. Here the encoder does not push everything to a single point, but it does confine everything to a low-dimensional subspace of the embedding space, so that if we write Σ = Cov(Enc(x)) with eigenvalues λ₁ ≥ … ≥ λ_K, what we observe is

λ_k ≈ 0 for all k > r, where r ≪ K

meaning your representation is nominally 1024-dimensional while the model is genuinely using perhaps thirty of those directions and the remaining several hundred carry essentially no variance at all. What makes this so much worse than complete collapse is that nothing in your training run will tell you it is happening; the prediction loss looks entirely healthy, distinct views really are being mapped to distinguishable embeddings, the model genuinely is solving the task you set it, and you typically discover the problem much later when a linear probe underperforms for reasons you cannot localize to anything in particular.

It is also worth noticing that this is not a binary condition in the way complete collapse is. Dimensional collapse is a spectrum, every trained model sits somewhere along it, and the honest question is never whether your model has collapsed but rather how much of the embedding space you paid for is doing real work. Which hands us a considerably sharper question than the vague one we started this section with, because if what we actually care about is how variance gets distributed across the directions of the embedding space, then the thing we should be asking is: what is the right distribution of variance across those directions?

Hold onto that, because it turns out to have a precise answer, and pinning that answer down is the whole of Section 6.

3. Connecting probability with optimization

We closed the last section with an objective of the form f(θ) + α·g(θ), and I asked you to take on trust that this is the natural shape for a JEPA to take. That was slightly unfair, because written that way the constraint term looks like something we bolted on because we needed it. So before we go hunting for the right g, it is worth showing that this shape is not a convenience at all; it drops out of probability theory on its own, and once you see where it comes from, the presence of a constraint term stops being a design decision and becomes something you could not avoid even if you wanted to.

3.1 Starting from Bayes

Start from the most familiar statement in probability:

P(A | B) = P(B | A) · P(A) / P(B)

Nothing controversial here, it is a rearrangement of the definition of conditional probability, and A and B are just labels. That last part matters more than it usually gets credit for, because it means we get to decide what goes in those slots depending on what we are actually trying to learn.

In our setting there are two objects in play; the data X that we observed, and the parameters θ of the model we are fitting. Putting those in gives us

P(θ | X) = P(X | θ) · P(θ) / P(X)

and each piece here has a job. P(θ|X) is the posterior, the thing we are trying to maximize, and it asks how plausible a particular set of parameters is given the data we have. P(X|θ) is the likelihood, which asks how probable our observed data would be under a given setting of the parameters. P(θ) is the prior over parameters, what we believed about reasonable models before seeing anything. And P(X) is the evidence, which is a normalization term.

We can drop the evidence immediately, and it is worth being precise about why. P(X) does not contain θ anywhere, so its gradient with respect to θ is exactly zero, which means it shifts the value of the objective without ever changing where the optimum sits. It is not being neglected as small, it genuinely has no say in the answer.

3.2 Taking logs

The remaining expression is a product, which is awkward to optimize; gradients through products get messy and probabilities underflow the moment you multiply enough of them together. Since log is monotonic, whatever maximizes a quantity also maximizes its log, so we lose nothing by taking logs and gain a sum:

log P(θ | X) = log P(X | θ) + log P(θ)

And there it is. The objective we want, on the left, is the sum of a fit term and a prior term, and nobody put that second term there deliberately; it fell out of Bayes and there is no version of this equation that does not contain it.

That second term is one you have used your whole career, usually without calling it a prior. Put a Gaussian prior on the weights and −log P(θ) becomes a squared penalty on their magnitude, which is weight decay. Put a Laplace prior on them and you get L1. The entire discussion we had in the regularization article about penalizing large weights and constraining model complexity was, underneath the geometry, an argument about how to choose P(θ).

3.3 Where LeCun’s framing changes the picture

So far this is standard, and it gets us a constraint on the parameters. But the thing we care about in self-supervised learning is not really the weights, it is the representation, and this is where LeCun’s way of setting the problem up becomes useful.

The move is to stop treating the representation as an intermediate quantity that happens to appear inside the network, and start treating it as part of what we are inferring. Instead of asking about θ alone, we ask about the pair:

P([θ, Z] | X) = P(X | [θ, Z]) · P([θ, Z]) / P(X)

Same equation, same evidence term dropping out for the same reason, same logs:

log P([θ, Z] | X) = log P(X | [θ, Z]) + log P([θ, Z])

I should be upfront that there is a modelling stance hidden in this substitution rather than a derivation. In a JEPA the encoder is deterministic, so Z = Enc(X; θ) is fully determined once you fix θ and X, and treating Z as something to be inferred alongside θ is a choice we are making rather than something the mathematics forces on us. But it is the right choice, and it is more or less the whole of LeCun’s position; the representation is the object of interest and the weights are merely how we happen to parameterize it. If that is what you believe, then the representation belongs in the objective, and Bayes has no objection to putting it there.

Now flip the sign, so that we are minimizing a loss like everywhere else in this article rather than maximizing a probability. Negating the log posterior gives us

−log P([θ, Z] | X) = −log P(X | [θ, Z]) − log P([θ, Z]) + log P(X)

and since log P(X) contains neither θ nor Z, it shifts the value without moving the optimum, so it survives in the equation but disappears from the minimization:

argmin over (θ, Z) of [ −log P([θ, Z] | X) ] = argmin over (θ, Z) of [ −log P(X | [θ, Z]) − log P([θ, Z]) ]

where the first term is f(·), the fit, and the second is g(·), the prior.

Finally we attach a coefficient, because in practice we do not want the fit and the prior weighted equally. This last step deserves an honest label, since the moment α appears we are no longer computing a posterior and the equality would be a lie; what we are doing is defining a training objective that is inspired by one:

L(θ, Z ; X) ≜ −log P(X | [θ, Z]) + α · ( −log P([θ, Z]) )

with the semicolon separating what we optimize over, θ and Z, from what the loss is conditioned on, X.

min over (θ, Z) of L(θ, Z ; X)

which is exactly the f + αg we wrote down at the end of Section 2, now arrived at rather than asserted. Setting α to one recovers exact Bayesian inference; anything else is a deliberate statement about how seriously we take the prior relative to the data, and that knob is the same Lagrange multiplier that appeared in the regularization article when a hard threshold got softened into a penalty.

3.4 What lives inside the joint prior

The interesting term is P([θ, Z]), and the reason it is interesting is that it is a joint prior over both objects at once. It is not a prior on weights with a representation term stapled on; it is a single statement of belief about what a good (parameters, representation) pair looks like, and that gives us more room than we are used to having.

There are three kinds of belief you can express in there.

The joint prior gives three distinct places to put a belief. Almost all classical work sits in the leftmost panel; every SSL constraint we are about to look at sits in the middle one.

You can place belief on θ alone, which is classical regularization and is what almost all of the literature has historically done. Weight decay, L1, spectral norms, anything that says a good model has small or well-behaved parameters. This constrains the machine.

You can place belief on Z alone, which is far less explored and is where this entire article is heading. Statements like “the representation should not collapse”, “the representation should use its dimensions evenly”, “the representation should follow some particular distribution”. This constrains what the machine produces, which is a genuinely different thing to constrain.

Or you can place belief on the coupling between them, which is the part that only exists because we kept the prior joint. Anything that says a certain kind of representation should only arise from a certain kind of parameterization lives here, and it is largely unexplored territory.

One detail about the Z part is worth pausing on, because it explains something that confuses people the first time they implement any of this. A belief about the representation is almost never a belief about an individual embedding vector; it is a belief about the distribution of embeddings that the encoder induces across the data. Collapse is not a property any single vector can have. A lone embedding sitting at the origin is unremarkable, and it is only when every embedding sits at the origin that anything has gone wrong. Hence any term in this slot has to be computed over a population rather than per-sample, which is precisely why every anti-collapse mechanism in the literature operates on a batch and why none of them can be written as a per-example loss.

The orange point is the same embedding in all three panels. Nothing about it tells you whether the representation has collapsed; only the cloud around it does.

3.5 Why the constraint carries the inductive bias

Now Yann’s framing from the introduction stops being a slogan and becomes something you can point at in the equation.

Recall from Section 2 that the fit term is trivially satisfiable. A constant encoder drives the predictive loss to zero and learns nothing, which in the language we have just built means f(·) does not determine the solution; there is an enormous family of (θ, Z) that optimize it perfectly and they differ wildly in how useful they are. The fit term tells you which solutions are admissible. It says nothing about which one you should end up at.

The prior term is what chooses. Which means whatever you put in that slot is not a safety mechanism protecting your training run, it is the thing that decides what your model actually learns. Two JEPAs with identical predictive losses and different priors are not two implementations of one idea with different stability characteristics; they are different learning algorithms that will converge to genuinely different representations.

That is an uncomfortable thought when you set it next to the list from the introduction, because none of those methods wrote their prior down. Stop-gradients, EMA teachers, whitening layers, negative samples, variance-covariance penalties; every one of them is an attempt to express a belief about Z without ever stating what that belief is, and several of them are not even in the loss but are instead surgery performed directly on the gradient flow. We have been choosing the most important part of the objective by trial and error, and then describing the result as a stability trick.

That is the gap LeJEPA walks into. Its claim is that the Z part of this prior has a correct form, that the form can be derived rather than guessed, and that once written down properly it makes every mechanism that was approximating it unnecessary. The next two sections look at what those mechanisms actually are and where each one falls short, and then we can go looking for the answer.

4. The heuristics zoo, and what “under-specified” really means

We now have a slot in the objective that needs filling, and we have the observation that every existing method is filling it without saying so. It is worth going through those methods properly, not to catalogue them but to ask the same question of each one; what belief about Z is this mechanism trying to express, and where does the expression fall short? Because the pattern that emerges is remarkably consistent, and it is what tells us what we should be asking for in the first place.

Most of these mechanisms we have met before, since I went through the contrastive, clustering and distillation families in some detail in the SSL primer; what I want to do here is different, which is to re-read each of them through the lens we just built and ask what belief each one is really trying to express.

It helps to keep the base form from the last section in view while we do this, because every method below is an answer to the question of what g should be, and a few of the answers turn out to be that there isn’t one:

L(θ, Z ; X) = f(θ, Z ; X) + α · g(θ, Z), where g(θ, Z) = −log P([θ, Z])

4.1 Two perspective of SSL tasks

Before going mechanism by mechanism, it is worth borrowing the taxonomy LeCun uses in the same lecture series, because it explains why the zoo has the shape it does rather than being an arbitrary collection of tricks.

The framing is energy-based. Define some scalar E(x) that is low for configurations the model considers plausible and high for the ones it does not; for our purposes E measures how compatible two views are under the current encoder, so low energy means “these belong together”. Training then has exactly one job, which is to push energy down on the data.

And here is the difficulty that generates everything else. Pushing energy down on the data is trivially easy. A model that assigns low energy to everything achieves it perfectly, and that is the same statement as collapse written in a different notation; a flat energy surface and a constant encoder are the same failure. So the entire problem of self-supervised learning is not making energy low where it should be, it is making energy high everywhere else.

There are exactly two ways to do that, and every method we are about to look at is one or the other.

Orange points are the data. Both families end up with energy that is high away from the data; they differ entirely in how they get there.

Contrastive methods push the energy up explicitly. You pick some points, you increase their energy, and the surface acquires shape because you carved it. The variants differ only in how the points get chosen; you can push up everywhere at once, which is maximum likelihood and which needs either a tractable partition function or a variational approximation because “everywhere” is an integral over the whole space. You can push up at chosen locations, which covers most of what we actually use, including metric learning and siamese networks, noise contrastive estimation, contrastive divergence, and, viewed from the right angle, the discriminator in a GAN. Or you can train a function that maps off-manifold points back onto the manifold, which is what denoising autoencoders do and what masked models like BERT do, where the corrupted input is the high-energy point and the reconstruction is the push back down.

Regularized and architectural methods limit the volume of low-energy space instead. Nobody is pushed up anywhere. Instead you arrange matters so that there is only so much low energy to go around, and the model is forced to spend it where the data actually is. These come in variants too. You can build the machine so the low-energy region is bounded by construction; PCA has a fixed number of components, K-means a fixed number of centroids, a normalizing flow is volume-preserving by design, and in none of these is there a term in the loss doing the work, since the architecture simply makes the failure unreachable. You can add a regularization term that measures the volume of low-energy space and penalizes it, which is sparse coding, sparse autoencoders, VAEs with their KL term, and VQ-VAE with its finite codebook. You can make the output as insensitive as possible to the latent, which is what contracting and saturating autoencoders do by penalizing the Jacobian. Or you can shape the surface locally, minimizing the gradient and maximizing the curvature around data points, which is score matching.

Mapping this back onto our equation

Now read both families through the objective from Section 3, because the taxonomy predicts exactly what we are about to find.

Contrastive methods put the push-up directly into g, and we will see this concretely in a moment when InfoNCE splits apart; the denominator is a sum over chosen points whose energy is being raised, and it is a real term with a real value. The trouble is that the sum has to run over the points you chose, so the cost scales with how many you need, and in high dimension the number needed to meaningfully shape a surface grows uncomfortably. That is not an implementation detail you can optimize away, it is intrinsic to the decision to push up pointwise in a large space.

Architectural methods have no g at all, which is precisely the g ≡ 0 we are going to keep running into. The constraint has been moved out of the objective and into the shape of the model or the behaviour of the optimizer, where it cannot be inspected, weighted, or ported to a new setting.

Regularization-term methods are the interesting branch, because they keep the constraint in the loss without paying the pointwise cost. Rather than raising energy at specific locations, they make a statement about the distribution of low energy as a whole, which is a global claim rather than a pointwise one, and that is why it can be cheap. It is also exactly the slot we identified at the end of Section 3.

Which tells us where this article is going. LeJEPA lives in the regularization branch, and it inherits the obvious question that the branch has never answered properly; if your term is supposed to measure something about the distribution of the representation, what exactly should it measure? Sparse coding says the representation should be sparse. VAEs say it should look like a chosen prior, which is closer to the mark than anything else on this list and we will come back to it. VICReg says its covariance should be the identity. All three are guesses.

The rest of this section walks through the guesses and where each one falls short, and then we go looking for a derivation instead.

4.2 Negative samples

The oldest answer, and in some ways still the most honest, is to simply tell the model that different things should not land in the same place. Contrastive methods like SimCLR and MoCo pull the two views of an image together while explicitly pushing every other image in the batch away, which turns the anti-collapse requirement into something you can write down; a collapsed encoder maps everything to one point, so if you penalize proximity between unrelated samples then collapse is directly and heavily punished.

The positive pair is pulled together while every other sample in the batch is pushed away, which is where the quadratic cost comes from.

Contrastive learning does something rather elegant on our equation, in that f and g are both hiding inside a single expression. Expand the InfoNCE loss and they separate cleanly:

L_InfoNCE = −log [ exp(sim(z, z⁺)/τ) / Σ_j exp(sim(z, z_j)/τ) ]

= − sim(z, z⁺)/τ + log Σ_j exp( sim(z, z_j)/τ )

where the first term is f(·), the fit, and the second is g(·), the constraint.

The numerator pulls the positive pair together and is the fit term; the denominator sums over every other sample in the batch and is the constraint. So contrastive methods do have a real g, they simply never separated it out and named it. The cost is visible right there in the sum, which runs over the whole batch for every sample and gives us O(B²K).

The belief being expressed here is that the representation should spread its mass out rather than concentrate it, and as beliefs go that is entirely reasonable. The trouble is what it costs to say it this way. Every sample has to be compared against every other sample in the batch, which is a quadratic operation, and worse, the quality of the signal depends on how many negatives you have available; too few and the constraint barely bites, so you end up needing enormous batches or an auxiliary memory bank of stored embeddings to make it work at all. You are also, in a subtle way, lying to the model, since two different photographs of the same breed of dog are treated as things that must be pushed apart when they arguably should not be. The constraint is real and statable, it is just expensive and slightly wrong.

4.3 Teacher–student networks with EMA

BYOL and DINO take an entirely different route, and this is the one I find hardest to justify from first principles. You keep two copies of the network, a student that receives gradients normally and a teacher whose weights are an exponential moving average of the student’s, and the student is trained to predict the teacher’s output. The teacher never receives gradients at all; it simply trails behind.

Two ways of avoiding collapse without writing anything into the objective. In both cases the constraint lives in the gradient flow rather than in the loss.

Written against our base form, the objective here is only the fit term, with the target produced by a second, slower parameter set:

L_BYOL = ‖ Pred(Enc_θ(x_v)) − Enc_ξ(x_v′) ‖², ξ_t = τ ξ_{t−1} + (1−τ) θ_t

g(·) ≡ 0

There is nothing at all in the constraint slot. The entire anti-collapse behaviour lives in that second expression, which is an update rule rather than a term in the objective, and τ is a schedule rather than a trade-off.

What belief about Z does this express? I genuinely struggle to answer that, and I do not think the answer exists in a clean form. The mechanism is temporal rather than distributional; because the teacher lags, the target the student chases is always slightly stale, so the collapsed solution is a fixed point the optimizer cannot walk into quickly enough to reach before other structure has formed. Collapse is avoided by making the descent path awkward rather than by making collapse expensive.

And it works, extremely well, which is exactly what makes it such a clean illustration of the argument from the previous section. Here is a mechanism carrying an enormous amount of the inductive bias in some of the best vision models we have, and if you ask what property of the representation it guarantees, nobody can tell you. There is no theorem. There is an EMA decay rate that needs to be scheduled by hand, and if you get the schedule wrong the whole thing falls over.

4.4 Stop-gradients

SimSiam pushed this further and showed, somewhat startlingly, that you can drop the EMA teacher entirely and keep the collapse-avoidance, provided you place a stop-gradient on one branch so that gradients only flow through the other.

The same situation, made starker:

g(·) ≡ 0, ∇̃_θ L ≠ ∇_θ L

The constraint is not so much absent from the objective as the objective is absent from the optimization; we descend a modified gradient for which no corresponding loss exists.

I want to flag what has happened here, because it is easy to skate past. A stop-gradient is not a term in the loss. It does not appear in the objective, it has no value, it cannot be traded off against anything. It is an instruction to the autograd engine to discard part of the true gradient, which means the quantity actually being descended is not the gradient of the loss you wrote down. Whatever implicit objective SimSiam is optimizing, it is not L_pred, and there is no expression for it anywhere.

In the framing from Section 3 this is not a badly-specified prior; it is not a prior at all. There is nothing in the P([θ, Z]) slot, and the constraint has instead been smuggled into the optimizer.

4.5 Whitening

Barlow Twins and W-MSE take the most transparent approach of the lot. If dimensional collapse means the embedding dimensions are redundant and confined to a subspace, then simply require them not to be redundant; drive the cross-correlation matrix between the two views’ embeddings towards the identity, so that each dimension carries its own information and none duplicates another.

Barlow Twins drives the cross-correlation matrix towards the identity. The object being optimized is K × K, which is the entire cost problem in one picture.

This is the first method that puts something explicit into the slot:

g_BT(Z) = Σ_i (1 − C_ii)² + λ Σ_{i≠j} C_ij²

where C = cross-corr( Zᵛ, Zᵛ′ ) is a K × K matrix.

A genuine, differentiable, statable term at last. The K × K sitting in the shape of C is exactly the cost problem.

This is a real belief about Z, stated in the loss, differentiable, and defensible. The problem is what it costs and how it behaves. Building a correlation matrix over a K-dimensional embedding is a K×K operation, so both memory and compute grow quadratically in the embedding dimension, which is precisely the direction you want to scale. Several variants additionally need matrix inverses or square roots to perform the whitening, and those are numerically delicate in a way that shows up as instability at exactly the moment your model gets large enough to be interesting.

4.6 VICReg, and the thing that is actually wrong

VICReg deserves the most attention here, both because it is the most explicit of the lot and because it will reappear later as a special case of what LeJEPA does.

Its loss has three named terms; an invariance term that pulls the two views together, a variance term that forces each embedding dimension to maintain a minimum standard deviation across the batch, and a covariance term that pushes off-diagonal covariances towards zero. Put together, the second and third are a direct instruction about the shape of the embedding distribution; every dimension should carry variance, and no two dimensions should be correlated. Written in the language of Section 3, VICReg says the representation should have roughly identity covariance.

That is a genuine, written-down belief about Z, which already puts it ahead of everything above. And it still is not enough, for a reason that has nothing to do with cost.

In our notation its constraint term is the most explicit of the lot:

g_VIC(Z) = Σ_k max( 0, γ − σ_k(Z) ) + Σ_{i≠j} [ Cov(Z) ]²_ij

the first sum being the variance term and the second the covariance term.

Identity covariance constrains the first two moments of a distribution and says absolutely nothing about the rest of it.

And now the failure can be stated precisely rather than merely argued, because both terms depend on Z only through Cov(Z), which means:

Cov(P) = Cov(Q) ⟹ g_VIC(P) = g_VIC(Q)

That single line is the whole problem. Any two distributions sharing a covariance matrix are completely indistinguishable to this penalty, no matter how different they are in every other respect.

This is what “under-specified” means, and it is worth being concrete, because the phrase gets used loosely. A criterion is under-specified when a solution can satisfy it completely while remaining degenerate in a way the criterion cannot see. And there are many distributions with mean zero and identity covariance:

1. A standard Gaussian has identity covariance.

2. So does a distribution concentrated on the corners of a hypercube; all mass sitting at ±1 in each coordinate, nothing anywhere in between.

3. So does a heavy-tailed distribution where most embeddings sit near the origin and a handful sit very far away, arranged so the second moments work out.

1. So does a mixture of a few tight, widely separated clusters.

Every one of those passes VICReg’s check with a perfect score. Several of them are terrible representations. A distribution living on hypercube corners has a discrete, spiky structure that any downstream linear probe will struggle with; a heavy-tailed one has most of its samples crowded into a region where they cannot be told apart while its variance budget is spent on outliers. The covariance matrix cannot distinguish any of these from a Gaussian, because covariance is blind to everything above the second moment.

So the criterion is satisfied and the model is still bad, and crucially, nothing in your loss curve will indicate this. You have not been told you are wrong, honestly you have been told nothing at all about the thing you failed to constrain.

Four distributions with identical mean and covariance, computed rather than asserted. Three of them are poor representations, and no covariance-based penalty can tell any of them apart from the first.

4.7 The pattern

Setting all five side by side, and specifically setting their g terms side by side, the failures sort themselves into three kinds, and most methods manage to exhibit more than one. Two of the g terms are identically zero, one is O(B²) in the batch, one is O(K²) in the embedding dimension, and the only cheap and well-specified-looking one turns out to be blind to everything above the second moment.

The five mechanisms, their constraint terms, and the three ways of falling short.

The three failure kinds are worth naming, because they are exactly what we will demand the fix avoid:

1. Quadratic in something you want to scale. Batch size for contrastive methods, embedding dimension for whitening. Both are dimensions you would like to grow, and both are where the cost lands.

2. Under-specified. The criterion can be fully satisfied by a degenerate solution, which is VICReg’s problem and, in a looser sense, everybody’s.

2. Not in the loss. Stop-gradients and EMA schedules are interventions in the training procedure rather than statements in the objective, which means they cannot be reasoned about, traded off, or ported to a new setting without re-tuning from scratch.

Notice that none of these is a complaint about performance. All of these methods work; several of them produce state-of-the-art representations, and I have written admiringly about models built on them. The complaint is that we cannot say what any of them guarantees, which means we cannot improve them except by trying things, and we cannot know whether we are near the best available answer or nowhere close.

Which finally lets us ask the right question. Rather than continuing to patch, what would we actually want from a term in that slot if we could specify it freely? That is the next section.

5. The Ideal Case

We have spent three sections building a frame and then finding that nothing quite fits in it. Section 3 gave us a slot in the objective and told us what kind of object belongs there; Section 4 went through everything currently occupying that slot and found each occupant wanting. Which puts us in the position I always find most useful at this point in a paper, namely to stop reading it entirely and ask what we would want if we could simply specify it.

The nice thing is that Section 4 has already done most of the work, because each failure inverts directly into a requirement. So let us collect those first, then add the one requirement that the failures cannot possibly tell us about, which turns out to matter more than the rest combined.

5.1 What the failures demand

One. The constraint must actually exist. It sounds absurd to have to say this, and yet two of the five mechanisms we looked at had g ≡ 0. Whatever we end up with has to be a term in the loss with a value you can print, a gradient you can inspect, and a coefficient you can turn. Not a schedule, not a second network, not an instruction to the autograd engine.

Two. It must be fully specified. This is the VICReg lesson and the most important of the lot. We saw the failure stated exactly; Cov(P) = Cov(Q) implies g(P) = g(Q), which means any two distributions agreeing on their second moments are invisible to each other under the penalty. What we want is the opposite property, that g(P) = g(Q) should force P and Q to genuinely be the same distribution rather than merely agreeing on a couple of summary statistics. Anything constraining finitely many moments fails this by construction.

Three. It must be linear in both batch size and embedding dimension. Contrastive methods are O(B²K) and whitening is O(BK²), and those are quadratic in precisely the two directions you would like to grow. Any constraint requiring all pairs of samples, or a matrix over all pairs of dimensions, hits a wall exactly when the model becomes interesting.

Four. One knob. This follows from requirement one rather than standing alone; once the constraint has a form, a single coefficient in front of it is available, and the trade-off becomes a number instead of a tuning ritual. The difference between one hyperparameter and six is not really about convenience, it is about whether you can move the method to a new domain without redoing the search from scratch.

Five. Bounded, well-behaved gradients. Stability should come from the shape of the objective rather than from discovering the schedule that avoids the instability. A constraint whose gradients blow up when embeddings wander somewhere unusual will accumulate warmups and clipping and decay curves until you are back where you started.

5.2 A requirement the failures cannot give us

Those five are all defensive. Every one is derived from something that went wrong, which means collectively they describe what a constraint must avoid rather than what it should be. Satisfy only those and you get a cheap, stable, fully-specified penalty that pins the representation to some particular shape, with nothing whatsoever telling you which shape.

Recall too where the energy framing left us at the end of 4.1. Both families exist to make energy high away from the data, and we concluded that LeJEPA belongs in the regularization branch, which limits the volume of low-energy space rather than pushing up pointwise. But “limit the volume” is a statement about how much, not about where or in what shape, and sparse coding, VAEs and VICReg are three different guesses at the missing part. None of the five requirements above can adjudicate between them, because all three are perfectly capable of being cheap, stable and in the loss.

So here is the sixth requirement, and it is the one that actually generates an answer rather than filtering candidates:

Six. The representation should be optimal for downstream tasks we have not seen yet.

This is easy to nod at and harder to take seriously. The entire premise of self-supervised pretraining is that you do not know what the model will be used for. You train on unlabelled data, freeze the encoder, and later somebody attaches a linear probe for a classification problem you never anticipated, or runs k-NN over the embeddings, or fine-tunes on a small labelled set from a domain you have never seen. At the moment of pretraining the downstream label y does not exist. It is not merely unknown to us, it has not been chosen by anyone.

Which makes the question we are really asking a slightly strange one, worth stating carefully:

which distribution should Z follow to minimize expected risk,

when the task is drawn from a family we cannot observe?

Notice what has happened over the course of three sections. We started with a defensive question, how do we stop the model from cheating, and by taking the requirements seriously we have arrived at an entirely different one, which is what shape a representation ought to be. The first question has many acceptable answers, which is precisely why the literature contains so many mechanisms. The second question, remarkably, has exactly one.

5.3 What we are really asking for

Putting all six together, the thing we want in that slot is not a penalty at all. It is a target distribution.

That reframing is what the whole article has been building towards, so let me state it plainly. If the constraint must fully specify the distribution rather than a handful of its moments, then it cannot be a list of desirable properties like “high variance” or “low correlation”. It has to name a distribution Q and measure how far the induced representation distribution p_z sits from it:

g(θ, Z) = D( p_z ‖ Q )

for some target Q and some notion of distance D. Requirement six then tells us how to choose Q, since it should be whichever distribution minimizes expected risk over unseen tasks. Everything else on the list becomes a constraint on D; it must be linear in K and B, differentiable with bounded gradients, and it has to genuinely distinguish distributions rather than just their moments.

Which splits the problem cleanly in two, and that split is the structure of the rest of this article:

1. Which Q? A question about statistics and downstream risk, and the answer is provable rather than chosen. This is Section 6.

2. Which D? A question about how you compare distributions in high dimension without the cost exploding or the gradients misbehaving. This is Section 7.

It is worth flagging how unusual this separation is. In most papers the objective and its implementation are tangled together and you cannot discuss one without the other. Here the target and the divergence are genuinely independent design decisions; you could keep LeJEPA’s Q and swap its D, or the reverse, and both halves would still make sense. That separability is usually a sign the framing is correct.

It also explains why VAEs came closest of anything in Section 4. A VAE with a KL term to a Gaussian prior has exactly this shape, a target distribution and a divergence to it. What it lacks is any argument that its particular Q is the right one, since the Gaussian there is chosen for the reparameterization trick and analytic convenience rather than derived from downstream risk, and its D is a KL that needs the latent to be sampled rather than deterministic. The form was right. The justification was missing.

5.4 A note on what we are giving up

Before we go looking for Q, it is worth being honest that requirement six carries a real cost, and the paper is upfront about it.

Optimizing for an unknown task means optimizing for the worst case over a family of tasks, and a worst-case optimum is by construction not the best answer to any particular one. If you happen to know your downstream application is medical imaging, or document retrieval, or one specific classification problem, then a representation shaped for that task will beat a task-agnostic one. We are deliberately choosing not to bet, and the price of never betting is that you never win big on any single hand.

Whether that trade is worthwhile depends entirely on what you are building. For a foundation model, whose whole premise is serving tasks that do not exist yet, it is obviously right. For a system with one known downstream use it may not be, and I will come back to this in the Thoughts section, because I do not think it is as settled as the paper implies.

With that noted, let us go and find Q.

6. Methodolgy : Finding Q (why the isotropic Gaussian?)

This is the section the whole article has been walking towards, so we are going to take it slowly and derive rather than assert. The claim we are chasing is that there is exactly one distribution the embeddings should follow, and it turns out to fall out of two separate arguments that answer two genuinely different questions. People conflate them constantly, so it is worth naming them up front.

The first argument asks about the shape of the covariance and answers isotropic. The second asks about everything above the second moment and answers Gaussian. Neither one gets you there alone, and the gap between them is precisely where VICReg lives.

6.1 Setting up a question that can actually be answered

We cannot minimize risk on a task we have not seen. That is not a technical obstacle, it is a logical one; there is no y, so there is no risk to minimize. What we can do is minimize risk over a whole family of possible tasks, and specifically over the worst member of that family, which is what requirement six was really asking for.

Before we can compare candidate distributions we need to fix what is held constant, otherwise the question is ill-posed. Consider what happens if we do not: take any embedding distribution and multiply every vector by a thousand. Every downstream probe becomes trivially better conditioned, the effective noise shrinks relative to the signal, and we have learnt nothing except that bigger numbers are easier to fit. So we constrain the total variance,

tr( Cov(Z) ) = Σ_k λ_k = κ

and hold κ fixed throughout. Now the question becomes sharp and non-trivial. We have a fixed budget of variance, and the only question is how to distribute it. Every candidate distribution below has the same total energy; they differ only in geometry.

Notice this is already an interesting reframing. We are not asking how much information the representation should carry, we are asking how it should be arranged.

6.2 Linear probing gives us isotropy

Start with the most common way anybody evaluates a frozen encoder, which is to fit a linear probe on top of it. The probe solves ridge regression:

β̂ = argmin over β of ‖ y − Zβ ‖² + λ ‖β‖²

I want to pause on this equation for a second, because we have met it before. This is exactly the objective we built from scratch in the regularization article, fit term plus L2 penalty, with λ the Lagrange multiplier that emerged when a hard constraint on the weights got softened into a penalty. That article was about what regularization does to a model. This section is about what it does to your evaluation, and the two turn out to be the same mathematics pointed in different directions.

[EMBED: Regularization : What? Why? and How? (Part -1)]

The solution is standard:

β̂ = ( ZᵀZ + λI )⁻¹ Zᵀ y

Now, why are we looking at this at all? Because β̂ is an estimator, and estimators have bias and variance, and both of those turn out to depend on the geometry of Z in a way we can compute. If we can show that anisotropy hurts both, we are done with the first half of the argument.

Step 1: what the bias looks like

Assume the task is genuinely linear in the embedding, so y = Zβ* + ε with E[ε] = 0 and β* the true parameters, and take Z to be mean-centred. Then

E[β̂] = ( ZᵀZ + λI )⁻¹ ZᵀZ β*

and the bias is the difference between that and the truth:

Bias = [ ( ZᵀZ + λI )⁻¹ ZᵀZ − I ] β*

This is opaque as written, so let us do the thing that always helps with expressions full of matrix inverses, which is to move into the eigenbasis. Write the empirical second moment as ZᵀZ ≈ N·Σ where Σ = Cov(Z), and eigendecompose Σ = UΛUᵀ with eigenvalues λ₁ … λ_K. In that basis everything becomes diagonal and the whole expression collapses to a per-coordinate statement:

Bias_k = [ Nλ_k / (Nλ_k + λ) − 1 ] β*_k = − s_k · β*_k , s_k = λ / (Nλ_k + λ)

This is the whole thing, so let us read it carefully. Each coordinate of the estimate is shrunk towards zero by a factor s_k that depends only on that coordinate’s eigenvalue. When λ_k is large the denominator is large, s_k ≈ 0, and that coordinate is recovered faithfully. When λ_k is small s_k ≈ 1, and that coordinate is shrunk almost to nothing regardless of what the true β*_k was.

The shrinkage factor as a function of a direction’s eigenvalue. Directions carrying little variance are erased from the estimate entirely, whatever the task needed from them.

So the ridge penalty is not shrinking your estimate uniformly. It shrinks hardest along the directions where your representation has the least variance, which is a completely reasonable thing for it to do, since those are the directions where the data tells you least, and it is catastrophic if the task happens to depend on one of them.

Step 2: taking the worst case

Squaring and summing gives the total squared bias:

‖Bias‖² = Σ_k ( λ / (Nλ_k + λ) )² · β*²_k

Now we invoke requirement six. We do not know β*, and we are not allowed to assume anything about it, so we ask what the worst task of a given size looks like. Fix ‖β*‖ and let an adversary point it wherever it does most damage; obviously it points along the coordinate with the largest shrinkage, which is the one with the smallest eigenvalue:

max over ‖β*‖ = c of ‖Bias‖² = ( λ / (Nλ_min + λ) )² · c²

And now the optimization writes itself. Minimizing the worst-case bias means maximizing λ_min, subject to the eigenvalues summing to κ. There is no cleverness required here; if you have a fixed total to distribute and you want to make the smallest share as large as possible, you split it evenly:

max λ_min subject to Σ λ_k = κ ⟹ λ₁ = λ₂ = … = λ_K = κ/K

Which is isotropy. Not chosen, not preferred; forced, by the combination of a fixed variance budget and a refusal to assume anything about the downstream task.

Step 3: the variance says the same thing

Bias was one half of it. For the variance, set λ = 0 so the algebra stays clean, and use the standard result for OLS:

Var(β̂) = σ² ( ZᵀZ )⁻¹ = (σ²/N) Σ⁻¹

Total estimator variance is the trace, which in the eigenbasis is just a sum of reciprocals:

tr( Var(β̂) ) = (σ²/N) · Σ_k 1/λ_k

Minimize that subject to Σλ_k = κ. Lagrange gives

∂/∂λ_k [ Σ_j 1/λ_j + μ Σ_j λ_j ] = − 1/λ_k² + μ = 0 ⟹ λ_k = 1/√μ for every k

which is the same value for every k, so once again all eigenvalues equal and once again isotropy.

There is something worth staring at in that sum of reciprocals, though, because it is doing more than delivering the same answer twice. As any λ_k → 0, the sum diverges. Not degrades, diverges. And a λ_k heading towards zero is precisely the dimensional collapse we defined back in Section 2. So this expression is telling us that dimensional collapse does not merely waste capacity, it sends the variance of every downstream estimator to infinity. The failure mode we started the article worrying about and the geometry we are now deriving are the same phenomenon seen from two angles.

6.3 Why isotropy is not enough

We now have a result, and it is a real one; the covariance should be a scaled identity. It is also, unfortunately, exactly the result VICReg already had.

Look back at the four distributions from Section 4.6. All four have identity covariance, so all four are perfectly isotropic, so all four are equally optimal by everything we have derived so far. Three of them are terrible representations.

The reason our argument cannot see the difference is structural rather than accidental. A linear probe interacts with the embedding only through ZᵀZ, which is to say only through the second moment. Anything the second moment cannot see, a linear probe cannot see either, and so no amount of care in the linear analysis will ever tell us about the shape of the distribution beyond its covariance.

Hence if we want to constrain the rest of the distribution, we need a probe that actually depends on the rest of the distribution. Which is not a contrivance, incidentally; k-NN and kernel probes are standard ways to evaluate frozen encoders, arguably more common than linear probes in retrieval settings. So let us redo the analysis with one of those and see what falls out.

6.4 Nonlinear probing gives us Gaussian

Take a radius-based k-NN probe, which predicts by averaging the labels of everything inside a ball:

ŷ(q) = ( 1 / |N_r₀(q)| ) · Σ_{n ∈ N_r₀(q)} y_n , N_r₀(q) = { n : ‖z_n − q‖ ≤ r₀ }

The key difference from before: this estimator depends on which points happen to be nearby, and that is a statement about the local density of the embedding distribution rather than about its global second moment. That is exactly the sensitivity we were missing.

Step 1: where the bias comes from

Let g(z) = E[y | z] be the true regression function. The bias of the probe at a query point q is

Bias(q) = E[ g(z) | z ∈ N_r₀(q) ] − g(q)

that is, the difference between the average of g over the neighbourhood and its value at the centre. Taylor expand g about q:

g(z) = g(q) + ∇g(q)ᵀ (z−q) + ½ (z−q)ᵀ H (z−q) + …

Now here is the crucial observation, and it is worth thinking about before we compute anything. If the points in the neighbourhood were distributed symmetrically about q, the linear term would average to exactly zero. For every point displaced in one direction there would be a matching point displaced the other way, and ∇g would contribute nothing to the bias. The neighbourhood average would be unbiased to first order and only the curvature term would survive.

So the linear term only survives when the neighbourhood is lopsided, that is, when there are systematically more points on one side of q than the other. And a systematic imbalance of points across a small ball is exactly what it means for the density to have a gradient there.

Left, a flat density: the neighbours sit symmetrically about the query and their displacements cancel. Right, a density gradient: the ball fills unevenly and the average is pulled towards the dense side.

Step 2: making that precise

Expand the density around q in the same way, using ∇p = p · ∇log p:

p(z) ≈ p(q) · ( 1 + ∇log p(q)ᵀ (z−q) + … )

Taking the density-weighted expectation of the displacement over the ball, and using the standard fact that for a uniform ball of radius r₀ in R^K we have E[(z−q)(z−q)ᵀ] = r₀²/(K+2) · I:

E[ (z−q) | z ∈ N_r₀(q) ] ≈ ( r₀² / (K+2) ) · ∇log p(q)

So the displacement does not average away. It averages to something proportional to the log-density gradient, and the bias picks up a term

Bias(q) ≈ ( r₀² / (K+2) ) · ∇log p(q)ᵀ ∇g(q) + ( a second O(r₀²) term from the curvature H )

Read that plainly: the probe is biased in proportion to how fast the density is changing. In a region where embeddings pile up on one side, the neighbourhood average is pulled towards that side, and the prediction inherits the pull.

Step 3: integrating over all query points and all tasks

Square the bias and integrate over the distribution of query points. The ∇g factor depends on the task, which we do not know, so we treat it as an unknown of bounded size and absorb it into a constant τ_g. What is left is:

ISB ∝ ( r₀⁴ / (K+2)² ) · τ_g² · J(p) , J(p) = ∫ ‖∇log p(z)‖² p(z) dz

That integral has a name. J(p) is the Fisher information of the distribution with respect to location, and it measures how sharply the density varies relative to its own magnitude. A density with steep gradients, spikes, sharp edges or isolated clusters has large J(p). A smooth, gently-varying density has small J(p).

Everything else in that expression is fixed by the setup; the radius, the dimension, the unknown task. The only thing we control is J(p), and it is the only thing left to minimize.

Step 4: which distribution minimizes Fisher information?

We now have a clean, self-contained mathematical question with nothing to do with machine learning:

min over p of J(p) subject to Cov(p) = Σ

And it has a classical answer. For any distribution with covariance Σ,

J(p) ≥ tr( Σ⁻¹ ) , with equality if and only if p = N(0, Σ)

The Gaussian is the unique minimizer. Not the best among the ones we tried, not optimal under extra assumptions; unique, given the covariance constraint.

The four distributions from Section 4.6, evaluated by numerical integration. Covariance could not tell them apart at all; Fisher information separates them by nearly three orders of magnitude, and only the Gaussian attains the bound.

Combine that with 6.2, which told us Σ must be a scaled identity, and the answer is fully pinned down:

Q = N( 0, σ² I_K )

Every degree of freedom in the choice of target distribution has now been used up. The shape of the covariance came from the linear probe, and everything above the second moment came from the nonlinear one.

6.5 Why this is the right answer from two directions

I find a result more convincing when a second, unrelated argument arrives at the same place, and here one does.

Ask a different question: given a fixed variance budget, which distribution assumes the least about anything? That is the maximum entropy question, and its answer under a covariance constraint is famously the Gaussian. So the distribution that minimizes Fisher information and the distribution that maximizes entropy, under the same constraint, are the same distribution.

That coincidence is not a coincidence, and it has a clean interpretation. Both quantities are measuring structure. Entropy asks how much structure the distribution has, Fisher information asks how sharply that structure varies, and under a variance budget the flattest, least committed, most spread-out option optimizes both.

Which finally gives us the intuition to carry away from this whole section. You are building a searchlight before you know where you will need to look. Any asymmetry in the beam is a bet about where the target will be. Isotropy is the decision not to bet on direction; Gaussianity is the decision not to bet on location within a direction, by refusing to concentrate mass anywhere in particular. The two theorems in this section are the two halves of declining to gamble, and requirement six was the statement that we are not in a position to gamble at all.

Three ways to spend a fixed illumination budget. The first two are excellent if the tasks land where you guessed and useless otherwise; only the third makes no assumption at all.

It also retroactively explains the four distributions from Section 4.6. The hypercube corners and the cluster mixture both make an enormous bet, namely that the downstream task cares about exactly those four regions and nothing in between; if it does they are excellent, and if it does not they are useless. The heavy-tailed distribution bets everything on the centre. Only the Gaussian bets nothing, and that is precisely why it wins the worst case.

6.6 Where this leaves the objective

We came into this section with a target-shaped hole:

g(θ, Z) = D( p_z ‖ Q )

and we can now fill in half of it:

g(θ, Z) = D( p_z ‖ N(0, σ² I) )

which, plugged back into the objective from Section 3, gives us the shape of LeJEPA:

L(θ, Z ; X) = L_pred(θ ; X) + α · D( p_z ‖ N(0, σ² I) )

One thing remains, and it is not a small thing. We have to actually compute D, which means measuring the distance between an empirical distribution of a few hundred embedding vectors and a Gaussian, in a space of a thousand dimensions, differentiably, cheaply, and with bounded gradients. That is a genuinely hard problem and it is where most of the paper’s engineering lives.

That is Section 7.

7. Methodology : Finding D (measuring the distance to a Gaussian)

We know what we want the embedding distribution to be. Now we have to write down something differentiable that says how far we are from it, and this turns out to be considerably harder than it sounds.

Here is the problem in its concrete form. Your batch gives you maybe 256 embedding vectors sitting in 1024 dimensions. From that you have to produce a single scalar saying how un-Gaussian this cloud of points is, differentiate it, and do so cheaply enough to run on every step of training. Anyone who has tried to estimate a density in a thousand dimensions from a few hundred samples will already be uncomfortable, and they should be; that problem is genuinely hopeless. So the interesting question is how the paper avoids having to solve it.

7.1 Distribution matching as a hypothesis test

The usual instinct here is to reach for a divergence. KL, Wasserstein, MMD, pick one, minimize it. The paper does something different, and the reframing is the first good idea in this section.

Rather than asking how far apart two distributions are, ask a statistician’s question:

H₀ : p_z = Q versus H₁ : p_z ≠ Q

This is a hypothesis test. You compute a test statistic T from your samples, which is a single number summarizing how much evidence the data provides against the null, and if it exceeds a critical value you reject. Statisticians have been designing these for a century and they come with properties that divergence estimators generally do not; known null distributions, provable consistency, and a clear notion of what the statistic is measuring.

Then comes the move that makes it useful to us. The test statistic is already a scalar function of the embeddings, so just use it as the loss.

g(θ, Z) = T( { f_θ(x_n) } )

Minimizing T is minimizing the evidence against the null, which is to say training the encoder until a statistician could no longer distinguish its embeddings from Gaussian samples. That is a much more precise statement of intent than “penalize the covariance”, and it is exactly the “fully specified” requirement from Section 5 restated in a form you can compute.

It also gives us something no other method in Section 4 could offer, which is a scale for the number. A test statistic can be compared against a critical value, so you can ask whether your representation is acceptably Gaussian at some significance level rather than just noting the penalty went down.

7.2 The obstacle: high dimension

Unfortunately, multivariate normality tests are not cheap. Mardia’s test needs third and fourth moment tensors, Henze–Zirkler needs all pairwise distances, and essentially every classical option is at least quadratic in sample count and often far worse in dimension. Worse than the cost is the sample complexity; to say anything reliable about the shape of a density in K dimensions you need a number of samples that grows exponentially in K. With 256 samples in 1024 dimensions, the empirical distribution is not an approximation to p_z, it is 256 isolated points in a space so large that they are all effectively equidistant from one another.

Estimating a density directly needs sample counts that grow exponentially with dimension; a univariate test on each slice needs a fixed number regardless of K.

So a direct test is not merely expensive, it is statistically meaningless at the batch sizes we can afford. We need a way of not looking at the high-dimensional density at all.

7.3 Sketching: reading a distribution off its shadows

The escape is a theorem from 1936, and it is the most elegant thing in the paper.

Cramér–Wold. For two random vectors X, Y in R^K:

⟨a, X⟩ ≜ᵈ ⟨a, Y⟩ for every unit vector a ⟺ X ≜ᵈ Y

In words: if every one-dimensional projection of two distributions matches, the distributions themselves match. Nothing is lost by only ever looking at shadows, provided you look from every angle. The paper proves a hyperspherical variant of this (Lemma 3) restricted to unit-norm directions, which is what we actually sample.

Take a moment with how strong that is. A distribution over a thousand dimensions is an enormously complicated object, and this says it is completely determined by a family of one-dimensional distributions, each of which is a thing you can plot on a line and reason about with century-old statistics.

The procedure follows immediately. Sample M unit directions, project the batch onto each, run a cheap univariate normality test per direction, and aggregate. Theorem 2 in the paper establishes that taking the worst direction gives a valid test, with the right level under the null and power tending to one under the alternative:

T_A(Z) = max over a ∈ A of T( { aᵀ z_n } )

And then the paper immediately does something else. The actual SIGReg definition averages over directions rather than maximizing:

SIGReg(A, Z) = (1 / |A|) · Σ_{a ∈ A} T( { aᵀ z_n } )

The stated reason is worth dwelling on, because it is a good illustration of the difference between a statistic and a loss. The maximum is the correct test, since you reject the null if any direction provides evidence against it. But as an objective it produces a sparse gradient; only the single worst direction receives any signal on each step, and the other M−1 projections contribute nothing. Averaging spreads the gradient across every direction at once, and since we are optimizing rather than deciding, that is the behaviour we want. The theory justifies the max; the implementation uses the mean.

Why this beats the curse of dimensionality

This deserves care, because the argument is not the naive one. You might hope that since each univariate test needs a sample size independent of K, everything is fine. But the real question is different: if we only constrain M directions, what happens along all the directions we did not sample?

The paper answers with smoothness. If the embedding density lies in a Sobolev space H^α, then satisfying the test exactly on M directions bounds the expected discrepancy over all directions:

E_a ∫ | φ_a(t) − φ_N(t) |² dt ≤ C(K, α) · |A|^(−2α / (K−1)) · (‖ · ‖_{H^α})

The exponent is what matters. The bound decays as |A| grows, and it decays faster when α is larger, meaning smoother embedding densities are constrained more effectively by the same number of directions. The paper’s conclusion is that |A| = O(K) directions suffice for an ε-approximation when α is large, and it argues that deep networks produce naturally smooth densities, both from the architecture and from the implicit and explicit regularizers already in play. Since the target itself is a smooth isotropic Gaussian, α grows quickly during training, which makes the bound tighten as you converge.

There is a second argument, and it is the one I find more persuasive because it is almost free. SGD resamples the directions at every step. So while any single minibatch constrains only M directions, the cumulative number of directions applied over training grows linearly with the number of steps. The paper reports that |A| as small as 16, resampled every step, comfortably outperforms a fixed set of thousands held constant throughout training.

So the curse is beaten twice over; once by smoothness, which says a modest M constrains the whole sphere, and once by resampling, which says the effective M is the number you draw times the number of steps you take.

A fortunate interaction with Section 6

Something quietly convenient happens here. Slicing requires knowing the projected target distribution for each direction a. For a general target that is its own computation, potentially different for every direction. But our target is the isotropic Gaussian, and for z ~ N(0, σ²I) with ‖a‖ = 1:

aᵀ z ~ N(0, σ²) for every a

Every projection of an isotropic Gaussian is the same univariate Gaussian.

So there is nothing to compute per direction; the target is one fixed 1-D distribution and every one of the M tests compares against it. Back in Section 5 I claimed Q and D were independent design decisions. They are, in the sense that you could swap either, but they are not indifferent to each other, and the Q that the risk argument forced on us happens to be the one that makes slicing essentially free.

7.4 Choosing the univariate test

We have reduced everything to: given N real numbers, how Gaussian are they? The paper works through three families and rules out two of them.

Option one: moments

Take the Jarque–Bera test, which compares skewness and kurtosis, and extend it with the first two moments to get a standard-Gaussian test rather than a Gaussian-of-any-mean-and-variance test. This is moment matching over the first four moments, and moment matching is a respectable tool.

Two problems, and the paper formalizes both. The first is identifiability. Theorem 3 states that minimizing a weighted sum of squared differences over any finite number of moments does not imply the distributions are equal:

Σ_{k=1..K} c_k ( m_k(P) − m_k(Q) )² = 0 does not imply P = Q

Which is precisely VICReg’s failure from Section 4, generalized. The obvious response is to use more moments. Which runs into the second problem: gradient norms scale as O(k) and the variance of Monte Carlo gradient estimates grows as O(k² m_{2(k−1)}) for the k-th moment, because moments are polynomial in the data and raising values to a power of k explodes quickly. The paper’s phrasing is that stability and identifiability cannot be achieved simultaneously with moments, and that is the whole objection.

Option two: the empirical CDF

Cramér–von Mises, Anderson–Darling, Watson. These compare the empirical CDF against the target and are properly consistent, so unlike moments they are not under-specified. The problem is mechanical. All of them require sorting, and sorting breaks the embarrassingly parallel nature of SGD, particularly across multiple GPUs where synchronization is needed. Sorting and order statistics are also non-differentiable, so gradient-based optimization needs a relaxation, and every available relaxation introduces its own hyper-parameters, which defeats the entire motivation. The paper also excludes Kolmogorov–Smirnov specifically because its ℓ∞ norm produces sparse gradients, and reports Shapiro–Wilk as unstable in practice.

Option three: characteristic functions

The empirical characteristic function is the Fourier transform of the density, estimated as a simple average:

φ̂_X(t) = (1/n) Σ_j e^{i t X_j}

Three properties make this the right choice. It is identifiable, since the map from distribution to characteristic function is a bijection, so a test built on φ encodes the whole distribution rather than a summary. It is differentiable, with no sorting anywhere. And it is distributed-friendly, because an average over samples composes across devices via a single all-reduce.

The Epps–Pulley test compares the two in weighted ℓ² norm:

EP = N ∫ | φ̂_X(t) − φ(t) |² w(t) dt

with the weight function typically Gaussian, w(t) = e^{−t²/σ²}, and σ commonly set to 1.

The boundedness is where this becomes decisive, and the paper proves it as Theorem 4:

| ∂EP / ∂z_i | ≤ 4σ² / N , | ∂²EP / ∂z_i² | ≤ C √π σ³ / (2N)

Both the gradient and the curvature are bounded by constants depending only on the batch size and the weight bandwidth σ, and on nothing about the data. Compare that to moments, where the gradient of the k-th moment grows like |x|^{k−1} and a single drifting embedding can dominate everything. Chain rule then gives a stable bound on the gradient with respect to θ. Stability here is a property of the function rather than something tuned in, which is requirement five from Section 5 met by construction.

One honest detail the paper is upfront about. Using minibatches introduces a bias in the estimator, quantified in Theorem 6:

E[ L̂_n(θ) ] = L(θ) + (1/N) ∫ w_s(t) ( 1 − |φ_P(t)|² ) dt

So both the loss and its gradient carry an explicit O(1/N) bias. The authors report this as minimal in practice even at batch sizes as small as 16, and note that U-statistic debiasing or sample splitting would remove it, neither of which they pursue.

7.5 Putting SIGReg together

Assembling everything, SIGReg draws M unit directions, projects the batch onto each, computes the Epps–Pulley statistic of each projection against N(0, 1), and averages:

SIGReg(A, Z) = (1 / |A|) Σ_{a ∈ A} EP( { aᵀ z_n } )

The integral inside EP is estimated by quadrature. The paper finds the simple trapezoidal rule sufficient with as few as 17 knots, and exploits the symmetry of the integrand to double the effective number of knots for free.

Check it against the Section 5 requirements. It is a term in the loss with a value and a gradient, so requirement one holds. It is built on characteristic functions, which are identifiable, so requirement two holds. Its cost is linear in minibatch size, so requirement three holds. It sits behind a single coefficient, so requirement four holds. Its gradients and curvature are bounded by Theorem 4, so requirement five holds. And its target came from the risk argument in Section 6, so requirement six holds.

Six for six, which is the first time anything in this article has managed that.

Why it distributes so well

The only quantity crossing between devices is the empirical characteristic function, which is an average over samples. Averages compose; each device computes its partial mean and one all-reduce combines them. Nothing depends on gathering the embeddings themselves. Contrastive methods, by contrast, need negatives from across the batch and therefore must move embeddings between devices, which is a far larger and more awkward transfer.

VICReg falls out as a special case

A satisfying way to close, and the mechanism is simpler than I expected. It is not an approximation or a limit of the Epps–Pulley statistic; it is a different choice of T. Substitute

T( { x_n } ) = mean( { x_n } )² + ( std( { x_n } ) − 1 )²

into SIGReg, keep everything else, and in the limit of many slices you recover VICReg. The paper proves in Appendix B.14 that this enforces exactly E[Z] = 0 and Cov(Z) = I; the argument is a neat one, showing that if every unit direction has zero mean and unit variance in projection then the mean vector must vanish and the covariance must be the identity, by testing the standard basis vectors and then their pairwise sums.

So VICReg is LeJEPA with a degenerate statistical test, and by Theorem 3 that test is exactly the kind that admits shortcut solutions. The paper is blunt about it, advocating strongly against the setting for that reason.

7.6 What actually gets trained

Before moving on, it is worth saying plainly that the objective is the paper. Everything we have built over the last four sections, the joint prior, the isotropic Gaussian target, the sliced characteristic-function test, is the contribution. What follows is scaffolding, and the paper treats it that way too; there is no clever architecture here, no novel augmentation scheme, no training trick. That is rather the point. If the objective is right, the rest should be boring.

The encoder is arbitrary.

f_θ appears nowhere in the derivation except as the thing producing Z, so LeJEPA imposes no architectural requirement whatsoever. The paper pretrains roughly 50 architectures from 8 families, taking every model in timm under 20M parameters, and all of them learn useful representations. A method carrying architectural anti-collapse machinery inherits assumptions about the architecture; a constraint stated purely on the output distribution does not care what produced it.

No predictor.

This is stronger than “optional”. Prior work established that predictors in image JEPAs, absent genuinely asymmetric views, exist primarily to prevent collapse, and removing them collapses the encoder to chance level. With SIGReg in place the paper removes both the predictor and the teacher–student architecture without collapse, and recommends training without a predictor. A teacher–student setup still gives ViTs a small boost, but as stochastic weight averaging rather than as a collapse guard.

No register tokens.

Recent vision models need register tokens to avoid training instabilities. The paper’s position is that those instabilities come from poorly conditioned objectives, and LeJEPA is stable with or without them.

Views follow the DINO setup.

Eight views total, two global at 224×224 and six local at 96×96. Nothing exotic, and the configuration common in prior work transfers directly.

Training data.

ImageNet-10, ImageNet-100 and ImageNet-1k for the main runs, and the in-domain experiments pretrain directly on the target corpus, including Galaxy10 (11,000 samples, 10 classes), Food101, Flowers102, CIFAR-10 and CIFAR-100. Worth noticing what is absent: no curated billion-image collection and no data-filtering pipeline, which is a meaningful contrast given that recent progress in this area has come substantially from data curation rather than from objectives.

The recommended settings.

Unusually for this kind of article, I can just list them, because there are so few: λ = 0.05, two global and six local views, batch size at least 128, 1024 slices, 17 quadrature knots over an integration domain of [−5, 5], AdamW with standard linear warmup and cosine annealing on the learning rate and no scheduler on weight decay. The ablations show none of these choices is delicate; the paper’s own summary is that while some settings slightly improve performance, none of them leads to catastrophic collapse.

7.7 The complete objective

Two pieces remain to be written down properly, and the first is more interesting than it looks.

The prediction loss.

The natural way to write “all views should agree” is as an average over all pairs of views:

L_pred = (1/V_g) Σ_{v=1..V_g} (1/V) Σ_{v′=1..V} ‖ z_{n,v} − z_{n,v′} ‖²

which looks like it costs V_g × V comparisons. It does not. The paper shows (Appendix B.6) that this collapses exactly into a distance-to-the-mean:

L_pred = (1/V) Σ_{v′=1..V} ‖ μ_n − z_{n,v′} ‖² , μ_n = (1/V_g) Σ_{v=1..V_g} z_{n,v}

The algebra is a pleasant little identity; expanding the squared norms, the cross terms reorganize into the mean and the pairwise structure disappears. So every view predicts the average of the global views, and the cost is linear in the number of views rather than quadratic. There is no predictor network anywhere in this expression.

The total loss.

The two terms combine as a convex combination rather than an unbounded penalty:

L_LeJEPA = λ · (1/V) Σ_v SIGReg( { z_{n,v} } ) + (1 − λ) · (1/B) Σ_n L_pred( { z_{n,v} } )

with λ = 0.05 recommended as a robust default.

Note that SIGReg is applied per view and then averaged, rather than to all views pooled together. And note the shape of the combination: λ sits in [0, 1] and trades the two terms against each other directly, which is a slightly different object from the Lagrange multiplier we derived in Section 3. The Bayesian derivation still tells you why there are exactly two terms and what each one is; the convex form is the practical parameterization, and it has the pleasant property that λ is bounded and interpretable rather than needing a scale search.

That is the whole method. One predictive term, one constraint term, one hyperparameter. No teacher, no stop-gradient, no predictor, no schedule, no whitening, no negatives, no register tokens.

It is worth pausing on how much of that objective we derived rather than read. The form came from Bayes in Section 3, the target came from downstream risk in Section 6, and the divergence came from the requirements list in Section 5 plus a theorem from 1936. The only genuinely free choices left are λ, the number of slices, and the quadrature grid, and of those only λ meaningfully changes what the model learns.

Next, what this looks like as code.

8. SIGReg in code

The paper gives its implementation in full, and it is short enough to walk through line by line. What follows is Algorithm 1 from the paper, split into pieces, with Algorithm 2 for the full loss at the end.

8.1 The signature

def SIGReg(x, global_step, num_slices=256):
    # x is an (N, K) tensor of embeddings
    dev = dict(device=x.device)

Everything the constraint needs is in that first argument plus a step counter. No teacher weights, no momentum buffer, no queue of past embeddings, no temperature schedule. The global_step is not used for annealing anything; it is a random seed, and the next fragment explains why.

8.2 Directions

g = torch.Generator(**dev)
g.manual_seed(global_step)
proj_shape = (x.size(1), num_slices)
A = torch.randn(proj_shape, generator=g, **dev)
A /= A.norm(p=2, dim=0)

Gaussian samples, normalized to the unit sphere. That is the standard way to draw uniformly from a sphere, and it is genuinely just randn; the paper discusses low-discrepancy sequences in its appendix but notes they do not change the asymptotic bounds, and the shipped implementation does not use them.

The seeding is the detail worth noticing. Every device seeds its generator with the same global_step, so every device draws the same directions on the same step without any communication. And because the seed advances every step, the directions are resampled continuously, which is exactly the compounding-coverage effect from 7.3. One line buys both distributed consistency and the escape from the curse of dimensionality.

8.3 Projection

x_t = (x @ A).unsqueeze(2) * t              # (N, M, T)

The projection itself is x @ A, taking an (N, K) batch to (N, M). After this, K never appears again. Everything downstream is one-dimensional statistics on M columns. The unsqueeze and multiply broadcast each projected scalar against the integration grid, giving an (N, M, T) tensor which is the largest thing this function allocates, and which is linear in all three of its dimensions.

8.4 The integration grid

t = torch.linspace(-5, 5, 17, **dev)
exp_f = torch.exp(-0.5 * t**2)

Seventeen points spanning [−5, 5]. The exp_f term does double duty; it is both the characteristic function of the standard Gaussian, e^{−t²/2}, which is our target, and the Gaussian window used to weight the integrand. That the two coincide is a convenience of targeting a Gaussian.

Seventeen knots sounds coarse for a numerical integral. The ablation shows the choice barely matters, and the reason is instructive: the characteristic function is accurate near zero, so the moments of the distribution are well characterized even over a modest integration range.

8.5 The empirical characteristic function

ecf = (1j * x_t).exp().mean(0)
ecf = all_reduce(ecf, op=”AVG”)

The first line is the empirical characteristic function, computed as a complex exponential averaged over the batch, collapsing (N, M, T) to (M, T). Using complex tensors directly keeps the real and imaginary parts together, and both carry information; the real part sees deviations that are symmetric about the mean, the imaginary part sees asymmetric ones.

The second line is the entire distributed story. An empirical characteristic function is a mean over samples, and means compose across shards, so one all-reduce of an (M, T) array, a few thousand complex numbers, gives exactly the result you would have obtained on a single device. Nothing depends on N, K, or the device count. This is why the paper can claim fifty lines.

8.6 Comparing to the target

err = (ecf - exp_f).abs().square().mul(exp_f)
N = x.size(0) * world_size
T = torch.trapz(err, t, dim=1) * N
return T

The weighted squared discrepancy, integrated over frequency by the trapezoidal rule and scaled by the global batch size. Note that N is the total across devices, not the local batch, which matters because the Epps–Pulley statistic is defined with that scaling.

This is also where boundedness becomes visible as code. ecf is an average of complex exponentials, so its magnitude is at most one no matter what the embeddings do; an embedding sitting at 10⁶ contributes a point on the unit circle, exactly like every other embedding. Put a kurtosis penalty in the same position and that embedding contributes 10²⁴. No single input can dominate this statistic, which is Theorem 4 arriving as an operational fact.

One thing this function does not do is take a maximum. The return value is a per-direction vector, and the averaging happens in the caller. As we saw in 7.3, that is deliberate.

8.7 The full loss

def LeJEPA(global_views, all_views, lambd):
    g_emb = forward(torch.cat(global_views))
    a_emb = forward(torch.cat(all_views))
    centers = g_emb.view(-1, bs, K).mean(0)
    a_emb = a_emb.view(-1, bs, K)
    sim = (centers - a_emb).square().mean()
    sigreg = mean(SIGReg(emb, global_step) for emb in a_emb)

    return (1 - lambd) * sim + lambd * sigreg

Eight lines. The centers variable is μ from 7.7, the mean of the global-view embeddings, and sim is every view’s squared distance to it, which is the collapsed all-pairs loss. There is no predictor and no second network; forward is called twice on the same encoder. For non-ViT architectures the paper simply sets global_views = all_views and the local-view machinery disappears.

SIGReg is applied per view and averaged, and the two terms combine convexly through λ. That is the entire training step, minus the optimizer.

8.8 Complexity, totalled

The paper states the implementation is O(N) in both time and memory with respect to the minibatch size, and reports wall-clock timings. Reading the operations directly: the projection is O(NKM), the characteristic function is O(NMT), and the reduction is O(MT), giving

O( N . M ( K + T ) )

linear in batch size, linear in embedding dimension, linear in slices. To put numbers on it, at N = 512, M = 512 and 16 integration points the forward-backward pass costs about 0.47 ms on a V100, and holding N fixed while raising M from 512 to 8192 takes it only to about 0.67 ms. Against an encoder forward pass, this is a rounding error.

9. Results & Outcomes

A derivation can be elegant and still lose to a pile of heuristics that somebody tuned for two years, so what follows is the empirical case. The paper’s validation spans more than ten datasets and over sixty architectures at scales approaching two billion parameters. I have grouped the findings into five, roughly in order of how much each one surprised me.

9.1 Stability across hyper-parameters and architectures

LeJEPA trains stably across every hyper-parameter the authors varied and across roughly 50 architectures from 8 families, using a single value of λ.

The architecture sweep is the cleanest demonstration. The authors took every model in the timm library under 20M parameters, 50 of them spanning ConvNeXt, EfficientNet, Inception, LeViT, MaxViT, MaxxViT, ResNet and ViT, pretrained each on ImageNet-10 with the same recipe, and evaluated with a frozen backbone. All 50 land between 91.5% and 95% top-1, which is a remarkably narrow band for a set of models that different. The paper’s practical note is that architectures which do well under supervision also do well here, so ResNets and ViTs are the sensible starting point over specialized designs like EfficientNet.

On hyper-parameters, the headline is that nothing is delicate. Varying λ across two orders of magnitude on ImageNet-100 with a ResNet-50 leaves performance stable, with the optimum drifting slightly upward as the number of views increases. On ImageNet-1k with a ViT-Large/14, the ablations cover the Epps–Pulley integration domain, the number of quadrature points, the number of slices, batch size, projector dimension and register tokens. The paper’s own summary is the right one to quote in spirit: while some settings slightly improve performance, none of the choices leads to a catastrophic collapse.

Two details worth pulling out. Competitive performance holds at batch sizes as small as 128 on ImageNet-1k, which suggests noticeably lower memory requirements than methods needing large batches for negatives. And the number of slices has only a modest effect, with 512 already competitive and 1024 recommended.

The removal results are the ones that most directly vindicate the theory. Prior work established that predictors in image JEPAs and teacher–student architectures exist primarily to prevent collapse, and that removing them yields encoders at chance level. With SIGReg in place the paper removes both and trains fine. A teacher–student setup still helps ViTs slightly, but as stochastic weight averaging rather than as a collapse guard. Register tokens, introduced elsewhere to stabilize vision transformers, turn out to be unnecessary too, and the paper’s read is that those instabilities were symptoms of poorly conditioned objectives rather than facts about the architecture.

9.2 The training loss predicts downstream accuracy

LeJEPA’s training loss correlates with downstream linear-probe accuracy at roughly 85% Spearman out of the box, rising to nearly 99% under a simple rescaling, which makes it usable for label-free model selection.

This is the result I did not expect, and the most practically useful one in the paper. Ordinarily an SSL pretraining loss tells you almost nothing about representation quality; a contrastive loss can fall steadily while the representation quietly loses rank, which is exactly the dimensional collapse from Section 2 and is invisible in the curve. The paper notes that in recent JEPA models the training loss may not even decrease monotonically. So the standard workflow is to checkpoint, freeze, fit a probe on labelled data and read the accuracy, meaning every model-selection decision inside a supposedly unsupervised procedure secretly requires labels.

The rescaling is worth stating because it is so simple. Correlate the loss divided by a power of the trade-off parameter:

C(α) = ρ_s ( train_loss / λ^α , test_accuracy )

At α = 0 you recover the plain training loss, already strongly correlated. At α ≈ 0.4 the correlation reaches nearly 99% across multiple datasets and architectures, with the individual curves running from 0.60 to 0.95 on ResNet-18 with Flowers102, and 0.90 to 0.98 on ViT-S/8 with ImageNet-10. The entire set of points is collected across variation in learning rate, weight decay, epoch count and λ, so this is not a within-run artifact.

Why does LeJEPA behave differently? The answer follows from Section 6, and it is worth stating because the paper reports the observation more than it explains it. The SIGReg term is not a proxy for downstream performance, it is a direct estimate of the quantity the theory says governs downstream performance. If the isotropic Gaussian argument is correct, the loss ought to predict the probe. That it empirically does, at 99% correlation, is the strongest available evidence that the argument describes something real rather than merely being internally consistent.

There is a nice secondary observation in the loss plane. Different values of λ trace out trade-off fronts between the two terms that produce similar downstream accuracy, and those fronts are linear and point towards the lower-left corner. Which is to say the two terms are not fighting each other in any complicated way; you want both small, and λ decides how you get there.

9.3 In-domain pretraining outperforms frontier transfer

On Galaxy10, LeJEPA pretrained directly on the target data with a small backbone outperforms DINOv2 and DINOv3 transferred from natural images, across every data regime from one sample per class to the full set, under both frozen-backbone linear probing and full finetuning.

The setup deserves stating because it is deliberately unflattering to LeJEPA. Galaxy10 is a galaxy-morphology task with 11,000 training samples across 10 classes, visually and statistically remote from natural images. The paper pretrains a variety of backbones on it for 400 epochs using default hyper-parameters, no tuning, and compares against DINOv2 (trained on LVD-142M) and DINOv3 (trained on LVD-1.7B).

Reading the frozen-backbone column, which is where representation quality shows most directly: LeJEPA ResNet-34 reaches 78.17 against DINOv3’s 71.38 and DINOv2’s 67.62 on the full set, and the gap holds at every sample count down to one per class. Under full finetuning the same ResNet-34 reaches 83.28 against 81.60 and 78.34. A 21M-parameter model trained from scratch on eleven thousand images beats a model pretrained on 1.7 billion.

The effect is not confined to astronomy. Across six datasets the paper shows in-domain LeJEPA competitive or better, including Flowers102 with only 1,020 training samples, where a ResNeXt-26 reaches 82.19 from scratch. The one place frontier transfer still wins comfortably is where you would expect it to, on natural-image datasets like CIFAR and ImageNet-10 where I-JEPA pretrained on ImageNet-22k has an enormous head start.

Why this matters beyond the leaderboard: the received wisdom is that in-domain SSL on a small dataset is not worth attempting, because you cannot out-train a foundation model pretrained on a billion curated images, so you transfer and finetune instead. That wisdom formed when SSL required careful per-dataset tuning and collapsed if you got it wrong, which made in-domain pretraining an expensive gamble on a dataset too small to justify it. Take away the tuning and the gamble disappears, which suggests the received wisdom was partly an artifact of the methods rather than a fact about the data.

9.4 Scaling across data and models

On ImageNet-1k, LeJEPA reaches 77.1% online linear probe with a ViT-Large (0.3B) and 78.5% with a ConvNeXtV2-Huge (0.6B), and transfers better than I-JEPA while using a smaller model and a third of the training schedule.

The transfer comparison is the more informative one. Against I-JEPA with a ViT-Huge (0.6B), across eight datasets spanning textures, objects and fine-grained categories, LeJEPA ViT-L (0.3B) averages 79.48 to I-JEPA’s 78.50 in the all-shot setting, while training for 100 epochs against I-JEPA’s 300. That is roughly a threefold reduction in training cost with a model half the size, and the margins are widest exactly where you would hope, on fine-grained tasks; DTD 78.30 against 73.32, Food101 82.05 against 81.02, Flowers102 91.21 against 86.47.

On the largest run, a ViT-gigantic at 1.8B parameters trains with a smooth, stable loss curve and no special handling, which the authors attribute directly to the bounded-gradient guarantee from Theorem 4.

Worth framing carefully, because this is easy to over- or under-sell. These are not state-of-the-art numbers in absolute terms; heavily tuned DINOv2 and DINOv3 systems reach higher on ImageNet. What they are is competitive numbers produced by an objective with one hyper-parameter and no heuristics, at a fraction of the training budget. The useful comparison is not whether 78.5% wins the benchmark, but how much was given up by deleting the teacher, the predictor, the stop-gradient, the schedulers, the whitening and the register tokens, and the answer appears to be less than one would expect.

9.5 Emergent semantic structure

PCA of last-layer features from an ImageNet-1k pretrained LeJEPA ViT-Large shows clear object–background separation, and thresholding the attention maps yields object segmentation and tracking across video frames without any segmentation supervision.

Projecting the first three principal components to RGB gives images in which warm colours consistently pick out foreground objects and cool colours pick out background and foliage. This is the DINO-style qualitative check, and LeJEPA passes it.

The video result goes further than I expected. Thresholding the self-attention of the [CLS] token produces binary masks that track salient objects across frames with reasonable temporal coherence, again with no segmentation labels anywhere in training.

This is soft evidence and I would not lean on it hard, but it is a good consistency check. Nothing in SIGReg mentions semantics; the constraint only says the embedding cloud should be an isotropic Gaussian. Semantic organization appearing along the principal directions anyway is what you would hope for if the predictive term is doing its job and the constraint is genuinely neutral about what gets represented, caring only about how the mass is arranged. The authors’ framing is that a stability-focused objective does not sacrifice semantic richness, and on this evidence that holds.

9.6 Limitations of the evidence

Three things the results do not establish, which are worth stating plainly.

1. Vision only. Every experiment is on images. The Section 6 theory makes no reference to modality, since the risk argument concerns probes and distributions rather than pixels, so there is no principled barrier to audio, text or time series. Nobody has demonstrated it, and “should transfer” has a poor track record here.

2. No matched-compute comparison against the strongest baselines. The I-JEPA comparison is careful and favourable, but DINOv2 and DINOv3 are compared only in the in-domain setting where they are transferring rather than training in-domain. The experiment that would settle the general question is DINOv3’s exact data, architecture and compute with only the objective swapped, and nobody has run it.

3. “One hyper-parameter” is generous. True of the objective, and it excludes learning rate, weight decay, batch size, augmentation policy and architecture, all of which still require choosing; the paper cross-validates learning rate and weight decay in its architecture sweep. The accurate claim is that LeJEPA eliminates the hyper-parameters specific to preventing collapse, which is a real reduction but not a single knob in total.

What I will like to point out is that the limitation I expected to find is not there as I assumed the number of slices would be an underexamined knob hiding a lot of sensitivity, since it controls how well the Cramér–Wold approximation holds. The paper ablates it directly, bounds it theoretically, and shows that resampling makes small values work.

10. Thoughts

1. The heuristics were never wrong, they were unnamed.

Every mechanism from Section 4 was groping towards a constraint on the embedding distribution without stating what it was. VICReg got the second moments, whitening got the correlation structure, negatives got the spread. All three are shadows of “be isotropic Gaussian” seen from different angles. The field was doing empirical science on an objective nobody had written down, and each heuristic was a partial measurement of the same object.

2. The loss does not predict performance, it measures it.

The 99% correlation is under-sold in the paper. Operationally it means architecture search, hyper-parameter selection and early stopping without touching a label, which removes the most annoying dependency in the pretraining workflow. But the reason matters more: SIGReg estimates the quantity that provably bounds downstream risk. If that framing is right, any objective built on a proven risk bound should behave this way, and the loss-tells-you-nothing era was a symptom of under-specified objectives rather than a fact about self-supervision.

3. Gaussianize once, or Gaussianize always?

In TurboQuant the central move was a Hadamard rotation applied to KV-cache vectors so that, by CLT, their coordinates become approximately Gaussian, because a known distribution lets you precompute an optimal Lloyd-Max codebook. The whole method manufactures Gaussianity at inference because the model did not provide it.

LeJEPA trains the representation to be Gaussian in the first place. So: if you train with SIGReg, do you still need the rotation? A model whose hidden states are natively isotropic Gaussian is already in the distribution Lloyd-Max wants; skip the Hadamard, skip the online scale estimation, quantize directly.

Both papers want the same target for the same reason and arrive from opposite ends, one post-hoc because it cannot change the model, one during training because it can. Both lean on random one-dimensional projections to make a high-dimensional problem tractable; TurboQuant to preserve inner products under sign quantization, LeJEPA to test distributional fit. The same tool pointed in opposite directions.

4. Isotropic is right only if you are genuinely not betting.

The isotropic Gaussian is optimal for worst-case risk over an unknown task family, which is exactly right for a foundation model and not obviously right when you know a lot about your downstream use. There, worst-case optimality means declining to use information you have. SIGReg does not care what Q is; the machinery works for any target whose one-dimensional projections you can write down. A task-informed, deliberately anisotropic Q is the most immediately actionable gap I can see.

5. What happens in language?

Every experiment is vision and the theory is modality-agnostic, so text is the obvious move. I suspect it is harder than it looks: LLM embeddings are famously anisotropic, the cone effect, and there is an argument that the anisotropy encodes real structure about token frequency rather than being pathology. Forcing isotropy might destroy something the model was using. Or the cone is dimensional collapse we have been rationalizing for years because we had no principled reason to object. Section 6 gives us that reason now.

6. The contribution is the question.

The contribution is not SIGReg, and not the isotropic Gaussian result either. It is asking what distribution should the embeddings follow, which nobody had posed in that form. Once asked, the answer is a fairly standard exercise in bias-variance and Fisher information, and the implementation is a statistical test from 1983 wrapped in a theorem from 1936. None of the pieces is new; the framing is. Worth remembering when reading papers for techniques to borrow, since the technique here is a fifty-line function, but the thing worth stealing is the habit of asking what the correct answer would look like before going to find one.

11. Conclusion

We talked about self-supervised learning, and specifically about the part of it nobody could justify. Stop-gradients, EMA teachers, whitening, negatives, register tokens; all found by trial and error, all kept because they prevented collapse.

The route out was taking Yann’s framing seriously. If SSL is constrained optimization, the constraint is not a safety mechanism, it is the thing that decides what the model learns. Bayes made that precise in Section 3, giving a joint prior with two slots — one for parameters, which is classical regularization, and one for the representation, where every anti-collapse mechanism had been operating without saying so.

That turned the question from how do we stop the model cheating into what shape should a representation be, and the second has exactly one answer. A linear probe forces isotropy, since a variance budget spread unevenly always leaves some task badly served. A nonlinear probe forces Gaussianity, since neighbourhood-averaging probes are biased in proportion to Fisher information, which the Gaussian uniquely minimizes. Isotropy declines to bet on direction, Gaussianity declines to bet on location.

Enforcing it was harder. Density estimation in a thousand dimensions from a few hundred samples is hopeless, and the escape was Cramér–Wold: a distribution is determined by its one-dimensional shadows. Project onto random directions, run a characteristic-function test on each, average. The result is fully specified, bounded in gradient by construction, linear in everything you want to scale, and about twenty lines of PyTorch. Then delete the machinery, because none of it does anything the constraint does not.

The empirical case holds. Fifty architectures on one hyper-parameter. A loss predicting downstream accuracy at 99%, so model selection without labels. Eleven thousand galaxy images beating a model pretrained on 1.7 billion.

What strikes me most is how little here is new. Epps–Pulley is from 1983, Cramér–Wold from 1936, Fisher information older still. The contribution is a question nobody had asked in that form, and everything else follows from asking it properly.

Whether SIGReg becomes standard I have no idea. The durable part is the move underneath it: when something in your pipeline works and you cannot say what it guarantees, that is not a trick you have found, it is a specification you have not written.

References:
LeJEPA paper: https://arxiv.org/pdf/2511.08544
Code: https://github.com/rbalestr-lab/lejepa
Yann's NYU lecture:

That's all for today.
Follow me on LinkedIn and Substack for more such posts and recommendations, till then happy Learning. Bye👋

Large language models are slow when they generate text, and the reason catches most people off guard. It is not that the math is hard. The problem is that the model has to read its entire multi-hundred-gigabyte brain out of memory once for every single word, and it cannot start the next word until the current one is finished. Speculative decoding is the trick that breaks this one-word-at-a-time barrier, and it does so without changing the model’s output at all.

This post builds the idea up from the ground, then puts it into production. I start with why generation is slow, work out the core “guess-and-verify” mechanism, prove it is mathematically exact, dig into the economics that decide how much speedup you actually get, walk the method families (n-gram lookup, Medusa, EAGLE), and finish by actually deploying EAGLE3 on vLLM and reading the speedup off a real GPU.

Table of Contents

1. Why generation is slow

2. The core idea: draft, then verify

3. Why it’s exact: the correctness guarantee

4. The economics: alpha, tau, and K

5. System realities and limitations

6. The landscape of methods

7. Putting it in production: EAGLE3 on vLLM

8. Takeaways

9. What I didn’t cover

Code Repo

https://github.com/Mayankpratapsingh022/LLM-Inference-Playbook

Speculative decoding makes a large model generate 2–3× faster while producing the exact same text it would have written anyway. It is close to a free lunch, paid for with spare compute the GPU was already wasting, but it is a bet, and Section 7 shows a real run where it did not pay off.

1. Why generation is slow

1.1 Generation happens one token at a time

An LLM generates text autoregressively: it predicts one token, appends it to the input, and runs the whole model again for the next one. There is no shortcut here. Each new word literally depends on the one before it.

Figure 1: In autoregressive decoding, the large model predicts only the next word, which is then appended to the context. The model runs again to predict the subsequent word, and this process repeats many times to generate the complete sequence.

Notice what this costs. To produce the phrase “a small red,” the giant model runs three separate times. To write a 500-token answer, it runs 500 times, strictly in order. Every run is a full forward pass through every layer of the network.

1.2 Training is parallel, inference is sequential

This sequential bottleneck is unique to generation. During training, the model already knows the whole target sentence, so it can score every position at once in a single parallel forward pass. During inference, the future words do not exist yet, so the model is stuck in a step-by-step loop.

Figure 2: A comparison of Transformer processing modes during training and inference. During training (Panel A), the full sentence is known, enabling the model to score multiple positions simultaneously in a single parallel forward pass. In contrast, during inference (Panel B), future words are missing, forcing the system into a sequential, step-by-step generation loop where each new token must be predicted one at a time.

The transformer is perfectly capable of massive parallelism. We just cannot use it at generation time, because we are missing the very thing we are trying to produce.

1.3 The autoregressive barrier

Here is another way to see the problem. At any moment, the model has a valid KV cache for everything it has already written (the past), and a complete blind spot for everything it has not (the future). It can generate exactly one token to push the boundary forward, and only then can it look one step further.

Figure 3: An illustration of the autoregressive barrier during LLM inference. Because the model cannot see future words, it uses past context to actively generate the immediate next token, and must complete this step before moving to the next slot. Speculative decoding attempts to bypass this strict sequential bottleneck by filling the future unknown slots with rapid guesses.

That last sentence is the whole thesis of this post in one line:

1.4 The real culprit is memory bandwidth, not compute

Here is the part that trips people up. Each generation step is slow not because the GPU is busy doing math. It is slow because the GPU has to physically stream the model’s entire set of weights, every single parameter, out of high-bandwidth memory (HBM) and into the compute cores, and it does this for every token it produces.

Figure 4: Key system observations highlight the inefficiencies of autoregressive decoding, noting that every new word necessitates another full pass through the large model. Because the same massive weights are repeatedly loaded, each additional word incurs another computationally expensive step, ultimately causing long answers to generate slowly.

For a 70B-parameter model, that is roughly 140 GB read from memory per token. Three tokens means moving 420 GB. The arithmetic units, meanwhile, sit mostly idle. They finish long before the next batch of weights even shows up. This is what people mean when they say generation is memory-bandwidth bound.

This one fact is the foundation of everything that follows:

And we can feed it more tokens at once. Remember Figure 4, Panel A: verifying many tokens in parallel costs almost the same as generating a single one, because either way you pay to load the weights exactly once. That asymmetry, where one weight-load can check many tokens, is the gap speculative decoding pries open.

2. The core idea: draft, then verify

2.1 A small model guesses, a large model checks

Speculative decoding adds a second, much smaller draft model (M_q​) alongside the big target model (M_p​). The draft model is cheap to run, so it can rattle off several guesses for the next few tokens very quickly. The target model then verifies all of those guesses in a single parallel pass, the same pass that in normal decoding would have produced just one token.

Figure 5: In the speculative decoding pipeline, a small model (M_q​) rapidly guesses future tokens while a large model (M_p​) verifies those guesses in parallel. Correct guesses are committed to the KV Cache, and rejected ones are corrected to preserve M_p​'s exact distribution.

The draft model produces K candidate tokens in roughly 4×T_fast​ (its drafting latency). Then the target model spends one expensive pass, T_slow​, confirming the whole batch instead of producing a single token. When the guesses are good, you get several tokens for the price of one big step.

2.2 Verification accepts a prefix, then corrects

How does verification decide what to keep? The target model scores the drafted sequence left to right and accepts tokens until the first mismatch. Everything before the mismatch is committed. The mismatched token is replaced with the target model’s own correct token. Everything after it is thrown away.

Figure 6: The token verification process. The target model evaluates the drafted sequence sequentially from left to right. It successfully verifies and accepts "a" and "small". However, upon rejecting the token "blue", all subsequent guesses in the drafted sequence are automatically discarded.

This "first mismatch wins" rule is what keeps the output identical to plain decoding. The moment a guess strays from what the big model would have said, we stop trusting the draft and hand control back to the target. And notice the worst case: even if every single guess is wrong, the verification pass still hands back one correct token (the corrected one), so we never do worse than plain autoregressive decoding.

2.3 The speedup, in numbers

Let me put concrete costs on it. Say one big-model pass costs 100 units and one draft pass costs 10.

Standard decoding of five tokens is five serial big passes: 5×100=500

Figure 7: Cost comparison and best-case performance of speculative decoding versus standard autoregressive decoding. Standard decoding requires paying for five large, sequential passes to generate five tokens, resulting in a baseline cost of 500. Conversely, speculative decoding employs a computationally cheap draft model to guess five words simultaneously. In an ideal scenario where the draft model guesses all words correctly, the target model verifies the entire block once in parallel, which reduces the total operational cost from 500 to 150.

In the best case, speculative decoding pays for five cheap draft guesses plus one verification pass: (5×10)+100=150. Same five tokens, less than a third of the cost. The catch, which the rest of this post is really about, is that this best case only shows up when the draft model guesses well. How often it does is the whole game.

3. Why it’s exact: the correctness guarantee

The most important property of speculative decoding is that it is not an approximation. The text you get out is drawn from exactly the same probability distribution as if the target model had generated every token itself. This is what separates it from lossy tricks like just using a smaller model directly.

The mechanism is a clever twist on rejection sampling, introduced for this setting by Leviathan et al. (2023) and Chen et al. (2023). Let q(x) be the draft model’s probability for a proposed token x, and p(x) the target model’s probability for that same token. We accept the draft token with probability:

Read it intuitively. If the target likes the token at least as much as the draft did (p≥q), we keep it outright. If the target likes it less, we keep it only proportionally, and sometimes reject it, so we do not over-represent tokens the draft was too eager about.

When a token is rejected, we do not just resample blindly. We sample the replacement from the normalized residual distribution, which fills in exactly the probability mass the acceptance step left out:

Together, the accept rule and the residual-resample rule provably rebuild the target distribution p, token for token. The draft model only ever affects speed. It never affects what actually gets written.

4. The economics: alpha, tau, and K

Speculative decoding is only as good as its guesses. Three quantities decide how much real speedup you see: the acceptance rate α, the expected accepted tokens per round τ, and the lookahead budget K.

4.1 Alpha (α): how often guesses are accepted

Alpha is the probability that any given drafted token survives verification. It is the single most important efficiency number, because rejected tokens are wasted work.

Figure 8: The Role of Alpha (α) in Speculative Execution. Alpha serves as the fundamental efficiency metric in speculative decoding, quantifying the rate at which the target model accepts drafted tokens. High α values indicate strong distribution alignment and yield significant inference acceleration. Conversely, low α values reflect computational waste, as misaligned draft tokens trigger early rejections and serial fallbacks.

When α is high, the draft and target models "think alike," long runs of guesses get accepted, and you approach that ideal 150-cost case. When α is low, you reject early and often, and most of your drafting effort goes in the bin.

4.2 Tau (τ): expected tokens per verification

So how many tokens do you actually commit per round, on average? If each of the K drafted tokens is independently accepted with probability α, the expected number of accepted tokens per round is a geometric series:

Figure 9: With K=5 as the lookahead budget for the number of draft tokens proposed during each validation step, the variable α defines the independent probability that a draft token matches the target distribution, which determines τ, the estimated average of accepted tokens per round. Achieving a larger τ is beneficial because it reduces the frequency of costly target-model runs, directly lowering system latency.

The numbers tell the story. A weak drafter (α=0.30) commits only about 1.43 tokens per expensive pass, barely better than plain decoding. A strong drafter (α=0.85) commits about 4.15 tokens per pass, and that is where the real speedups come from. Acceptance compounds across slots, which is why the bars decay from left to right. To get slot 5 right, slots 1 through 4 all have to be right first.

4.3 Lookahead K (γ): how far to guess

K (sometimes written γ) is how many tokens the draft model proposes before each verification. It is a balance, not a “bigger is better” dial.

Figure 10: The parameter K determines the number of tokens drafted per step. A value too small limits speedup, while a value too large wastes compute on discarded tokens. To maximize efficiency, optimal systems dynamically tune K based on the token acceptance rate.

If K is too small, you verify too often and leave easy speedup unclaimed. If K is too large, you draft far past the first inevitable mismatch and burn compute on tokens you will only throw away. The sweet spot depends on α: confident drafters can afford a deeper lookahead.

4.4 Adaptive K: tune the lookahead to the text

Since the ideal K depends on how predictable the upcoming text is, good systems do not fix it. They use an adaptive controller that reads the local difficulty of the stream, usually through the entropy H(x) of the model’s distribution, and sets the lookahead to match.

Figure 11: Dynamic lookahead overcomes the limitations of a static K budget by adjusting draft length based on textual confidence, allowing farther drafts for easy text and shorter drafts for hard text, to minimize wasted compute and memory bandwidth.

On boilerplate or highly predictable text ("in the logs we found a critical error in the system"), the controller drafts far ahead, because the guesses will mostly be accepted. On genuinely uncertain text ("therefore the proof…"), it drafts cautiously, because deep guessing would just produce garbage to discard.

5. System realities and limitations

Speculative decoding looks like a clean win on paper. In a real serving system it comes with sharp edges worth knowing about before you reach for it.

5.1 The two models run in lockstep

The draft and target models do not work at the same time. They take turns. While the draft model is busy proposing, the expensive target model sits idle. While the target verifies, the draft model sits idle.

Figure 12: Speculative decoding's sequential workflow. In Phase 1, the large, expensive target model sits idle while the draft model generates its guesses one by one. In Phase 2, the bigger issue, drafting stops entirely while the draft model waits for the slow target model to finish its verification.

This back-and-forth is why naive implementations sometimes fall short of their theoretical speedup. There is real "bubble" time where one of the two models is doing nothing at all.

5.2 It only helps when the GPU has spare compute

Remember that the whole trick leans on spare compute (Section 1.4). That spare capacity exists at low batch sizes, when one or a few requests leave the GPU’s math units idle. At high batch sizes, many requests already saturate the GPU, and speculation’s extra drafting work just steals memory bandwidth away from useful generation.

Figure 13: Because speculative decoding only reduces latency when GPUs have spare compute, and actively steals memory bandwidth under high load, modern serving engines dynamically toggle the feature based on real-time queue depth to maximize overall system throughput.

This is why production serving stacks treat speculation as a latency optimization for light load, switching it off automatically when the queue fills up and the system flips into throughput-maximizing mode.

5.3 Limitation 1: the tokenizers must match

Because verification compares the draft’s token against the target’s token, both models have to speak the exact same vocabulary. A draft model that splits “active” into [act] and [-ive] cannot be checked against a target that emits a single [active] token. The two sequences would never line up.

Figure 14: Limitation 1, tokenizer identity. The draft and target models must share an identical tokenizer and vocabulary; otherwise their token sequences cannot be aligned during verification.

In practice this means your draft model is usually a smaller member of the same model family as the target (say, a 7B drafting for a 70B with a shared tokenizer), which limits your options.

5.4 Limitation 2: you pay for a second model in VRAM

Speculative decoding puts two models on the GPU at once. The draft model is small, but it is not free, and the memory it takes up is memory you cannot hand to the KV cache or to bigger batches.

Figure 15: Limitation 2, dual-model VRAM overhead. Co-locating both models in VRAM (for example ~14 GB for a 7B draft plus ~140 GB for a 70B target) severely reduces the capacity left for KV caches and high batch sizes.

5.5 Limitation 3: rejection triggers a serial fallback

When a drafted token is rejected, every guess after it gets discarded and the system falls back to generating that step the slow way. On hard text, rejections come often, so you keep re-drafting and re-verifying, and the speedup for that step evaporates.

Figure 16: Limitation 3, distribution alignment and the rejection loop. When the draft and target distributions diverge, proposals get rejected and discarded, and the system falls back to a serial re-drafting loop that erases the speedup for that step.

All three limitations point back to the same truth: speculative decoding is a bet. It pays off handsomely when the draft is well-aligned and the GPU has room to spare, and it quietly costs you when those conditions fail.

6. The landscape of methods

“Speculative decoding” is now a family of techniques, not a single algorithm. They differ mainly in where the draft tokens come from.

Figure 17: A timeline of speculative decoding methods, including the separate vanilla draft model, model-free n-gram lookup, and single-model multi-head drafting (Medusa).

I will focus on the three most foundational and widely used approaches:

1. Vanilla draft-target (Section 2): two separate models, where the draft proposes and the target verifies. This is the classic formulation.

2. N-gram prompt lookup (6.1): no draft model at all. The drafts are copied from repeated text in the prompt.

3. Medusa (6.2): a single model grows extra heads that predict several future tokens at once.

The figure also shows Eagle (feature-level tree drafting), self-speculation (drafting by skipping the target's own middle layers), and streaming overlap (pipelining draft and verify). I introduce the core idea behind EAGLE in 6.3, because it is the method that now dominates production serving; its tree-based refinements (EAGLE-2/EAGLE-3), self-speculation, and streaming overlap are more advanced and belong to a future post.

6.1 Deep dive: n-gram prompt lookup

The simplest drafter is no model at all. In a lot of real tasks (summarization, document Q&A, code editing, retrieval-augmented generation) the output copies long spans verbatim from the input. Prompt Lookup Decoding (Saxena, 2023) takes advantage of this directly.

Figure 18: N-gram prompt lookup. Instead of a neural drafter, the system scans the past context for a repeat of the most recent n-gram suffix and proposes the tokens that followed it last time. The target model verifies these guesses in a single pass. This is model-free and adds zero VRAM, but its hit rate falls on creative, non-repetitive text.

The mechanism is just string matching:

1. Take the last few generated tokens (the active n-gram suffix, say if key in).

2. Scan backward through the prompt and prior output for an earlier occurrence of that same n-gram.

3. If you find one, propose the tokens that followed it last time (cache: return) as the draft.

4. Verify those proposals with the target model exactly as before.

Because no neural network is involved in drafting, this approach uses zero extra VRAM, which sidesteps Limitation 2 completely, and it works with any decoder model out of the box. The trade-off is that its hit rate is only as good as the repetition in your text. It shines on structured, input-grounded generation (2–4× speedups are common) and does almost nothing for open-ended creative writing.

6.2 Deep dive: Medusa

Medusa (Cai et al., 2024) takes a different route. Instead of a separate draft model, it bolts extra decoding heads onto the target model itself. One backbone pass produces several future-token predictions at once, and they all get verified together. No second model, no separate tokenizer to babysit.

Multiple heads predict multiple positions

A normal LLM has one “LM head” that predicts the very next token (position +1). Medusa adds extra heads on top: head 1 predicts position +2, head 2 predicts +3, head 3 predicts +4. All of them read from the same backbone hidden state, computed in a single forward pass.

Figure 19: Medusa augments a model with extra decoding heads. After the backbone runs once, the original LM head predicts the next position while each added Medusa head predicts a position further out, and each head keeps its top-2 candidates.

Because each head keeps its top-2 candidates, the heads together describe many possible continuations, not just one. That branching is handled by a tree.

Tree attention verifies many candidates at once

With each head offering 2 candidates, the continuations form a small tree. Two heads give 2×2=4 candidate paths, and a third head would give 8. Medusa packs all these paths into a single sequence and uses a specially shaped tree attention mask, so each candidate token only attends to its own ancestors. The branches stay logically separate even though they ride through one shared forward pass.

Figure 20: Medusa’s candidate tree and tree-attention mask. The top-2 outputs of each head combine into multiple candidate paths (top-2 × top-2 = 4, or 8 with a third head). The attention mask ensures each row attends only to its ancestors and itself, so independent branches do not contaminate each other.

Here is the payoff. All of these candidate paths get checked in one backbone pass. The target verifies the whole tree at once and commits the longest path that survives, and that committed path immediately seeds the next tree.

Figure 21: A single backbone pass with tree attention verifies every candidate path at once. Here the path "returned · the · cached" is accepted, 3 tokens committed in one pass, and it seeds the next round's tree.

Two training recipes: Medusa-1 vs Medusa-2

How you train those heads decides the trade-off between quality and speed.

Figure 22: Medusa's two training procedures. In Medusa-1 the backbone is frozen and only the heads are trained, so accepted tokens match the original model exactly (lossless). In Medusa-2 the backbone is fine-tuned alongside the heads using a careful recipe (a combined loss, split learning rates, and warming up the heads first), which yields more accurate heads and higher speedup while keeping quality intact.

• Medusa-1 freezes the backbone and trains only the heads. Since the backbone is never touched, the verified output is lossless, identical to the original model. It is cheap to train (one GPU) and a safe default.

• Medusa-2 fine-tunes the backbone and heads together. This makes the heads more accurate, so more tokens are accepted per step, but it risks degrading the base model, which is why it needs the careful recipe shown above.

The lever that connects training to speed is simple:

better heads → more tokens accepted per step → higher speedup.

Figure 23: Why Medusa-2 is faster comes down to one lever, head accuracy. More accurate heads get more tokens accepted, which means more tokens per step and higher speedup. Numbers from Cai et al. (2024), Vicuna-7B on MT-Bench, both variants trained on the same ShareGPT data for 2 epochs.

6.3 EAGLE: drafting features, not tokens

Everything so far, vanilla draft-target, n-gram lookup, Medusa, shares one habit: the drafter guesses raw tokens, and verification compares those exact tokens against the target’s. That is fine on predictable text. It falls apart exactly where language is hardest.

Why token-level guessing stalls

Recall from Section 4 that speedup lives and dies by the acceptance rate α. The trouble is that α is not a constant, it collapses in high-entropy contexts, the moments where many different next tokens are all reasonable.

Figure 24: In a high-entropy context, several next tokens are plausible at once. The target's distribution p favours "the" while the smaller draft's q favours "a". Because verification accepts only up to the first token-level mismatch, this single disagreement rejects the draft early and cuts the accepted prefix short.

Look at what happens. The target model spreads its probability p across returned, caused, the, and more, landing on "the". The little draft model, working with a coarser view of the world, puts its mass q on "a". Neither is "wrong", both are grammatical, but the exact-token verifier does not care about that. The first token-level mismatch ends the accepted run right there.

The cost compounds through the same geometry seen in Section 4. If entropy pushes per-token acceptance down to, say, α=0.4, then with a lookahead of K=5:

Barely more than one token committed per expensive verification pass. The more the draft's distribution diverges from the target's, the shorter the accepted prefix, and the smaller the real speedup. Guessing the exact word is simply too brittle a target when many words are plausible.

EAGLE’s fix: predict the feature, not the word

EAGLE (Li et al.) makes one deceptively simple change. Instead of training the drafter to predict the next token, it trains a lightweight module to predict the target model’s next feature, the hidden state from the layer just below the LM head. Features are continuous, smooth, and far more predictable than the discrete lottery of exact tokens, so the drafter agrees with the target much more often.

Figure 25: EAGLE drafts at the feature level. It takes the target's own hidden feature (Layer N−1 output) plus the embedding of the token that actually followed, concatenates them, and runs a tiny autoregressive head (one FC layer + one decoder layer) to predict the next feature. That predicted feature loops back in as the next input, and is also decoded to a draft token through the frozen target LM head. The target model itself is never modified.

Read the pipeline left to right:

1. Reuse the target’s features (left). The target model is frozen. EAGLE taps the hidden state from its second-to-last layer, the “feature” that summarises everything the model knows right after the word is. Nothing about the big model changes.

2. Anchor each guess with the real next word (middle). EAGLE feeds the drafter the feature and the embedding of the token that actually came next (is → blue). That next word is what tells the drafter which branch of reality actually happened (blue, not falling), so its predictions stay grounded instead of drifting.

3. Autoregress on features (right). A tiny head, one fully-connected layer plus a single decoder layer, takes the concatenation [featis ; emb(blue)] and predicts the next feature. That predicted feature loops straight back in as the input for the following step (this is the autoregressive part), and in parallel it is pushed through the frozen target LM head to read off an actual draft token (today).

Because the drafter learned to imitate the target's own internal trajectory, its proposals line up with what the target would have produced, so acceptance climbs and τ with it. And crucially, since draft tokens are still decoded through the target's real LM head and verified the usual way, EAGLE keeps the exactness guarantee from Section 3 intact. The feature trick buys speed, never correctness.

From theory to deployment. EAGLE's descendant, EAGLE3, is exactly what I deploy in Section 7 of this post, on a real Llama-3.1-8B with vLLM. If you want the flags, benchmarks, and measured numbers, jump to 7.

6.4 The production reality: where the speedup goes

The clean 1150-vs-500 arithmetic from Section 2 is the theoretical story. On a real serving stack, two things bend those numbers: fixed per-step system overheads, and the batch size the GPU is running at. This section is the honest accounting.

Ideal vs measured: the system tax

Even in the friendliest case, a single request at batch 1 with every drafted token accepted, you do not get the speedup the token math promises. Real engines pay fixed costs the idealized picture ignores.

Figure 26: The gap between the paper speedup and the measured one, at batch 1. The theoretical ideal (left) accepts all four drafted tokens for a 2.5× win. The production timeline (right) inserts real per-step costs, kernel-launch overhead and KV-cache lookups, that the token arithmetic never counted, dragging the effective speedup down to roughly 1.8×. (Illustrative numbers.)

The extra blocks on the right, kernel-launch overhead and KV-cache lookup, are the "system tax." They are largely fixed per step, so they do not shrink just because your draft was good. A 2.5× algorithm becomes a 1.8× system. This is why a speedup you can prove on your own hardware is worth more than any number on a slide.

Batch size decides everything

Section 5.2 made the claim that speculation only helps when the GPU has spare compute. Here is the mechanism, drawn out across the two regimes.

Figure 27: The same speculative machinery in two regimes. At batch 1 (latency-bound) the GPU is mostly idle, so the draft model's extra work hides inside that idle time and verification is nearly free, this is where speculation shines. At batch 128 (throughput-bound) the GPU is already saturated, drafting kernels serialize and compete for bandwidth, and the step is held hostage by its slowest request (the straggler effect), leaving effective speedup somewhere between −5% and +10%.

The two lanes could not be more different:

• Lane 1, batch 1 (latency-bound). The compute engine sits mostly idle waiting on memory. Drafting slots into that idle time for free, and verification overlaps without adding wall-clock. Spare compute is exactly the currency speculation spends, and here it is abundant.

• Lane 2, batch 128 (throughput-bound). Now 128 requests already saturate compute and memory bandwidth. Every extra drafting kernel is no longer free, it serializes and steals bandwidth from useful generation. Worse, a batched step only finishes when its slowest member finishes, so one request that rejects and re-drafts (the straggler) stalls the whole batch. Net effect: anywhere from a small loss to a modest gain.

This is precisely why production stacks toggle speculation on at light load and off under saturation, the dynamic switch introduced in Section 5.2.

Following one decoding step’s clock

One more view makes the batch-size story concrete: take a single full decoding step, call it 100% of its wall-clock time, and watch how that time redistributes as the batch grows.

Figure 28: The composition of one decoding step's wall-clock time at three batch sizes. As the batch grows from 1 to 128, target verification's share climbs (52% → 63% → 74%) while the draft model's share shrinks (24% → 17% → 10%). (Illustrative numbers.)

The trend is the whole point. At batch 1, the draft forward pass is a real slice of the step (24%), so making it cheap or making its guesses land pays off handsomely. By batch 128, verification dominates (74%) and drafting is a rounding error (10%). As the batch grows, verification eats the step and the draft’s share collapses, so the headroom speculation can possibly recover shrinks right along with it.

7. Putting it in production: EAGLE3 on vLLM

Everything above was the idea. Now I run it. I deploy the method that won production, EAGLE3, on vLLM, serving Llama-3.1-8B on a single 48 GB GPU, and measure the speedup against a real baseline. Same model, same prompts, one flag flipped.

This half is also a notebook. Every command below is a runnable cell in the companion vLLM notebook.

Code Repo

https://github.com/Mayankpratapsingh022/LLM-Inference-Playbook

The subsection numbers (7.1, 7.4 and so on) match the notebook, so you can read a part here and run the same thing there. Spin up the pod from 7.2, upload the notebook into the pod's JupyterLab, and step through it cell by cell while you read along. I tried hard to make it as simple as it can possibly be: open it, run each cell, and read the numbers off.

I use an 8B target on one affordable card so you can reproduce it for about a dollar. Fair warning, and it is the honest point of 7.7: on this small-model, single-GPU setup the speedup did not materialize, and that section explains exactly why.

7.1 What we are building

The experiment is deliberately simple to reason about:

• Target model: meta-llama/Llama-3.1-8B-Instruct.

• Hardware: a single 48 GB RunPod GPU (an A40, L40S, RTX 6000 Ada, or A6000). One card is plenty, see the memory note in 7.2.

• Engine: vLLM, the most widely deployed open inference server.

• Method: EAGLE3, the current state-of-the-art (7.3 explains why I don’t bother with the rest).

• What we toggle: a baseline run with plain autoregressive decoding, then the same model with EAGLE3 turned on. That’s the whole A/B.

• Workload: chat traffic (ShareGPT-style), which is decode-heavy and the most common production shape.

• What we measure: output throughput, acceptance length, latency percentiles, and cost per million tokens.

7.2 Renting the server on RunPod

RunPod rents GPU pods by the minute, which makes it ideal for a benchmark you will tear down afterward.

You can use any service, maybe Massedcompute or my personal favourite Modal.

Before you start the clock, here is roughly how long the whole experiment takes on a single 48 GB card, running baseline vs EAGLE3 on vLLM. Plan for about 45 to 60 minutes end to end. A single 48 GB GPU on RunPod runs about $0.40 to $0.90 per hour, so the whole run costs roughly $1.

The weights download only once and live on the /workspace persistent volume, so if you come back later that day you skip the ~10-minute download. Your two biggest sources of variance are HuggingFace download speed and whether you benchmark at --request-rate inf (fast) or a finite rate (slower).

A memory reality check (read this first)

Llama-3.1-8B in bf16 is about 16 GB of weights. Add the KV cache, activations, CUDA graphs, and the small EAGLE3 draft head (only a percent or two of the base model), and the whole thing still sits well under 48 GB. That means one card does the job, with no tensor parallelism and no quantization to fuss over.

Any of these single 48 GB GPUs on RunPod work:

• A40 (48 GB) is the cheapest option and the one this guide assumes by default.

• L40S, RTX 6000 Ada, or A6000 (48 GB) all work too and are a bit faster; pick whichever is available and affordable.

The commands in this guide all use -tp 1 (a single card). If you happen to be on a smaller 24 GB GPU like a 4090 or L4, add --quantization fp8 to fit, at a small quality cost.

Spin up the pod

RunPod pod configuration. Pick a single 48 GB card (A40 by default, or L40S / RTX 6000 Ada / A6000) with the GPU count set to 1. Then choose a recent CUDA/PyTorch template, a large enough container disk, a persistent volume for the model weights, and an exposed HTTP port for the inference server.

1. In the RunPod console, go to Pods → Deploy, select a 48 GB GPU (an A40 by default, or an L40S / RTX 6000 Ada / A6000), and leave the GPU count at 1. As the memory check above explains, 8B fits on one card, so there is no tensor parallelism to set up.

2. Choose a recent PyTorch / CUDA template (or a vLLM template from the RunPod Hub).

3. Give it a persistent volume of ~60 GB mounted at /workspace so the 8B weights survive a restart, and set container disk to ~40 GB.

4. Under Expose HTTP Ports, add 8000 (vLLM’s default).

5. Deploy, then connect via the web terminal or SSH.

One-time setup inside the pod

# Hugging Face access (Llama is a gated model, so accept the license on HF first)
pip install -U "huggingface_hub[cli]"
huggingface-cli login   # paste your HF token

# Keep all downloads on the persistent volume
export HF_HOME=/workspace/hf
mkdir -p "$HF_HOME"

Now you are ready to install an engine and serve the model.

7.3 Why EAGLE3 is the one to deploy

The theory above introduced a whole family of speculative decoding methods. In production today, EAGLE3 is the default choice, and the others are mostly of historical or niche interest. Here is the honest comparison so you know why we are skipping them rather than just being told to.

EAGLE3 wins on every axis that matters in production. It reaches the highest acceptance rates, so you get the most tokens per verification pass. It costs almost no extra VRAM, because it’s a small head rather than a second model. It has no tokenizer-matching problem. And, best of all for our purposes, pre-trained EAGLE3 heads are published for popular targets including Llama-3.1-8B, so there is nothing to train. Reported speedups reach roughly 2.4× on 8B (and up to ~4× on larger models) without changing the output.

What a “pre-trained EAGLE3 head” actually is

Every command in this guide points at two model names: the target I serve, and a draft head. It’s worth being clear about what that second name is, because it looks like a model but behaves very differently.

A draft head is a small checkpoint hosted on Hugging Face, named in the usual owner/name form.

Three things follow from that name, and they are the whole reason EAGLE3 is easy to deploy:

• It is tiny, not a second model. The head is only a percent or two of the target’s size. It’s the guesser, and the engine downloads it and bolts it onto the frozen target for you. That’s why speculation here costs almost no extra VRAM.

• It is tied to one specific target. An EAGLE3 head is trained against a single model and only works with that model, because it learned to predict that model’s internal features. EAGLE3-LLaMA3.1-Instruct-8B works with Llama-3.1-8B and nothing else. This is exactly why, when this guide moved from a 70B target to 8B, the draft head had to change too.

• “Pre-trained” means you skip the hard part. Someone already trained this head and published it, so you just reference it by name. If no head existed for your model, you would train one yourself (the EAGLE repo has recipes), but for Llama-3.1-8B one already ships, so there is nothing to do but point at it.

The workflow I’ll follow

I deploy on vLLM (7.4) and follow one simple loop so the comparison is honest:

1. Deploy the model without speculation and benchmark it. This is your reference: the plain tokens/sec and latency of the model as-is. (7.4 to deploy, 7.6 to benchmark.)

2. Re-deploy with EAGLE3 and run the identical benchmark. (7.4.)

3. Compare. The ratio of step-2 throughput to step-1 throughput is your speedup; the acceptance length tells you whether EAGLE3 is actually working. (7.7.)

7.4 vLLM implementation

Install vLLM:

pip install -U vllm

vLLM enables speculation through a single --speculative-config JSON blob (offline, the equivalent speculative_config={...} dict). Everything else about serving stays the same.

Baseline (no speculation)

This is the run every speedup is measured against. Deploy it and benchmark it first (7.6) before you touch EAGLE3.

VLLM_USE_V1=1 vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --seed 42 \
  -tp 1 \
  --gpu-memory-utilization 0.90

The server exposes an OpenAI-compatible API at

http://localhost:8000/v1.

Turning on EAGLE3

Now the same model, same flags, plus one --speculative-config. EAGLE3 attaches a small draft head that reuses the target model’s internal features:

VLLM_USE_V1=1 vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --seed 42 \
  -tp 1 \
  --speculative-config '{"model": "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", "num_speculative_tokens": 3, "method": "eagle3"}'

Each field in that JSON maps to a speculative-decoding idea:

• method: "eagle3" picks the drafting algorithm. It tells vLLM the draft head predicts the target’s internal features instead of raw tokens (the EAGLE idea from 6.3), which is what pushes acceptance so high.

• model is the pre-trained EAGLE3 draft head, the small network that does the guessing. The target model is never touched.

• num_speculative_tokens: 3 is how many tokens the head drafts ahead before the target verifies them. This is the lookahead K from 4, and it is the one dial you actually tune. Bigger is not automatically better: acceptance compounds, so each extra drafted token lands less often. Red Hat measured acceptance dropping off past two or three tokens on 8B (about 45% at 2 drafts, down to 28% at 4), so three is the sweet spot.

vLLM keeps this simple: you set how far to draft with this one number, and it sizes the draft tree (its depth, branching, and per-pass token budget) for you internally.

Two practical notes:

• On a single card there is nothing to shard, so the draft head runs alongside the target on the same GPU. On a multi-GPU box you would add "draft_tensor_parallel_size": 1 to keep the tiny head on one card.

• An alternative draft checkpoint is RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3.

Benchmark this exactly as you benchmarked the baseline (7.6), then compare.

Offline (Python) equivalent

If you would rather drive vLLM from Python instead of the server, the same config is a dict:

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    tensor_parallel_size=1,
    speculative_config={
        "method": "eagle3",
        "model": "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
        "num_speculative_tokens": 3,
    },
)

out = llm.generate(
    ["Explain speculative decoding in one paragraph."],
    SamplingParams(temperature=0, max_tokens=256),
)
print(out[0].outputs[0].text)

Smoke test

Check the server is answering before you benchmark it:

import openai

client = openai.Client(base_url="iframe.php?url=http%3A%2F%2F127.0.0.1%3A8000%2Fv1", api_key="None")
resp = client.chat.completions.create(
    model="meta-llama/Llama-3.1-8B-Instruct",
    messages=[{"role": "user", "content": "List 3 countries and their capitals."}],
    temperature=0,
    max_tokens=64,
)
print(resp.choices[0].message.content)

If this returns sensible text, your server is up and you can benchmark it.

7.5 How to benchmark it honestly

A speedup number means nothing without saying what you measured. Four metrics tell the whole story.

• Output throughput (tokens/sec). The headline. How many tokens the system produces per second under load. This is where EAGLE3 shows up biggest.

• Acceptance length, τ (and acceptance rate, α). From 4: τ is the average number of drafted tokens committed per verification pass, and α is the per-token acceptance probability. If τ is near 1, speculation is doing nothing; if it is 3–4, you are winning. This is the single best diagnostic, and vLLM reports it.

• Latency: TTFT and TPOT/ITL. Time to first token (TTFT) measures responsiveness; time per output token / inter-token latency (TPOT/ITL) measures how fast the stream flows. Report p50 and a tail percentile (vLLM prints p99).

• Cost per 1M tokens. The business case. Derived from throughput and the pod’s hourly price:

Why chat? Chat is decode-heavy: short prompts, long answers. That is exactly the regime where generation is memory-bound and speculation has spare compute to exploit. On prefill-heavy workloads (huge unique prompts, short answers) the win shrinks.

7.6 Running the benchmark

vLLM ships its own load generator. Run the baseline first, record the numbers, then re-launch the server with EAGLE3 and re-run the identical benchmark command.

The load test

Grab the ShareGPT dataset once:

wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json \
  -O /workspace/sharegpt.json

With the server running, point the benchmark at it:

vllm bench serve \
  --backend vllm \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --endpoint /v1/completions \
  --dataset-name sharegpt \
  --dataset-path /workspace/sharegpt.json \
  --num-prompts 500 \
  --request-rate inf

--request-rate inf floods the server to measure peak throughput; set a finite rate (e.g. --request-rate 8) and --max-concurrency to simulate steady traffic. The summary reports throughput, TTFT, and TPOT; vLLM's server logs/metrics report the acceptance length for the EAGLE3 run.

A minimal sanity script (optional)

For an intuition check before the heavy tooling, time tokens/sec from a single streamed request:

import time, openai

client = openai.Client(base_url="iframe.php?url=http%3A%2F%2F127.0.0.1%3A8000%2Fv1", api_key="None")
prompt = "Write a 400-word explanation of how a CPU cache works."

t0 = time.time()
n = 0
stream = client.chat.completions.create(
    model="meta-llama/Llama-3.1-8B-Instruct",
    messages=[{"role": "user", "content": prompt}],
    temperature=0, max_tokens=512, stream=True,
)
for chunk in stream:
    if chunk.choices[0].delta.content:
        n += 1
dt = time.time() - t0
print(f"{n} tokens in {dt:.2f}s = {n/dt:.1f} tok/s")

Run it against the baseline server, then against the EAGLE3 server. The difference in tok/s is your speedup (or, as 7.7 found, sometimes a slowdown) in one number. Use the real benchmark tools for anything you report.

7.7 Results

Here are the numbers from an actual run: Llama-3.1-8B on a single 48 GB Ada GPU, baseline vs EAGLE3, same prompts. I am reporting them exactly as they came out, because the honest result is more useful than a flattering one: on this setup, EAGLE3 was slower in both regimes. The tables show it, and the acceptance number explains it.

Table A: Single request (batch 1, greedy temperature=0), the latency regime where speculation is supposed to win

Table B: Under load (500 prompts, --request-rate inf, ~500 concurrent), the throughput regime

Table C: Acceptance (the diagnostic, from the EAGLE3 benchmark)

Table D: Latency under load

Table C is the whole story. Acceptance length came out at 1.81, barely above 1 and nowhere near the 3–4 you need for a win. When fewer than two of every three drafted tokens survive verification, the cost of drafting and verifying isn’t paid back, so throughput drops and latency rises. This is not a bug; it is speculation losing its bet, and the diagnostic tells you why, not just whether.

Three things drove acceptance down, and they are worth internalizing because they generalize:

1. A small target on a fast GPU is the hard case for speculation. An 8B model already decodes quickly (45 tok/s single-stream here), so the memory-bandwidth headroom that speculation trades compute for is small, and the draft head’s own forward pass eats most of any gain. EAGLE3’s economics are far better on large targets (70B), where each decode step is expensive and the draft is comparatively cheap. The published ~2.4× for 8B and ~4× for 70B come from bigger, more favorable setups than a single mid-range card.

2. The load benchmark ran non-greedy. vLLM’s benchmark no longer forces temperature=0, and stochastic sampling lowers acceptance because the target’s samples diverge from the draft’s guesses more often.

3. --request-rate inf saturated the GPU. At ~500 concurrent requests the card is compute-bound, so speculation’s extra work steals cycles from useful generation (7.8).

Could I make it win? Probably, but not for free. You would push acceptance up and stay in the latency regime: benchmark with --max-concurrency 1 --temperature 0, try num_speculative_tokens: 2 (shorter drafts fail less often), and confirm the draft head is actually attached (the acceptance line in the server log). The biggest lever, though, is the target size: run EAGLE3 against a 70B model and the same machinery usually comes out well ahead. On an 8B, expect the fight to be close, and sometimes a loss.

7.8 When speculation helps, and when it hurts

EAGLE3 is excellent, but it is still a bet, not a free switch, and 7.7 is proof. From the production literature (Red Hat’s gpt-oss and EAGLE3 studies, ) plus my own run:

It helps most when:

• The workload is decode-heavy (short prompts, long responses), like chat, reasoning, and agents.

• The target model is large (say 70B). Each decode step is expensive, so a cheap draft that lands even half its tokens is a clear win. This is the single biggest factor, and the one my 8B run was missing.

• Acceptance is high (τ≳3), which needs a well-matched draft head and, ideally, greedy or low-temperature decoding.

It hurts or does nothing when:

• The target is small and the GPU is fast (my case): baseline decoding is already quick, so there is little bandwidth headroom to buy back, and the draft’s own forward pass eats the gain.

• The GPU is already compute-saturated (high concurrency), so speculation’s extra drafting steals cycles from useful work. Some studies report gains even at ~200 concurrent requests, but my 500-concurrent run lost, so do not assume it, measure it.

• The workload is prefill-heavy (long unique prompts, little cache reuse), because the decode phase is too short to amortize the drafting.

This is why modern serving stacks treat speculation as a dynamic feature, on at light/medium load, automatically off when the queue saturates.

8. Takeaways

The idea:

• The problem is memory bandwidth. Generation is slow because the GPU re-reads the model’s entire weights for every token, leaving its math units idle. There is spare compute sitting there unused.

• The idea is guess-and-verify. A cheap drafter proposes several tokens, the expensive target verifies them all in one parallel pass, accepting a prefix and correcting the first mismatch.

• It is exact, not approximate. Modified rejection sampling (accept with min⁡(1,p/q), resample from the residual) guarantees the output matches the target distribution token for token. The draft only affects speed.

• The economics come down to α, τ, and K. High acceptance (α) and well-tuned, adaptive lookahead (K) push you toward the best-case speedup. Mismatched drafts and a bad K waste compute.

• Pick the drafter that fits your text. Use n-gram lookup for repetitive, input-grounded tasks with zero extra VRAM, Medusa for a single self-contained model, and EAGLE when you want the highest acceptance by drafting features instead of tokens.

• The real speedup depends on batch size. Fixed per-step overheads and, above all, GPU saturation erode the paper speedup: speculation is a low-load latency win, and its headroom collapses as verification comes to dominate each step.

The deployment:

• In production, deploy one method: EAGLE3. It ships with pre-trained heads for popular targets (including Llama-3.1-8B), and it is one flag in vLLM (--speculative-config), not a rewrite.

• 8B is a single-GPU job. One 48 GB card runs the target, KV cache, and EAGLE3 head at once, with no tensor parallelism or quantization.

• Read acceptance length first. τ tells you whether EAGLE3 is working before throughput does. My run’s τ=1.81 predicted the slowdown in one glance.

• It is conditional, and it can lose. A small model on a fast, saturated GPU is the hard case, and ours came out slower. The win grows with model size and acceptance. Measure your own A/B before you ship it.

9. What I didn’t cover

This post covered the foundations and a working deployment. Several powerful ideas were mentioned only in passing or left out entirely, worth a look if you want to go deeper:

• EAGLE-2 / EAGLE-3: dynamic candidate trees. Section 6.3 introduced EAGLE’s core idea, drafting at the hidden-feature level. The refinements build a dynamic candidate tree on top of those features, plus training-time tricks that push acceptance rates (and speedups) well past the vanilla EAGLE drafter.

• Self-speculation (layer-skipping / early-exit drafting). Turning the target model into its own drafter by skipping its middle layers, so you get drafts with no second model and no extra VRAM, sidestepping the dual-model cost from 5.4.

• Streaming / overlapped execution. Pipelining the draft and verify phases so the two models stop sitting idle in lockstep, closing the “bubble” time described in 5.1.

• Tree-based verification beyond Medusa (e.g. SpecInfer, Sequoia): smarter token trees, how to shape and size them, and how tree depth trades off against verification cost.

• Speculation at high batch sizes. How newer methods keep speculation useful in the throughput-bound regime instead of being toggled off the moment the GPU saturates.

• Designing and training the draft model. Distilling and aligning a drafter to maximize the acceptance rate α, and how draft–target alignment is actually engineered.

• Adjacent ways to break the autoregressive barrier, such as multi-token prediction (MTP), lookahead / Jacobi decoding, and how they relate to classic speculative decoding.

• Other serving stacks. How TensorRT-LLM and SGLang implement and auto-tune all of this (SGLang in particular has strong EAGLE3 support once you are on a compatible CUDA build).

That is where I am going to stop for today. There is clearly more here: dynamic EAGLE trees, self-speculation, streaming overlap, and now this semi-autoregressive line from DeepSeek. Maybe, if there is appetite for it, I will do a Part 2 that actually works through those. But this post is already long, and I would rather end it clean than stretch it thin.

If you liked the way this one reads, I have more in the same style on the way, and I think you will enjoy those just as much.

This is Part 1 of a longer run I am planning on LLM inference optimization, with more coming on CUDA, Triton.

Alongside the writing, I’m also building Audio Deep Learning projects and LLM projects, sharing and discussing them on LinkedIn and Twitter. If you’re someone curious about these topics, I’d love to connect with you all!

Mayank Pratap Singh

LinkedIn : www.linkedin.com/in/mayankpratapsingh022

Twitter/X : x.com/Mayank_022.

1. Introduction

2. How Do Models Learn Beyond Training?

3. Why RAG Isn’t Enough

4. Why Memory is Hard

   1. The Plasticity-Stability Tradeoff

   2. Catastrophic Forgetting

   3. The Write Problem

5. The Ideal Case

6. Quick Recap of Titans Architecture

7. Methodology

   1. The Neural Learning Module and the Nested Learning Paradigm

   2. Continuum Memory System (CMS)

   3. Self-Modifying Titans and the Hope Architecture

   4. Training Pipeline

   5. Inference-Time Optimization

   6. Muon and the Multi-Scale Momentum Optimizer (M3)

8. Results and Outcomes

9. Thoughts

10. Conclusion

1. Introduction

A while back, I wrote about Titans which is Google’s proposal for a model that could update its own memory during inference, without waiting for a retraining cycle. If you haven’t read that piece, I’d strongly recommend starting there, as what follows builds directly on it.

Titans was a significant result. It demonstrated that the rigid boundary between training and inference; something the field had treated as structurally fixed but it doesn’t have to be that way. By treating the hidden state of a neural module as a learnable projector rather than a static representation, Titans showed that a model could genuinely update what it knows at test time, not just retrieve from what it was trained on. The three-memory architecture : short-term attention, long-term updatable state, persistent fixed memory was a clean solution to a problem that RNNs, SSMs, and standard transformers had each failed to solve completely.

But there was a question the original Titans paper left open.

Titans showed you could learn at test time. What it didn’t fully answer was how to do this continuously across an ongoing stream of new information, over long horizons, without the memory degrading, overwriting itself, or losing coherence with what came before. A model that updates once at test time is impressive. A model that keeps updating, correctly and efficiently, across thousands of interactions is a much different and harder problem entirely.

The paper we are discussing today (also from Google) is the answer to that question. It introduces a framework for continuous inference-time learning a system where the model doesn’t just adapt to a single query, but genuinely accumulates knowledge across time, deciding what to remember, what to overwrite, and how to integrate new information into reasoning without destroying what already exists. This is done by combining self-modifying sequence model with the continuum memory system for a learning module, called Hope, showing promising results in language modeling, knowledge incorporation and reasoning.

A lot of jargons? Dont worry; That is what we are unpacking today.

2. How Do Models Learn Beyond Training?

Before we get into the architecture, we need to build a precise understanding of what it even means for a model to learn and how much of that learning can happen after the training run ends. This is not a philosophical question. It has direct architectural consequences.

2.1 The Spectrum

Most discussions treat model learning as a binary, i.e. either the model is training, or it is not; this framing is too coarse/sparse. The more useful way to think about it is as a spectrum, defined by a single axis: how much is the model allowed to change in response to new information?

At one extreme, you have a fully frozen model with weights fixed, no updates of any kind, every query answered from a static snapshot of what was learned during pretraining. At the other extreme, you have a fully dynamic model; one that continuously updates its own parameters in real time, integrating every new piece of information as it arrives.

In between these two poles live three distinct learning regimes, each with a different answer to how much change is permitted and where that change is allowed to happen.

2.2 In-Context Learning (Change Without Weight Updates)

The softest form of post-training adaptation is in-context learning. The model weights never change. Instead, the model conditions its behaviour on examples, instructions, or information provided directly in the context window; hence, effectively using the attention mechanism as a temporary, ephemeral form of memory.

Formally, given a context C=(c1​,c2​,…,ck​) of examples and a query q, in-context learning computes:

The model hasn’t learned anything in the weight-update sense. It has simply conditioned on more information. The moment the context window is cleared, everything is gone and the model returns to its prior state as if nothing happened.

This is powerful as Transformers are remarkably good at in-context learning, and it requires zero infrastructure beyond the model itself. But it has a hard ceiling: the context window. Once k exceeds the maximum sequence length, you cannot fit more information in. The model’s ability to adapt is bounded by how much it can see at once, not by how much it has experienced over time.

In-context learning is adaptation without memory. It is stateless across queries and boundless in its forgetting.

2.3 Continual Learning (Change Across Time, Without Forgetting)

Continual learning asks for something harder. The model should be able to learn from a sequential stream of tasks or data over time hence accumulating knowledge rather than replacing it. Each new piece of information should be integrated without destroying what came before.

The core challenge is catastrophic forgetting; which happens when a neural network is updated on new data via gradient descent, the weight updates that improve performance on the new task systematically degrade performance on old tasks. The gradients don’t know what the old weights were protecting. They just follow the loss surface of the current objective, overwriting whatever was previously encoded.

Formally, suppose a model with parameters θ is trained sequentially on tasks T1,T2,…,Tn. After training on Ti​, the parameters shift from θi−1​ to θi​. The catastrophic forgetting problem is:

Performance on old tasks Tj​ degrades severely after training on new ones. The model is plastic as it learns new things; but it is not stable as it cannot retain old things simultaneously/to a large extent.

This is the plasticity-stability tradeoff, and it is the central unsolved tension in continual learning. A model that updates too readily is plastic but unstable as it learns fast and forgets fast. A model that updates too conservatively is stable but rigid as it retains old knowledge but resists integrating new information. These two requirements pull in opposite directions, and there is no free solution that achieves both simultaneously.

2.4 Inference-Time Learning (Change During the Forward Pass)

Inference-time learning is the most architecturally radical of the three. The model updates its own parameters during inference, not before, not after, but as part of the forward pass itself. Each query doesn’t just consume the model’s knowledge; it potentially modifies it.

This sounds paradoxical as inference is supposed to be a read operation as you can pass an input forward through fixed weights and get an output. Making it also a write operation where the forward pass updates the very weights it is passing through requires rethinking what a forward pass even means.

Architecturally, this requires two things. First, a subset of the model’s parameters must be designated as updatable at inference time (not the full model), which would be prohibitively expensive, but a structured, bounded module whose updates are fast and local. Second, an update rule must be defined that operates at inference speed; not full backpropagation through the entire model, but a lightweight approximation that captures the signal of new information without the cost of a full training step.

The Titans architecture from the previous article was an early answer to both requirements. The neural memory module which is a small MLP whose weights serve as the memory state can be updated during inference via a self-supervised objective. The update is local, fast, and bounded. It doesn’t touch the full model. But as noted in the introduction, it handles a single update at test time. What happens when updates need to keep coming?

2.5 Three Points on the Same Spectrum

These three regimes are not competing paradigms. They are different answers to the same question : how much should a model be allowed to change in response to new information? when placed at different points on a single spectrum; it looks something like this:

The paper this article covers lives at the inference-time learning end of this spectrum; but it extends it further than Titans did, toward something closer to true continual learning. The model doesn’t just update once at test time. It keeps updating, across a stream of information, with an explicit mechanism for deciding what to retain and what to overwrite.

That is the regime the next sections build toward.

3. Why RAG Isn’t Enough

If the spectrum in the previous section is about how much a model is allowed to change, RAG occupies a curious position as it looks like it solves the memory problem, without actually touching the model at all.

3.1 What RAG Actually Does

Retrieval-Augmented Generation works by treating an external corpus as a lookup table. Given a query q, a retriever selects the top-k most relevant documents D=(d1,d2,…,dk) from a large external store, and these are concatenated into the model’s context:

This should look familiar as it is structurally identical to the in-context learning formulation from the previous section. RAG is not a new point on the spectrum. It is in-context learning with an automated retrieval step bolted on front. The model weights never change. The retrieved documents are not learned because they are simply read, once, for this query, and discarded the moment the next query begins.

This is why RAG feels like memory; but isn’t. It gives the appearance of a model that “knows” things beyond its training data, but the knowing is entirely externally held in a vector database, not in the model itself.

3.2 The Structural Limitations

Four limitations follow directly from this design.

Context window bounds. Retrieved documents still have to fit inside the model’s context window, alongside the query and the conversation history. As the corpus grows, the retriever has to be more selective; because the value of k stays roughly fixed even as the available information grows without bound.

Retrieval quality ceiling. The entire system is bottlenecked by the retriever’s ability to find the right k documents. If the retriever misses the relevant document because of an imperfect embedding, an ambiguous query, or information that wasn’t indexed well; then the model never sees it, no matter how capable the model itself is. The generator cannot compensate for what the retriever failed to surface.

No true integration. This is the subtlest limitation. Even when retrieval succeeds and the right document is in context, the model is not integrating that information into its world model. It is conditioning on it, for this one forward pass, the way it would condition on any other text in the prompt. The next query starts from zero again; hence, nothing about how the model reasons has changed, only what it was shown changed (abstract context).

Stateless across queries. Every query is an independent retrieval-and-generate event. There is no accumulation. If a user mentions an important fact in turn one, and the conversation moves past it in the context window by turn fifty, that fact has to be re-retrieved or it is gone as the model never actually held it in the first place. (arguably memory could be maintained; but it is just another layer of retrieval).

3.3 The Deeper Problem : Memory as Storage vs. Memory as Structure

These four limitations all trace back to a single design choice: RAG treats memory as something external to the model like a filing cabinet you consult, rather than something that shapes how you think.

This is worth sitting with, because it is easy to conflate “has access to information” with “remembers” and they are not the same thing. A filing cabinet gives you access to information. It does not change how you reason about the world. You can look something up, use it for the task at hand, and put it back, completely unchanged by having read it.

Human memory does not work this way, because when you learn something like a fact, a skill, a person’s name then it does not sit in a separate compartment that you consult on demand. It becomes part of the substrate you think with; hence it changes which associations come to mind, which patterns you notice, which conclusions feel obvious versus surprising.

3.4 The Question This Sets Up

This reframes the problem we actually need to solve. The question is not “how do we give a model access to more information than fits in its context window”; as that is the question RAG answers, reasonably well.

The question is: what would it mean for a model to actually remember something; basically not retrieve it from an external store and condition on it for one forward pass, but have it genuinely reshape how the model processes every subsequent query, the way a fact you learned last year quietly shapes a judgment you make today, without you consciously retrieving it at all? That is a question about architecture, not retrieval infrastructure.

4. Why Memory is Hard

We have established what memory should look like in principle: structure that shapes reasoning, not storage you consult on demand. Now we need to confront why building this is genuinely difficult, not just an engineering inconvenience waiting to be solved with more compute.

4.1 The Plasticity-Stability Tradeoff

Any system that learns continuously faces a tension between two properties it needs simultaneously.

Plasticity is the ability to change in response to new information. A model needs this, or it can never learn anything beyond what it saw during pretraining.

Stability is the ability to retain what was already learned. A model needs this too, or every new piece of information would simply erase the last one (remember catastrophic forgetting).

The problem is that these two properties pull directly against each other. Consider a model with parameters θ, currently encoding some knowledge K_old​, that receives a new piece of information requiring an update Δθ:

If Δθ is large enough to fully encode the new information, it is also large enough to overwrite whatever values in θ_old were responsible for K_old​. There is no inherent mechanism in this update rule that protects old knowledge while incorporating new knowledge. The same parameters that hold K_old​​ are the parameters being modified to hold the new information.

You can dampen Δθ to reduce the risk of overwriting. But a smaller Δθ also means less of the new information actually gets encoded. You have not solved the tradeoff. You have just chosen a different point on it. Push toward plasticity and you risk forgetting. Push toward stability and you risk never learning. There is no parameter setting that gives you both for free.

4.2 Catastrophic Forgetting

The plasticity-stability tradeoff is the abstract version of the problem. Catastrophic forgetting is what it looks like concretely, in a trained neural network, when gradient descent is the update mechanism.

Suppose a model has been trained on task T₁ and has converged to a region of parameter space θ₁​ that performs well on it. Now it receives new data from task T₂ and is updated via gradient descent:

This update is computed entirely with respect to T₁​. The gradient ∇_θ​L_T2​​ has no term, no penalty, no awareness whatsoever of L_T1​​. It simply descends the loss surface of the current task. If the region of parameter space that minimizes L_T2​​ is far from the region that minimized L_T1​​, the update walks the model directly away from where it needs to be for the old task.

The result is not gentle decay. It is sudden and severe. Performance on T₁​ can collapse after just a few gradient steps on T₂, even though nothing in the training procedure explicitly tried to destroy that knowledge. It is simply a side effect of optimizing for the wrong objective using parameters that were doing double duty.

This is the mechanism, not just the symptom. Any system that wants to learn continuously without forgetting needs to either protect specific parameters from being overwritten, route new information to a region of parameter space that doesn’t interfere with old knowledge, or find some other structural way to prevent the gradient from old and new objectives from cancelling each other out.

4.3 The Write Problem

Even if you solve the forgetting problem, you are left with a second, equally hard question: what is actually worth remembering?

Not all information that arrives during inference deserves to be written into memory. Some of it is noise. Some of it is redundant with what is already known. Some of it is relevant only to the current query and has no business influencing future reasoning. A memory system that writes everything indiscriminately will fill up fast, with most of its capacity spent on information that didn’t need to be retained in the first place.

This is a selection problem, and it is genuinely hard for a structural reason: deciding whether something is worth remembering often requires knowing whether it will be useful in the future, which you cannot know at the moment you encounter it. A fact that seems irrelevant today might be exactly what a future query needs. A fact that seems important today might never come up again.

A good memory system needs some signal, some proxy, for deciding what to write. This could be based on how surprising the new information is relative to what the model already expects, how often similar information has proven useful in the past, or how confidently the information can be verified. Whatever the signal, the system needs one, because the alternative, writing everything, is not actually a memory system. It is just an ever-growing, ever more diluted version of the context window problem RAG already has.

4.4 Where This Leaves Us

Put these three problems together and the difficulty of the task comes into focus. You need a system that is plastic enough to integrate genuinely new information, stable enough not to destroy what it already knows, and selective enough not to drown itself in irrelevant updates. None of these three requirements can be solved in isolation. A solution to plasticity that ignores stability gives you catastrophic forgetting. A solution to stability that ignores plasticity gives you a frozen model. A solution to both that ignores selectivity gives you a memory system that degrades under its own weight.

This is the real shape of the problem the next section’s ideal case has to grapple with. Not “can a model update at test time”, which Titans already showed was possible, but “can a model update continuously, indefinitely, while staying coherent.” That is a substantially harder bar to clear.

5. The Ideal Case

Let’s start with an unconstrained dream, the similar way we did for Fast BLT (our previous article).

The perfect continuously-learning memory system would do four things at once. It would update instantly the moment new information arrives, with no separate training step and no delay. It would never forget anything it had previously learned, no matter how much new information arrived afterward. It would know, automatically, what was worth remembering and what wasn’t, writing only the information that would matter later. And it would integrate everything it learned into how it reasons, not store it as a separate fact to be consulted, but let it quietly reshape every subsequent inference.

This is the ideal. Instant plasticity, perfect stability, perfect selectivity, full integration. And it is unachievable, for reasons that go deeper than current engineering limitations.

5.1 Why the Ideal Breaks Down

The reason traces directly back to section 4.1. Plasticity and stability are not two independent dials you can both turn to maximum. They are coupled through the same mechanism, the same set of parameters, the same update rule.

Recall the update equation:

For the system to be instantly and fully plastic, Δθ must be large enough to fully encode any new piece of information, regardless of how it relates to what is already stored. For the system to be perfectly stable, Δθ must never perturb the components of θ_old​ that encode previously learned knowledge. These two requirements can only both hold if every new piece of information happens to map onto a completely disjoint region of parameter space from everything learned before it, indefinitely, forever. There is no mechanism that guarantees this. Parameter space is finite. Information is not.

This is structurally the same kind of impossibility we encountered with the autoregressive dependency chain in the Fast BLT article. There, the constraint was that you cannot compute x_t​ before x_t−1​ exists, because the dependency is a logical fact about the factorization of the joint distribution, not an engineering limitation. Here, the constraint is that you cannot guarantee perfect plasticity and perfect stability simultaneously, because the parameters doing the remembering are the same parameters being asked to change. It is not a question of better hardware or more compute. It is a property of what it means to update a finite system indefinitely.

The selectivity requirement compounds this. Even if you could solve plasticity and stability, deciding what is worth remembering requires information about the future that is not available at write time, as we discussed in section 4.3. A perfect selector would need to know, in advance, which pieces of information will matter for queries that haven’t been asked yet. This is not a hard engineering problem. It is asking for foresight the system structurally cannot have.

5.2 The Real Question

Given that the ideal is unreachable, the question that matters is not “how do we build the perfect memory system.” It is: how close can we get, and what are we willing to give up to get there?

This reframes the entire design space. Just as Fast BLT was not trying to reach the impossible single-pass generation ideal but instead tracing a speed-quality frontier as close to it as possible, a continuous learning system is not trying to reach perfect plasticity and perfect stability simultaneously. It is trying to trace a frontier between them, recovering as much of both as the architecture allows.

We can write this formally as a constrained optimization, mirroring the structure from our Fast BLT article. Let P denote plasticity, how readily the system integrates new information, and let S denote stability, how well it retains old information. The achievable frontier is:

Every architectural choice, how memory is represented, how updates are computed, what gets written and what gets protected, places the system at some point on or below this frontier. The goal of any concrete design is to push as close to the frontier as possible, knowing the frontier itself never touches the ideal corner of perfect P and perfect S simultaneously.

5.3 What Determines How Close You Can Get

Two things determine where on this frontier a real system can land.

How the architecture separates what changes from what stays fixed. If updates are forced to touch the same parameters that encode everything else, plasticity and stability fight directly, as in section 4.1’s basic update equation. If the architecture instead routes new information into a structurally separate component, one that can change freely without disturbing a more protected store of existing knowledge, the conflict softens considerably. This is the architectural lever.

How selective the write mechanism is. A system that filters aggressively, writing only information that clears some bar of novelty or usefulness, makes far better use of its limited capacity for change than one that writes indiscriminately. This does not solve the foresight problem from section 4.3, but a good proxy signal, something correlated with future usefulness even if not a perfect predictor of it, can meaningfully improve where the system sits on the frontier.

These two levers, architectural separation of memory components, and selective writing, are exactly what the methodology section is going to dig into. The paper’s answer to “how close can we get to the ideal” lives in how it structures the neural language module, how the continuous memory system is organized, and how the self-modifying mechanism decides what to write and what to protect.

That is where we are headed next.

6. Quick Recap of Titans Architecture

If you’ve read the original Titans article, this section will feel familiar.

If you haven’t, here is quick summary you need to follow the rest of this piece.

6.1 The Core Idea : Memory as a Neural Module

Titans starts from a simple but consequential reframing: instead of treating memory as a fixed-size vector or buffer, treat it as the weights of a small neural network. The hidden state isn’t a static representation that gets overwritten and decayed the way an RNN’s hidden state does. It is a projector, a learnable function, with the capacity to encode high-resolution information precisely because it has parameters that can be optimized.

This single design choice resolves a tension we touched on in section 2. RNNs and SSMs achieve constant space complexity but suffer from lossy compression, since their fixed-size hidden state has to discard information to make room for new input. Transformers achieve high-quality representation through attention but pay for it with quadratic complexity, since every token attends to every other token. Titans’ neural memory module gets the constant space of an RNN while approximating the representational capacity of a much larger system, because the module’s weights, not just its activations, are doing the work of encoding information.

6.2 Three Memories, Three Roles

Titans separates memory into three distinct components, each with a different update rule:

Short-term memory is standard attention, operating over a limited window of recent tokens. It is precise and high-quality, but quadratic in cost, so it is kept short by design.

Long-term memory is the neural module described above. Its weights are updated during inference, via a self-supervised objective, allowing the model to incorporate information from far outside the attention window without paying attention’s quadratic cost.

Persistent memory is also neural, also low-dimensional, but never updated at inference time. It encodes stable, task-general knowledge that should not drift, regardless of what happens during a particular inference session.

The separation between in-context and in-weights memory follows directly from this structure. In-context memory lives in the attention window, the same mechanism behind standard in-context learning from section 2.2, bounded and ephemeral. In-weights memory lives in the long-term module’s parameters, persisting across the session and shaped by an actual update rule rather than just conditioning.

6.3 The Test-Time Update Mechanism

The mechanism that makes long-term memory updatable at inference time is a gradient step, computed on the fly, using a self-supervised objective often referred to as a surprise signal. When the model encounters information that its current memory module predicts poorly, the resulting error drives an update to the memory module’s weights:

​This is, structurally, the same update equation we wrote in section 4.1 and 5.1, applied specifically to the long-term memory parameters rather than the full model. The persistent memory and the bulk of the model’s other weights are left untouched. Only the long-term memory module is plastic at inference time, which is exactly the kind of architectural separation we flagged in section 5.3 as a lever for softening the plasticity-stability tradeoff.

6.4 Where Titans Leaves Off?

Titans demonstrated something genuinely new: a model can update its own weights mid-inference, using a fast, local, self-supervised signal, without touching the bulk of the network. That is inference-time learning, concretely realized, for the first time at this level of architectural cleanliness.

But the original paper was largely concerned with a single inference session. It showed the update mechanism works, that it improves long-range understanding, that it can be parallelized for training. What it did not fully address was what happens when the updates don’t stop. What happens after the hundredth update to the same memory module, the thousandth, across sessions, across days, across a genuinely continuous stream of new information. Does the surprise-driven update rule remain well-behaved indefinitely, or does it start to drift, saturate, or interfere with itself the way the gradient updates in section 4.2 interfered across tasks?

Titans showed you could learn at test time, but the question this new paper takes on is how to do that continuously, and how to do it efficiently enough that it remains practical at scale. That is exactly where the methodology section picks up.

7. Methodology

Now that we have revised Titan and under the capabilites and advancements we need on top of it to realize the idea of a strong continuual learning paradigm; let’s start with the methdology proposed in paper.

7.1 The Neural Learning Module and the Nested Learning Paradigm

Before this paper, the natural way to think about a deep network was as a stack of layers, each layer a fixed computational block, with a separate optimizer sitting outside the architecture, adjusting weights via backpropagation. Nested Learning (NL) asks you to discard that picture entirely.

The Reframe : Optimizers Are Memory Modules

The paper’s starting observation is unusual: a gradient-based optimizer, the thing that updates a model’s weights, is itself a kind of memory module. Consider standard SGD with momentum. At every step, it takes the current gradient, combines it with a running average of past gradients, and uses that combination to update the weights:

Look closely at the first equation. It is an associative memory update. The momentum term m_t​ is compressing a stream of gradients, the “context flow” of the optimization process, into a fixed-size representation, the same way the long-term memory module in Titans compresses a stream of tokens into its weights. Adam does the same thing with two separate compression terms, one for the gradient mean and one for its variance.

This is the central reframe of Nested Learning: training a neural network and updating a memory module are not two different things happening at different stages. They are the same operation, an associative memory compressing a context flow, applied at different places in the system and at different speeds.

Multi-Level Optimization

Once you accept that framing, a deep learning model stops looking like a flat stack of layers with one optimizer outside it, and starts looking like a set of nested optimization problems, each with its own context flow, each updating at its own frequency.

Formally, NL represents a model as a collection of levels ℓ=1,2,…,L, where each level has its own parameters θ(ℓ), its own local objective L(ℓ), and its own update frequency f(ℓ), the rate at which that level’s parameters are allowed to change:

where Ct(ℓ)​ is the context flow that level ℓ sees at time t, the stream of information it is being asked to compress. Crucially, an outer level’s parameters can themselves be the thing that defines an inner level’s context flow or objective, which is why these problems are described as nested rather than merely parallel. The outer loop shapes the conditions under which the inner loop operates; the inner loop’s behaviour, in turn, contributes to the signal the outer loop optimizes against.

This single formalism is expressive enough to describe a standard transformer. In a transformer, the attention block has, in this framing, an effectively infinite update frequency, f→∞, since it recomputes its in-context behaviour fresh on every forward pass with no persistence across calls. The MLP blocks, by contrast, have update frequency f→0, since their weights are frozen entirely after pretraining and never change again at inference time. The paper’s striking observation is that the entire spectrum we mapped out in section 2, from in-context learning to continual learning, is really just a description of where on this frequency axis different components of a model happen to sit.

Why Nesting? (Navigating Plasticity and Stability Through Frequency)

This is where the nested structure earns its purpose. Recall the plasticity-stability tradeoff from section 4.1: a single set of parameters cannot be simultaneously maximally plastic and maximally stable, because the same update that integrates new information is the update that risks overwriting old information.

Nested Learning’s answer is not to find a clever update rule that magically achieves both. It is to refuse the premise that there has to be a single set of parameters and a single update rate at all. Instead, the system is decomposed into multiple levels, each with its own frequency, and the tradeoff is resolved not within any single level but across the hierarchy of levels.

A fast inner loop, updating frequently, gives you inference-time adaptability. It can absorb new information immediately, the way the attention mechanism absorbs the current context, or the way Titans’ long-term memory module absorbs a surprising new fact mid-session. Because it updates fast, it is plastic.

A slow outer loop, updating only during training, gives you stable long-term structure. It changes rarely, and when it does change, it does so based on aggregated signal across enormous amounts of data, not any single new fact. Because it updates slowly, it is stable.

Neither loop, on its own, solves the tradeoff. But together, they don’t have to. The fast loop is allowed to be unstable because the slow loop isn’t relying on it to hold permanent knowledge. The slow loop is allowed to be rigid because the fast loop is handling everything that needs immediate adaptation. The tradeoff from section 4.1 is not eliminated. It is distributed across timescales, with each level responsible only for the part of the tradeoff it is suited to handle.

This is the structural lever we flagged in section 5.3 as the key to pushing closer to the frontier, separating what changes from what stays fixed. Nested Learning generalizes that separation from Titans’ three discrete memory tiers into a full spectrum of update frequencies, and it is this generalization that the rest of the methodology, the Continuum Memory System and the self-modifying Hope architecture, builds directly on top of.

7.2 Continuum Memory System (CMS)

The traditional view of memory in a transformer is binary: there is short-term memory, the attention mechanism operating over recent tokens, and long-term memory, the MLP’s frozen weights encoding everything learned during pretraining. This binary framing maps cleanly onto the two extreme frequencies from section 7.1; which is attention at f→∞, MLPs at f→0; hence nothing in between.

The Continuum Memory System (CMS) discards this binary entirely as Memory is not two discrete tiers. It is a spectrum of frequencies, each responsible for a different timescale of knowledge, each updating at a rate appropriate to what it is storing; hence we want to operate in a continuous space.

7.2.1 Memory as a Chain of Frequencies

Formally, CMS replaces the single MLP block in a transformer with a chain of kk k MLP blocks:

where f1​>f2​>⋯>fk​ and each block updates at a strictly lower frequency than the one before it. Given input x={x1,…,xT}, the output of the chain is computed sequentially:

Each block ℓ has its own parameters θ^(fℓ​), updated every C^(ℓ) steps (its chunk size), the number of tokens it processes before incorporating what it has seen into its weights:

The block updates in a batch; accumulating gradients over its chunk window and then applying them all at once. Between updates, its parameters are frozen, acting as a static memory for that timescale.

The standard transformer MLP is a degenerate case of this formulation: k=1, frequency of update zero, chunk size effectively infinite. CMS generalises this into a full spectrum by stacking blocks at different C^(ℓ) values — with the smallest chunk size at the fastest block and the largest at the slowest.

7.2.2 What Each Frequency Level Stores

This is the key intuition behind CMS: the chunk size C^(ℓ) determines the scope of what a block can know about.

A block with C^(ℓ)=1 updates on every single token. It is highly plastic as it responds immediately to every new piece of information; but its parameters reflect only the most recent token’s context. It is volatile by design. This is the fast, adaptive end of the spectrum.

A block with C^(ℓ)=T represents the entire sequence length and updates only once after processing everything. Its parameters compress the entire context into a single set of weights. It is stable and slow to change, but what it stores is a global summary of everything seen so far. This is the persistent end of the spectrum.

Between these extremes, blocks at intermediate chunk sizes store knowledge at intermediate timescales. A block with C^(ℓ)=128 has seen and integrated 128 tokens’ worth of context before updating its knowledge is local enough to be responsive, global enough to be coherent across a short segment.

Formally, each block’s parameters θ^(fℓ​) are a compressed representation of their context window:

This gives CMS a property the binary long-term/short-term model structurally cannot have: multi-timescale knowledge recovery. If a fast block forgets something because its chunk window has moved on, the information may still be encoded in a slower block whose window hasn’t expired yet. Knowledge can be partially recovered from a lower-frequency layer even after the fast layers have moved past it; which is something that is impossible in a system where the only options are ephemeral attention or frozen weights.

7.2.3 The Write Mechanism

Writing into a CMS block is gradient-based memory update, applied locally within each block’s chunk window. For a block at frequency f^ℓ​ processing chunk i, the write step is:

The gradient here is computed with respect to the block’s own local objective in language modeling, this is the next-token prediction loss evaluated on the chunk. Each block is, in the Nested Learning framing from section 7.1, an associative memory compressing its own context flow into its parameters. The surprise signal is how much the block’s current parameters fail to predict the incoming token which drives the update.

Crucially, this write is bounded in scope. The gradient only flows through the current block’s parameters, not the full model. The plasticity is local. The slower blocks below it are not touched. This is the architectural separation we identified in section 5.3 as the key lever for navigating the plasticity-stability tradeoff as fast blocks are free to be plastic because slow blocks are structurally isolated from their updates.

7.2.4 The Read Mechanism

Reading from CMS is simply the forward pass through the chain where input flows through MLP(f₁​) first, then MLP(f₂), and so on. Each block transforms the representation based on its current parameter state, which encodes whatever it has learned from its chunk history.

This is fundamentally different from attention-based retrieval. In attention, reading involves computing similarity scores between a query and every key in the context, then retrieving a weighted combination of values. The retrieved content is always a direct function of the current query and the stored keys and values.

In CMS, reading is a deterministic forward pass through a learned function. There are no explicit keys, no similarity scores, no direct addressing. The block’s parameters have already compressed the relevant patterns from its context window as the read operation simply evaluates that compression on the current input. It is implicit retrieval rather than explicit lookup.

The tradeoff is expressivity versus efficiency. Attention retrieval is content-addressable which you can retrieve exactly what you ask for, from anywhere in the context. CMS retrieval is structural as whatever you get is shaped by what the block has learned to encode, which is a compressed summary of the context, not a direct lookup into it. For long-range information that matters across many tokens, CMS’s compression is more efficient. For precise retrieval of specific earlier tokens, attention remains sharper.

7.2.5 Variants : How Levels Connect?

CMS is not a single architecture but a family, distinguished by how knowledge flows between frequency levels.

Nested CMS is the initial state of block ℓ+1 is meta-learned from block ℓ. Each block initializes from a stable learned prior rather than from zero:

After each chunk is processed and the block resets, it reinitializes to this meta-learned starting point rather than cold-starting. This gives higher-frequency blocks a stable prior that reflects everything the lower-frequency levels have consolidated into an approximation of how the human brain’s offline consolidation during sleep replays (in hippocampus) and reorganizes recent memories into longer-term cortical storage.

Sequential CMS are blocks that are chained but independent/isolated during initialization. The output of one block feeds into the next, but there is no meta-learning of initial states across levels. This is simpler to implement, but at the cost of the inter-level knowledge transfer that nested CMS provides.

Both variants contrast sharply with the binary long-term/short-term memory model. There is no hard boundary between ephemeral and persistent memory as only a gradient of update frequencies, each responsible for its own timescale, collectively providing the multi-scale memory management that the plasticity-stability tradeoff requires.

7.3 Self-Modifying Titans and the Hope Architecture

The Continuum Memory System from section 7.2 solved the representation problem: how to store knowledge across multiple timescales without forcing everything into a binary long-term/short-term split. But it left a second problem untouched: what drives the updates? Every CMS block updates via gradient descent on the next-token prediction loss, which is a fixed, hand-designed update rule. The model learns what to store, but the mechanism by which it stores things is prescribed externally, the same gradient descent formula applied uniformly across every block, at every frequency level.

The Self-Modifying Learning Module asks a more radical question: what if the model learned its own update rule, rather than having one imposed on it?

7.3.1 The Core Idea : Learning to Update

Recall from section 7.1 that NL frames a gradient-based optimizer as an associative memory module, something that compresses a context flow (the gradient history) into a fixed representation (the momentum terms or second moments) via gradient descent on a local objective. The update rule itself is therefore just another learned compression, operating at a higher level than the network weights it is updating.

This observation opens a design space that standard deep learning leaves completely unexplored: if the update rule is a compression, and if more expressive compressions are better, then the update rule itself can be parameterized and learned, in the same way we learn the network weights it applies to.

The Self-Modifying Learning Module (SMLM) operationalizes this. It is a sequence model that maintains two sets of parameters:

• Outer parameters (Θ), updated only during training via standard backpropagation on the task loss. These are slow, stable, and encode the model’s long-run structural knowledge.

• Inner parameters (ϕ), updated during inference according to a rule that is itself parameterized by Θ. These are fast, plastic, and encode the model’s short-run adaptation to the current context.

The key property is that the inner update rule is not fixed. It is a function gΘ​ learned by the outer loop:

where gΘ​ is a learned function parameterized by the outer weights Θ, taking the current inner parameters, the current input, and the current gradient as inputs, and outputting a parameter update. The outer loop trains Θ such that the resulting gΘ​ produces good inner updates across a distribution of tasks and contexts. The inner loop applies gΘ​ at inference time to actually adapt ϕ.

This is what “learning to modify itself” means concretely. The model is not just updating its weights at inference time. It is doing so according to an update rule that was itself optimized, end to end, to be a good update rule for this class of problems.

7.3.2 How a Self-Modifying Layer Works

To make this concrete, consider how a single SMLM layer processes an input sequence x1,x2,…,xT​.

At each step tt t, the layer maintains its current inner parameters ϕt​. The forward pass has two components that happen simultaneously.

The output computation applies ϕt​ to the current input:

This produces the layer’s output for this token, using whatever the inner parameters currently encode.

The parameter update applies the learned update rule to generate new inner parameters for the next step:

The gradient ∇ϕ​L here is computed with respect to a local self-supervised objective, something like predicting the current token given recent context, the same surprise signal from Titans’ memory update in section 6.3. The output y_t​ and the updated parameters ϕt+1​ are both produced in the same forward step. The layer has processed the input and updated itself simultaneously.

7.3.3 Why This Is Different from Fine-Tuning

The comparison to fine-tuning is natural but misleading. Fine-tuning updates a model’s weights on new data using backpropagation through the full network on a labeled objective, with an externally chosen optimizer running for multiple gradient steps over the full parameter space. It is a process that happens offline, between deployments, not during inference.

The SMLM update differs on every one of these dimensions.

Scope. The inner update touches only ϕ, the inner parameters of this specific layer. The outer parameters Θ are completely frozen during inference. The rest of the model is not involved. There is no risk of one layer’s adaptation cascading into unintended changes elsewhere in the network.

Speed. The update gΘ​(⋅) is a single forward pass through a learned function, not iterative gradient descent over multiple steps. It happens in the same computational budget as the forward pass itself.

The update rule itself. In fine-tuning, the optimizer is SGD or Adam, chosen by the engineer, fixed, the same for every problem. In the SMLM, the optimizer is gΘ​, learned end to end to be a good update rule for this specific class of tasks. This is the deeper difference: the model is not just updating its weights with an external tool. It has internalized the update process as part of its own learned behavior.

7.3.4 Gradient Flow Through a Self-Modifying Layer

This section requires careful treatment because the gradient flow is non-trivial in a way that matters for how the system is trained.

In a standard layer, the gradient of the loss with respect to the layer’s input flows through the layer’s fixed parameters. One backward pass, one gradient, one update.

In an SMLM layer, there are two pathways through which gradients flow:

Pathway 1: Through the output. The standard path. Given output yt​=fϕt​​(xt​), the outer training loss L_outer​ backpropagates through y_t​ to reach ϕ_t​ and Θ:

​Pathway 2: Through the update rule. The less obvious path. Because ϕt+1​ depends on Θ through the learned update rule gΘ​, and because y_t+1​ depends on ϕ_t+1​, the outer loss at step t+1 also flows back through the update rule itself to reach Θ:

​where ∂Θ∂ϕt+1​​ requires differentiating through gΘ​, which itself is a neural network with parameters Θ. The total gradient on Θ is the sum across both pathways, and across all timesteps:

This is what makes the outer loop training expensive and why it requires careful implementation. Computing path 2 requires differentiating through the update rule gΘ​, which means keeping the computation graph of the inner loop alive during the outer loop backward pass, a memory cost analogous to unrolling an RNN through time.

In practice, the paper truncates this to avoid the full unrolling cost, computing only a fixed number of inner steps before detaching the gradient, the same truncated backpropagation through time trick used in training recurrent networks. The truncation introduces bias into the outer gradient estimate but keeps the training computationally tractable.

7.3.5 Hope : Putting It Together

Hope (Higher-Order Parameter Evolution) is the concrete architecture that combines the Self-Modifying Learning Module with the Continuum Memory System into a single, unified sequence model.

The architecture replaces the standard transformer block with a Hope block, which consists of:

• A standard attention layer handling short-term in-context memory, unchanged from the transformer

• A CMS chain handling multi-timescale in-weights memory, updating at decreasing frequencies

• An SMLM layer whose inner update rule is learned rather than fixed, operating on top of the CMS chain’s output

Each component maps onto a different level of the nested learning hierarchy from section 7.1. Attention operates at infinite frequency, reading the current context with full precision and no persistence. The CMS chain operates at intermediate frequencies, encoding knowledge across windows of increasing length. The SMLM layer operates with a learned update rule that itself evolves during outer loop training, making the entire adaptation mechanism a product of optimization rather than engineering.

The result is a model where different parts of the parameter space change at different rates, for different reasons, according to rules that were themselves learned. Attention adapts every token, the CMS chain adapts every chunk, and the SMLM layer adapts according to a policy that the outer loop trained specifically to be useful for this task distribution. Each level handles the part of the plasticity-stability tradeoff it is best suited for, and no single level is asked to handle all of it alone.

7.4 Training Pipeline

Training a model that modifies itself at inference time is not straightforwardly reducible to standard supervised learning. The usual setup is based on fixed parameters, one forward pass, one backward pass, one gradient update and assumes the model being evaluated is the same model (static weight structure) throughout. Hope violates this assumption. The model at step t+1 has different inner parameters than the model at step t, because the forward pass at step t already updated them. Training a system like this requires a careful decomposition of what is being optimized, at what speed, and against what objective.

7.4.1 The Two-Loop Training Structure

Training Hope decomposes into two nested optimization loops, operating at different timescales and against different objectives. This maps directly onto the NL formalism from section 7.1, but now applied to the training procedure itself rather than just the inference procedure.

The inner loop runs at inference time. Given a sequence of inputs x1,x2,…,xT, the inner loop applies the learned update rule gΘ​ to update the inner parameters ϕ at each step:

The inner loop objective L_inner​ is a self-supervised signal, typically next-token prediction on the current chunk, the same surprise-based signal from Titans. This objective is local and computed on the fly. Neither labels are needed nor external supervision is required as the inner loop is entirely self-contained.

The outer loop runs during training. It optimizes the outer parameters Θ (which include the weights of the learned update rule gΘ​ and the slow CMS blocks) against a task-level objective L_outer​ computed after the inner loop has already run:

where ϕT​(Θk​) denotes the inner parameters after T steps of inner loop adaptation starting from Θk​. The outer gradient therefore has to propagate back through the inner loop updates, which is what creates the two-pathway gradient structure we described in section 7.3.4.

The key asymmetry is that the outer loop is training Θ to produce a gΘ​ that makes the inner loop useful, not just to minimize the loss on the current batch. The outer loop is optimizing for a property of the inner loop’s behavior, not just for a direct prediction objective.

7.4.2 The Outer Loop Objective

The outer loop objective has to capture what it means for the inner loop to have been useful. A naive choice would be to evaluate the task loss at the very end of the sequence, after all inner loop updates have been applied:

This is valid but myopic as it rewards gΘ​ only for what the inner loop managed to produce by the final token, giving no signal about the quality of intermediate adaptations. A better objective aggregates the task loss across all timesteps after the first inner update:

This rewards the learned update rule for producing good predictions at every step, not just at the end. It means the outer loop is training gΘ​ to be a good optimizer throughout the sequence, producing useful parameter updates early and consistently, rather than one that only converges to something useful after many steps.

The outer loop loss is evaluated against the standard language modeling objective on a held-out task distribution, something rich and diverse enough that the learned update rule gΘ​ has to generalize across contexts rather than overfit to a narrow task.

7.4.3 Keeping the Inner Loop Stable

The most practically delicate part of training Hope is ensuring the inner loop updates don’t destabilize during training. Because gΘ​ is itself being learned, there is no guarantee early in training that the update rule it produces is well-behaved. A poorly initialized gΘ​ can produce inner updates that send ϕ into a bad region of parameter space, which then produces a garbage outer gradient, which then produces a worse gΘ​, which produces even worse inner updates. The training dynamics can spiral.

Three mechanisms stabilize this.

Initialization. The inner parameters ϕ are initialized from the outer parameters Θ at the start of each sequence. This means the inner loop starts from a known good point, the outer loop’s current best estimate of a good parameter configuration, rather than from a random initialization. Even if the first few inner updates are poor, the starting point is already reasonable.

Update magnitude clipping. The learned update rule gΘ​ is regularized to produce updates bounded in norm:

for a small threshold δ. This prevents catastrophically large inner updates, regardless of what the raw output of gΘ​ would otherwise be. It is a hard constraint, not a soft penalty, so the inner loop cannot escape its stability envelope no matter how badly gΘ​ misfires early in training.

Truncated inner unrolling. As discussed in section 7.3.4, the outer loop gradient propagates back through the inner loop updates. If this unrolling runs for all T inner steps, the computation graph becomes enormous and the gradients become numerically unstable, exploding or vanishing through the long chain of inner updates. The paper truncates the unrolling to a fixed window K≪T, computing outer gradients only through the last K inner steps and treating earlier ones as detached constants. This introduces bias into the outer gradient estimate but keeps the training numerically stable and computationally tractable.

7.4.4 Data Ordering and Curriculum

In a standard transformer, training data order matters primarily for optimization efficiency. Curriculum learning, where easier examples come earlier, can speed up convergence, but the model’s final capabilities are largely order-agnostic because every forward pass sees the same fixed weights.

Hope breaks this. Because the inner loop updates ϕ as it processes a sequence, the order of tokens within a training sequence now affects the state of ϕ at every subsequent position. A sequence where easier, more predictable tokens come first produces a better-initialized ϕ by the time harder tokens arrive. A sequence where difficult tokens come first may send ϕ off in a bad direction before the inner loop has had time to stabilize. This has two practical consequences:

Within-sequence ordering matters. Training sequences should be constructed so that each sequence has some internal coherence, so that the inner loop is adapting to a consistent local context rather than a random splice of unrelated documents. Random shuffling of tokens within sequences, standard in many language model training pipelines for deduplication purposes, is harmful here because it destroys the local context structure that the inner loop is trying to adapt to.

Cross-sequence curriculum is beneficial. At the sequence level, there is evidence that presenting simpler, more internally coherent sequences early in training helps gΘ​ learn a stable update rule before being asked to handle difficult, rapidly-shifting contexts. Once gΘ​ has learned the basic mechanics of a useful update rule on easier material, it generalizes more reliably to harder sequences where the inner loop has to navigate more complex adaptation trajectories.

The interaction between data ordering and inner loop behavior is one of the less well-understood aspects of this training paradigm, and it represents one of the open engineering questions that differentiates training Hope from training a standard transformer at scale.

7.5 Inference-Time Optimization

The self-modifying mechanism described in 7.3 is theoretically clean, but it carries a practical cost that has to be confronted directly: running a gradient update inside every forward pass is expensive, and if that cost is not managed carefully, the latency overhead makes Hope undeployable regardless of how good its memory system is.

This section is about how the paper brings that cost down to a level that is practical at inference time, what approximations are made to get there, and what the residual overhead looks like relative to a standard transformer.

7.5.1 The Raw Cost of an Inner Update

Before discussing approximations, it is worth being precise about what the inner loop update actually costs in its full, unapproximated form.

At each token xt​, the full inner update requires three things: a forward pass through the inner parameters ϕt​ to compute f_ϕt​​(xt​), a backward pass through ϕt​ to compute the gradient ∇ϕ_Linner(ϕt;xt), and a forward pass through the learned update rule gΘ​ to compute the parameter update Δϕt​.

For a standard transformer forward pass with L layers and hidden dimension dd d, the cost is O(Ld²) per token. The full inner update, because it requires an additional backward pass through ϕ and a forward pass through gΘ​, adds a constant multiple of this cost per token, so the total per-token cost is O(cLd²) for some c>1 that depends on the size of gΘ​.

For large d and deep gΘ​, this multiplier becomes prohibitive. A model that is 3 times more expensive per token than a standard transformer is not a practical inference system, regardless of its quality. The inference-time optimizations are therefore not optional refinements but necessary conditions for Hope to be deployable at all.

7.5.2 Approximating the Gradient

The most expensive component of the inner update is the backward pass to compute ∇ϕ​L_inner​. Full backpropagation through ϕ costs roughly twice the forward pass through ϕ in FLOPs and requires storing all intermediate activations, which adds memory pressure proportional to the depth of ϕ.

The paper replaces the full gradient with a lightweight approximation computed directly from the forward pass activations, without a separate backward pass. The key observation is that for the local self-supervised objective used in the inner loop, the gradient signal is dominated by the prediction error at the output layer, and the deeper gradient terms are relatively small in magnitude for well-initialized ϕ.

Formally, the full gradient is:

​where the sum runs over all layers ℓ of the inner parameter network ϕ. The approximation retains only the last-layer gradient term and discards the deeper terms:

​This is computable directly from the output activation and the prediction error, without any backward pass through the earlier layers of ϕ. The approximation loses the signal from how the prediction error propagates back through the intermediate representations, and this is the primary quality cost of the inference-time optimization. For inner parameters that have already adapted to a reasonable local context, this approximation is tight, because the early-layer gradients are small relative to the output-layer gradient. For inner parameters that are far from a good local solution, early in a new context window for example, the approximation is looser, and the inner loop converges more slowly than it would under the full gradient.

7.5.3 Chunked Updates and Amortization

Even with the gradient approximation, computing an inner update at every single token is wasteful because adjacent tokens within a chunk are highly correlated. The update that ϕ would make on token t is almost identical to the update it would make on token t+1 if the two tokens carry similar information, so computing both updates independently squanders compute on a negligible difference.

Hope addresses this through chunked updates, which directly mirrors the CMS chunk structure from section 7.2. Rather than updating ϕ on every token, the inner loop accumulates a gradient estimate across a chunk of C tokens and applies a single update at the chunk boundary:

The gradient is averaged over the chunk, which has two benefits beyond compute savings. The averaged gradient is a lower-variance estimate of the true local gradient than any single token’s gradient, so the update is more stable. Additionally, averaging across a chunk means the update reflects the aggregate local context rather than any single token’s idiosyncratic signal, which makes ϕ more robust to individual noisy or uninformative tokens.

The amortized cost per token of the inner update is now 1/C​ of the per-token cost, since one update is shared across C tokens. For a chunk size of C=64, the overhead of the inner loop drops to less than 2% of the total forward pass cost per token, which is negligible in practice.

7.5.4 KV Cache Compatibility

A practical concern that is easy to overlook is how inner loop updates interact with the KV cache, the mechanism that makes autoregressive generation efficient by storing past keys and values to avoid recomputing attention over the full context at every step.

In a standard transformer, the KV cache is valid indefinitely because the attention weights are frozen. In Hope, the inner parameters ϕ change at each chunk boundary, so any KV cache entries computed under the old ϕ are no longer exactly valid under the new ϕ.

The paper handles this with a partial cache invalidation strategy. At each chunk boundary, only the KV entries corresponding to the most recent chunk are recomputed under the updated ϕ. Older KV entries, from chunks far back in the context, are retained without recomputation. The justification is that the inner update at chunk ii i represents a small perturbation to ϕ, so the keys and values computed under ϕ_i−1​ are close to what ϕ_i​ would compute for the same tokens, and the approximation error is bounded by the magnitude of the inner update, which is already constrained by the clipping from section 7.4.3.

This selective cache invalidation means the KV cache overhead of Hope is slightly larger than that of a standard transformer (because recent entries are periodically recomputed) but not dramatically so, since the cost is amortized across the chunk size C and the recomputed entries are only the most recent C tokens.

7.5.5 Net Overhead Relative to a Standard Transformer

With all three optimizations in place (gradient approximation, chunked updates, and selective cache invalidation), the per-token inference cost of Hope relative to a standard transformer at equivalent parameter count can be summarized as follows:

The 10-15% overhead figure assumes a moderately sized gΘ​ and a chunk size of 64. Larger chunk sizes reduce the overhead further at the cost of slower inner loop adaptation. Smaller chunk sizes increase the overhead but make ϕ more responsive to rapid context shifts.

The memory overhead is more significant than the compute overhead. Storing ϕ separately from Θ doubles the parameter count of the self-modifying module, and keeping the chunk gradient accumulation buffer in VRAM adds memory proportional to the chunk size and the inner parameter dimension. In practice, the total memory footprint of Hope is roughly 20-25% larger than a standard transformer at the same layer count and hidden dimension, with most of the excess coming from the dual parameter representation rather than the gradient buffers.

7.6 Muon and the Multi-Scale Momentum Optimizer (M3)

The methodology sections so far have focused on what the model learns and how it stores that learning like the CMS, the SMLM, the nested update structure; but, there is a fourth piece that sits underneath all of this and is easy to overlook: what optimizer is actually running the outer loop? The paper argues that the standard answer : Adam is not just suboptimal but structurally mismatched with the Nested Learning paradigm. This is where Muon enters.

7.6.1 What is Muon?

Muon stands for Momentum Orthogonalized by Newton-Schulz. It is an optimizer designed specifically for the 2D weight matrices of neural network hidden layers, and its core idea is deceptively simple: take the momentum gradient update that SGD-momentum would normally apply, and replace it with the nearest semi-orthogonal matrix to that update before applying it to the weights.

Formally, given a weight matrix gradient G∈R_n×m with singular value decomposition G=USV^⊤, the Muon update replaces G with:

The singular values in S which encode the magnitude of the gradient in each direction are discarded entirely. What remains is only the directional information: which directions in weight space the gradient is pointing, stripped of any scale differences between them.

This is computed not via full SVD, which is prohibitively slow, but via a Newton-Schulz matrix iteration that converges to UV^⊤ in just 5 steps and can be run stably in bfloat16:

where the coefficients (a,b,c)=(3.4445,−4.7750,2.0315) are tuned to maximize convergence speed. Each application of this quintic polynomial pushes the singular values of G′ closer to 1, and after 5 iterations the result is within a small tolerance of true orthogonalization. The FLOP overhead relative to standard training is at most Tm/B where T=5 is the number of Newton-Schulz steps, m is the model dimension, and B is the batch size in tokens in practice under 1% for typical LLM training scenarios.

7.6.2 Why Standard Optimizers Fall Short

To understand why Muon matters for Nested Learning specifically, you need to understand what the paper argues is wrong with Adam as a gradient compressor.

Recall from section 7.1 that NL frames gradient-based optimizers as associative memory modules, compressing a context flow (the gradient history) into a fixed representation (the optimizer state). Adam does this compression through two running statistics: the first moment mtm_t mt​ (a weighted average of past gradients) and the second moment v_t​ (a weighted average of past squared gradients):

The second moment term v_t​ acts as a per-coordinate scaling as it divides each gradient component by an estimate of its typical magnitude. This is beneficial when gradient components have very different scales, which is common in embedding layers and output heads. However, for the hidden 2D weight matrices of a transformer, the paper argues this coordinate-wise scaling is actively harmful. The gradient matrices for these layers typically have very high condition number as they are nearly low-rank, with most of the gradient signal dominated by a small number of directions. Adam’s coordinate-wise scaling does nothing to address this structural problem. It scales each entry independently, without any awareness of the matrix structure of the parameter it is updating.

The result is that the effective update applied by Adam to a 2D weight matrix is dominated by the same few high-singular-value directions in every step, with the rare but important directions; as the ones corresponding to small singular values consistently underweighted (mostly). The model learns the dominant directions fast and the rare directions slowly, regardless of how important those rare directions might be.

Orthogonalization fixes this by definition. By replacing G with UV^⊤, Muon equalizes all singular values to 1; hence every direction in weight space receives an update of equal magnitude, regardless of how large or small the corresponding gradient component was. Rare directions are no longer suppressed by the dominance of the common ones.

7.6.3 How Muon Connects to the NL Paradigm

The connection to Nested Learning is direct. In the NL framing, every optimizer is a memory module compressing gradient context into a parameter update. The question is which compression is optimal for which context flow.

The paper makes a specific argument: for the 2D hidden weight matrices, the relevant context flow is a sequence of gradient matrices with high condition number and near-low-rank structure. The optimal compression of this context flow, under a spectral norm objective rather than the element-wise L2​ objective that motivates Adam, is precisely the orthogonalized gradient. Adam is the optimal associative memory for element-wise L2​ regression on gradient vectors. Muon is the better-suited memory for the matrix-structured gradient context of hidden layers.

This is not just a philosophical point. It means that different components of the network should be optimized by different optimizers, chosen based on the structure of their gradient context flows like the corresponding embedding layers and output heads, whose gradients don’t have the same near-low-rank matrix structure, continue to use Adam. Only the hidden 2D weight matrices switch to Muon. The optimizer is no longer a single global choice but a per-component decision made on the basis of what each component’s gradient context looks like.

7.6.4 M3 : Multi-Scale Momentum Muon

The paper’s own contribution on top of Muon is the Multi-Scale Momentum Muon optimizer, abbreviated as M3. The motivation follows directly from the CMS design in section 7.2. If CMS showed that memory systems benefit from operating at multiple timescales simultaneously, the same logic applies to the gradient compressor which is the optimizer.

Standard Muon maintains a single momentum term at a single timescale:

M3 replaces this with a chain of momentum terms at different decay rates, mirroring the CMS chain structure:

where​ each momentum term (β1​>β2​>⋯>βK) has a different decay rate, compressing gradient history over a different timescale. The fast momentum terms (βk​ close to 0) respond quickly to recent gradients. The slow momentum terms (βk​ close to 1) compress a long history of gradient information, encoding stable long-run directional structure. The final update is a weighted combination of all timescales, orthogonalized before being applied.

The intuition is the same as CMS: no single timescale captures everything. A single β=0.9 momentum term compresses roughly the last 10 gradient steps. It loses information from step 100 and ignores information from step 2. M3 runs multiple compressors simultaneously and combines them, letting the optimizer retain signal from multiple depths of gradient history without having to choose a single window.

7.6.5 Where M3 Sits in the Full Hope Architecture

In the complete Hope architecture, the three optimization components map onto three distinct levels of the NL hierarchy:

• The inner loop (the SMLM update rule gΘ​) handles inference-time adaptation, running at the fastest timescale and updating on every chunk.

• The CMS chain handles multi-timescale memory storage, updating at intermediate frequencies across the sequence.

• M3 handles outer loop training, running at the slowest timescale and compressing gradient history across the full training run into stable weight updates

Each level compresses its own context flow at its own frequency. M3 is the outermost and slowest level (as we saw abaove) the one responsible for the kind of long-run structural learning that persists across training. Its multi-momentum design ensures that this compression is as expressive as possible, retaining signal from multiple depths of gradient history rather than collapsing everything into a single exponential moving average.

8. Results and Outcomes

The paper evaluates Hope across four task categories: language modeling, common-sense reasoning, long-context reasoning, and formal language recognition.

Baselines include Transformer++, RetNet, DeltaNet, Samba, Mamba2, TTT, and the original Titans, evaluated at 340M, 760M, and 1.3B parameter scales, all trained on 30B to 100B tokens from The Pile.

8.1 Language Modeling and Common-Sense Reasoning

At the 1.3B parameter scale, Hope achieves the highest average accuracy of 57.23% on reasoning tasks and the lowest perplexity of 15.11 on the Wiki benchmark, outperforming Transformer++, RetNet, DeltaNet, and the original Titans.

Across language modeling and common-sense reasoning tasks, Hope demonstrates lower perplexity and higher accuracy compared to modern recurrent models and standard transformers. The margin over Titans is meaningful because Titans already represented the state of the art among memory-augmented sequence models, so Hope’s improvement is not against a weak baseline.

The gain comes specifically from two places: the self-modifying update rule adapts the inner parameters to the local context more expressively than Titans’ fixed surprise signal, and the CMS chain provides multi-timescale memory that recovers information at longer ranges than Titans’ single long-term memory module.

8.2 Long-Context Reasoning

Hope showcases superior memory management in long-context Needle-In-Haystack (NIAH) downstream tasks, proving that the CMS offers a more efficient and effective way to handle extended sequences of information, outperforming Titans, TTT, and Mamba2.

NIAH is a particularly diagnostic benchmark for a continuously-learning memory system because it requires retrieving a specific piece of information planted early in a very long context, after processing a large amount of intervening noise. Standard transformers fail here because the relevant token falls outside the attention window.

RAG-based systems can fail when the retriever doesn’t surface the right document. Hope’s CMS chain handles it by encoding the planted fact into a slow-frequency block whose chunk window is large enough to span the distance between the fact and the query.

8.3 Formal Language Recognition

Hope achieves perfect accuracy on formal language recognition tasks where standard transformers fail, leveraging non-parallelizable computation enabled by its nested recurrent structure.

Formal language recognition is a stress test for computational depth rather than memory capacity. Certain formal languages, particularly those requiring counting or stack-based operations, cannot be recognized by a standard transformer because the attention mechanism’s parallelism prevents the kind of sequential state accumulation these tasks need. Hope’s nested recurrent structure, where the inner loop updates ϕ sequentially across each chunk, provides exactly the kind of iterative computation that transformers structurally cannot perform.

8.4 Muon and M3 : Optimizer Benchmarks

The paper evaluates M3 against AdamW and Muon on two axes: quality and training cost.

M3 consistently occupies the bottom-left quadrant; i.e. lower train loss and lower test loss simultaneously across multiple hyperparameter configurations. AdamW clusters further right and higher, indicating both worse generalization and worse training fit. Standard Muon sits between the two, better than AdamW but consistently outperformed by M3. The shaded pink region marks the Pareto-optimal zone; M3 configurations dominate it almost entirely across both the smaller and larger model scales shown.

When evaluating at scale; for 140M parameters M3 takes roughly 43 hours versus 38 hours for Muon and 41 hours for AdaMuon. At 1.3B parameters, M3 takes roughly 79 hours versus 70 hours for Muon and 75 hours for AdaMuon. The overhead is real but modest as it is approximately 10 to 15% above standard Muon and consistent with the theoretical estimate from section 7.5 that the multi-momentum bookkeeping adds bounded cost proportional to the number of momentum terms K.

The quality gain from Figure 11 outweighs the cost overhead from Figure 12 by a meaningful margin, which is the paper’s central argument for M3 as the right outer loop optimizer for Hope.

9. Thoughts

The Nested Learning paradigm is genuinely ambitious, and Hope validates it as a proof of concept. But several questions sit underneath the benchmark numbers that deserve honest examination.

• No comparison against a strong RAG baseline : Hope was not evaluated head-to-head against a strong RAG baseline on knowledge incorporation tasks. NIAH is a synthetic retrieval benchmark where the needle is planted deliberately. A real-world retrieval task with a well-tuned RAG system could tell a very different story, and the absence of that comparison leaves the paper’s central motivation empirically unresolved.

• The data ordering tax is a hidden deployment cost : Standard LLM training pipelines shuffle aggressively, deduplicate across document boundaries, and mix domains without regard for within-sequence coherence. Hope’s inner loop depends on that coherence to produce useful chunk-level updates. This is not a minor engineering note. It means adopting Hope at scale requires rethinking the data pipeline from the ground up, an invisible cost the benchmark numbers do not reflect.

• The chunk size selection problem has no principled answer right now : The CMS design shows that more frequency levels help, but the paper does not give a method for choosing the chunk size schedule. At 1.3B parameters with a small fixed set of chunk sizes, this is manageable. At 10B or 70B, with longer contexts and more diverse data distributions, the wrong frequency schedule could waste most of the multi-timescale benefit. This feels like the kind of hyperparameter that will require its own line of research before CMS is reliably deployable.

• M3 is the most immediately usable contribution : Unlike Hope, which requires rearchitecting the full model and training pipeline, M3 can be dropped into any existing training setup as a replacement for Adam on hidden layer parameters. The overhead is bounded and predictable, the quality gains are consistent across scales, and it requires no changes to inference. The community will likely adopt M3 independently of whether Hope itself gains traction, and the paper somewhat undersells this.

• The inner loop overhead scaling is uncharacterized beyond 1.3B. The 10 to 15% per-token overhead estimate from section 7.5 holds under specific assumptions about chunk size and model dimension. At larger scales, the interaction between growing model dimension, longer contexts, and the chunked update approximations introduces compounding factors that have not been studied. It is plausible that the overhead remains bounded, but it is equally plausible that it grows in ways the current analysis does not capture.

• The most profound implication is what NL says about the design space. If optimizers, architectures, and learning algorithms are all instances of the same underlying operation, then the field has been exploring a very small corner of what is possible. We have been treating the model and the thing that trains the model as categorically separate objects. Nested Learning dissolves that boundary. Hope is a first-order approximation of what emerges when you take that dissolution seriously. The more interesting architectures are the ones that come after it.

10. Conclusion

Nested Learning reframes the entire machinery of deep learning, not as a stack of layers trained by an external optimizer, but as a coherent system of nested optimization problems, each compressing its own context flow at its own frequency. Through this lens, the boundary between architecture and optimizer dissolves, the boundary between training and inference softens, and the design space for future models expands in a dimension the field has barely begun to explore.

In this article, we started by establishing why current models are structurally amnesiac as their knowledge is confined to either the immediate context window or the frozen weights of pretraining, and every mechanism built to work around this, from RAG to fine-tuning to in-context learning, is a patch on top of a system that was never designed to keep learning. We grounded the difficulty of the problem in the plasticity-stability tradeoff and the catastrophic forgetting problem, showing why continuous learning is not just an engineering challenge but a structural one.

We then built up the Nested Learning paradigm from first principles, revisiting the Titans architecture that first showed inference-time learning was possible, and tracing how Hope extends that foundation into a full multi-timescale system. We explored the Continuum Memory System, which replaces the binary long/short-term memory split with a spectrum of frequency levels each responsible for its own timescale of knowledge. We examined the Self-Modifying Titans module, whose inner loop learns not just what to update but how to update, making the adaptation mechanism itself a product of optimization. We introduced M3, which brings the same multi-timescale philosophy to the optimizer, replacing Adam’s single exponential average with a chain of momentum terms that compress gradient history at multiple depths simultaneously.

The key takeaway is that the static nature of current language models is not a fundamental constraint of the architecture class. It is a consequence of treating the model and its training procedure as two separate objects, one that learns and one that is learned. When you recognize them as the same kind of thing operating at different frequencies, a new design axis opens up. Hope is the first architecture to seriously inhabit that axis.

We are not yet at a model that learns the way a person does, continuously, selectively, without forgetting, weaving new knowledge into the substrate of old reasoning rather than retrieving it from an external store. But Nested Learning tells us what the path toward that model looks like. The question is no longer whether inference-time learning is possible. The question is how deep the nesting can go; this changes the entire framework and space on which we optimnize.

Paper : https://arxiv.org/pdf/2512.24695

That’s all for today.

Follow me on LinkedIn and Substack for more such posts and recommendations, till then happy Learning. Bye👋

This chapter covers

• Why video classification needs both spatial and temporal reasoning

• How a video clip becomes a sequence of patch tokens

• The self-attention equations used inside TimeSformer

• Five space-time attention patterns: space-only, joint, divided, sparse local-global, and axial

• Why divided space-time attention gives the best accuracy-efficiency tradeoff

• How TimeSformer compares with 3D convolutional video models

• How attention visualizations show what the model learns from video

Architectures like Vision Transformers and Convolutional Neural Networks treat an individual image as the core unit of visual understanding. Though analyzing them can be difficult, a single image remains a fixed grid of pixels. Video, however, adds the dimension of time; it is a sequence of images where meaning and labels heavily depend on observing how the scene evolves across frames.

TimeSformer, introduced by Bertasius, Wang, and Torresani in the ICML 2021 paper Is Space-Time Attention All You Need for Video Understanding?, asks a direct question: can we build a strong video recognition model using self-attention alone, without 3D convolutions? The answer from the paper is yes, provided that attention is arranged carefully over the space-time volume of a video clip.

The model starts from the Vision Transformer idea: split visual data into patches, embed those patches as tokens, add positional information, and feed the tokens to a transformer encoder. The difference is that a video has patches across both space and time. If every patch in every frame attends to every other patch, the computation becomes expensive very quickly. TimeSformer is mostly about how to make that attention usable.

Note

In this chapter, the word video means the visual stream only. TimeSformer, as presented in the paper, does not use the audio track. The task studied here is video classification: given a clip, predict one action label for the whole clip.

1.1 From images to videos

A still image classifier can sometimes solve a video task by looking at one frame. If the clip shows a cricket batter at the crease or a basketball player on the court, one frame may already contain enough context. But many action labels depend on motion. A frame from a push-up and a frame from a high plank can look nearly identical, while the sequence of frames makes the difference obvious.

Figure 1.1. Static image classification can fail when two actions share nearly identical individual frames. In this example, a push-up and a high plank may look similar in a single snapshot. A video model must use the order of frames to detect whether the body is moving through an action or holding a static pose.

Figure 1.1 captures the central problem. A video model has to answer two questions at once: what is visible? and how is it changing? The first question is spatial. It concerns objects, body pose, background, and scene layout. The second is temporal. It concerns motion, ordering, repetition, and cause-effect patterns across frames.

Figure 1.2. A video clip can be viewed as a sequence of images ordered along time. Each frame can be divided into patches, just as in a Vision Transformer, but the model now receives many grids of patches instead of one. The token sequence therefore grows with both image resolution and the number of frames.

The Kinetics datasets used in the TimeSformer paper contain short videos of human actions. They include labels such as headbanging, stretching a leg, shaking hands, tickling, robot dancing, and salsa dancing. Some labels are strongly tied to appearance; others require motion to disambiguate the action.

Figure 1.3. Examples from a Kinetics-style action recognition dataset. The model receives a short clip and predicts one action category for the whole clip. The examples show why both spatial context and temporal evidence matter: the person, object, and scene provide cues, but the action is often defined by movement across frames.

Figure 1.4. Some video labels can be predicted mostly from spatial context. A single frame of cricket or basketball already contains strong cues from the player, equipment, pose, and scene. TimeSformer still processes the clip as video, but examples like these explain why space-only attention can perform reasonably well on datasets with strong scene bias.

The opposite case is the reason TimeSformer is interesting. When appearance alone is ambiguous, the model needs to compare a region in one frame with related regions in later frames. This is where temporal attention enters.

Figure 1.5. Spatial attention and temporal attention answer different questions. Spatial attention compares patches within a frame, such as the relationship between a hand, body, and background. Temporal attention compares corresponding or related patches across frames, allowing the model to represent motion and change over time.

1.2 Turning a video into tokens

TimeSformer follows the same tokenization idea as ViT, but applies it to every frame in a video clip. Let the input clip be

where F is the number of frames, H and W are the frame height and width, and the last dimension stores RGB channels. The TimeSformer paper writes the dimensions as H X W X 3 X F; both notations describe the same video volume.

Each frame is divided into non-overlapping patches of size P x P. The number of spatial patches per frame is

Across F frames, the clip contains NF patch tokens. A learned classification token is prepended, so the total token count is

Figure 1.6. Counting the number of tokens in a video transformer. A frame is divided into spatial patches, and this patch grid is repeated for every sampled frame. The total sequence length is therefore the number of patches per frame multiplied by the number of frames, plus one class token for video-level classification.

The default TimeSformer setting in the paper uses clips of 8 x 224 x 224 and a patch size of 16 x 16 pixels. That gives

patches per frame and

tokens per clip. This is already much longer than the 197 tokens in a standard ViT image model with one 224 x 224 image and 16 x 16 patches.

For a patch at spatial index p and frame index t, the flattened RGB patch is

It is mapped to a D-dimensional embedding with a learned matrix E:

Here, e_p^space tells the model where the patch lies inside the frame, and e_t^time tells it which frame the patch came from. This separation is useful because a token has two coordinates: its spatial location and its temporal position.

1.3 Self-attention over a video

Once the clip is represented as tokens, TimeSformer uses the usual transformer ingredients: layer normalization, query-key-value projections, multi-head attention, residual connections, and an MLP. For head a in block l , query, key, and value vectors are computed as

The attention score between a query and a set of keys uses the same scaled dot-product rule as the original Transformer:

where d_h = D / A is the feature dimension of one attention head. The output for a query token is the weighted sum of value vectors:

The set

is the important part. It defines which tokens a query is allowed to attend to. TimeSformer compares several choices for this neighborhood.
After attention, the heads are concatenated, projected, and passed through a residual MLP:

The final video representation is taken from the output class token:

A small classification head maps y to the video classes.

1.4 Five attention patterns

The TimeSformer paper studies five ways of arranging attention over space and time. The simplest way to understand them is to imagine choosing one query patch and asking: which other patches are allowed to become keys for this query?

1.4.1 Space-only attention

Space-only attention treats each frame almost like an independent image. A query patch attends to patches in the same frame, plus the class token:

This needs only N+1 key comparisons per query patch. It is cheap, and it can work when static appearance gives most of the answer. But it cannot directly model how a patch changes across frames.

Figure 1.7. Space-only attention. The query patch attends to patches within the same frame, but it does not directly compare itself with patches in earlier or later frames. This pattern is efficient and can work for appearance-biased labels, but it discards explicit temporal reasoning.

1.4.2 Joint space-time attention

Joint space-time attention is the most direct extension of ViT to video. A query patch attends to every patch in every frame:

This gives the model full access to the video volume, but the cost grows quickly.

Each query compares against NF+1 keys, and the full attention matrix has size (NF+1) x (NF+1).

Figure 1.8. Joint space-time attention. The query patch can attend to all patches in all frames, giving the model a full global view of the clip. The drawback is computational cost: as resolution or video length increases, the number of pairwise comparisons becomes large enough to exhaust GPU memory.

For the default setting N=196 and F=8, joint attention gives each query 1569 possible keys. For a long clip with F=96, the same spatial resolution would give 18,817 keys per query. This is why a naive global video transformer is difficult to scale.

1.4.3 Divided space-time attention

Divided space-time attention is the main TimeSformer design. It factorizes the work into two smaller attentions inside each transformer block:

1. Temporal attention compares the same spatial patch location across frames.

2. Spatial attention compares patches within a frame.

For temporal attention, the neighborhood is

The output of temporal attention is fed into spatial attention, which then uses

In residual form, a divided TimeSformer block can be summarized as

Figure 1.9. Divided space-time attention. The query first attends along the temporal axis at the same spatial patch location, then attends spatially within a frame. This gives the model access to both motion and appearance while avoiding the full quadratic cost of joint space-time attention.

Compared with NF+1 comparisons per query for joint attention, divided attention uses only

comparisons per query patch. With N=196 and F=8, this is 206 comparisons instead of 1569. With F=96, it is 294 instead of 18,817. The saving becomes larger as clips get longer.

1.4.4 Sparse local-global attention

Sparse local-global attention approximates full space-time attention by combining a local neighborhood with a sparse set of global locations. The local part captures nearby motion and appearance. The global part samples farther patches without attending to every possible location.

Figure 1.10. Sparse local-global attention. The query attends densely to a local space-time neighborhood and sparsely to farther locations. This gives the model some long-range context without paying the cost of full joint space-time attention. In the TimeSformer experiments, however, this pattern did not match divided attention’s accuracy.

1.4.5 Axial attention

Axial attention decomposes the video volume into one-dimensional passes. Instead of attending over the full space-time grid at once, it applies attention over time, then width, then height. This is another way to reduce cost by factorizing a high-dimensional attention problem.

Figure 1.11. Axial attention. The model decomposes attention into separate passes along time, width, and height. This reduces the size of each attention operation, but it also imposes a stricter path for information flow through the video volume.

Figure 1.12 places the five attention blocks side by side. The diagrams also show where residual connections and MLP layers sit inside the block.

Figure 1.12. The five TimeSformer attention block designs compared in one view. Space-only uses one spatial attention layer, joint space-time uses one global attention layer, divided space-time uses temporal attention followed by spatial attention, sparse local-global combines local and sparse global passes, and axial attention applies attention along time, width, and height.

1.5 Why divided attention works well

At first, joint space-time attention sounds like the strongest option because it lets every token directly compare itself with every other token. The experiments tell a more practical story. Joint attention is expressive, but it is expensive, and it does not give the model separate parameters for temporal and spatial interactions. Divided attention uses separate temporal and spatial query-key-value projections, so it has more specialized capacity while still doing less pairwise work.

The cost difference becomes clear when the spatial crop or the number of frames increases.

Figure 1.13. Computational cost comparison from the TimeSformer paper. Joint space-time attention becomes expensive as spatial crop size or input length grows, and it runs out of memory in the larger settings shown. Divided space-time attention scales more gently because it factorizes attention into temporal and spatial stages.

The accuracy comparison also favors divided attention.

Figure 1.14. Video-level accuracy for the five attention schemes reported in the TimeSformer paper. Divided space-time attention obtains the best result on both Kinetics-400 and Something-Something-V2. Space-only attention performs well on Kinetics-400 but collapses on Something-Something-V2, which is a stronger test of temporal reasoning.

The table in Figure 1.14 is worth reading carefully. Space-only attention reaches 76.9% on Kinetics-400, which confirms that many Kinetics labels can be predicted from appearance. But it reaches only 36.6% on Something-Something-V2. That dataset contains labels where the same objects can appear in different actions, so temporal order matters much more. Divided attention reaches 78.0% on Kinetics-400 and 59.5% on Something-Something-V2, beating the other attention patterns in the paper’s comparison.

Note

The main lesson is not that temporal attention should replace spatial attention. The lesson is that video classification needs both, but not necessarily in one huge attention matrix. TimeSformer gets a strong result by making time and space separate operations inside the same block.

1.6 An implementation view

The implementation idea behind divided attention is simple once the tensor shape is clear. Suppose the token tensor, after removing or separately handling the class token, has shape

where B is the batch size. Temporal attention should see F as the sequence dimension while keeping a fixed spatial patch index. Spatial attention should see N as the sequence dimension while keeping a fixed frame.

Listing 1.1.

The tensor reshaping pattern behind divided space-time attention. Temporal attention groups batch and spatial patch index together, so attention runs across frames. Spatial attention groups batch and frame index together, so attention runs across patches inside one frame.

# x has shape [B, F, N, D]

# Temporal attention: sequence length is F.
x_t = rearrange(x, "b f n d -> (b n) f d")
x_t = temporal_attention(x_t, x_t, x_t)
x = rearrange(x_t, "(b n) f d -> b f n d", n=N)

# Spatial attention: sequence length is N.
x_s = rearrange(x, "b f n d -> (b f) n d")
x_s = spatial_attention(x_s, x_s, x_s)
x = rearrange(x_s, "(b f) n d -> b f n d", f=F)

Listing 1.1 is not a full TimeSformer implementation, but it captures the core trick. To apply temporal attention, we flatten the batch and patch dimensions into one larger batch and let the attention module process the frame dimension. To apply spatial attention, we flatten the batch and time dimensions and let the attention module process the patch dimension.

The full model also needs patch embedding, a class token, space and time positional embeddings, layer normalization, MLP blocks, and a classification head. These are inherited from the transformer and ViT design. TimeSformer’s main change is how attention is scheduled over a video.

1.7 What the model learns

Attention visualizations help make the model less abstract. When TimeSformer predicts a Something-Something-V2 action, the attended regions often track the object or the hand-object interaction that defines the action.

Figure 1.15. Space-time attention visualization from the TimeSformer paper. The highlighted regions show that the model focuses on task-relevant areas, such as the hand and manipulated object, rather than spreading attention uniformly over the frame. This is the behavior expected from a model that must reason about actions, not just objects.

The paper also visualizes learned video embeddings with t-SNE. Each point represents a video, and videos with the same action label share a color. The divided-attention model forms more separable clusters than ViT or a space-only TimeSformer.

Figure 1.16. Feature visualization from the TimeSformer paper. The divided space-time model forms more separated action clusters than a space-only model or a ViT baseline. This suggests that the temporal branch is not merely adding computation; it changes the learned representation in a way that separates actions more cleanly.

These visualizations do not prove causal reasoning by themselves, but they support the quantitative result: modeling time explicitly improves the representation for action recognition.

1.8 Comparison with 3D convolutional models

Before video transformers, strong video classifiers were usually built from 3D convolutions. A 3D convolution slides a local kernel over height, width, and time. This gives the model a built-in local motion bias, which is useful, but it also means long-range relationships must be built gradually through many layers.

TimeSformer takes a different route. Self-attention can connect distant patches directly, and divided attention keeps the cost manageable enough to process longer clips. The paper reports that the default TimeSformer has 121.4 million parameters but only 0.59 TFLOPs of inference cost for the reported setting, compared with 1.97 TFLOPs for a SlowFast R50 model in the same comparison.

The larger TimeSformer-L variant processes 96 input frames and reaches 80.7% top-1 accuracy on Kinetics-400 in the paper’s reported table.

The official TimeSformer repository provides pretrained model variants for Kinetics-400, Kinetics-600, Something-Something-V2, and HowTo100M, including longer-input models that cover much more temporal context than the default 8-frame setting. The Hugging Face Transformers documentation also exposes TimeSformer as a video classification model, which makes the architecture easier to reuse in modern transformer workflows.

The practical tradeoff is clear. TimeSformer benefits from image pretraining and large datasets, while 3D CNNs still offer strong inductive biases and can be easier to train in smaller-data settings. But the paper showed that a convolution-free video model was no longer just a curiosity. With the right attention pattern, transformers could compete seriously on video recognition.

1.9 Resources

- Video:

In this video, Dr. Sreedath Panat. explains how to implement TimeSformer from scratch and walks through the main building

blocks of the model.

TimeSformer from Scratch

- Paper:

Is Space-Time Attention All You Need for Video Understanding?

Summary

• A video clip can be represented as a sequence of frame-level patches. If each frame has N patches and the clip has F frames, the transformer receives NF+1 tokens, including the class token.

• Space-only attention compares patches within the same frame. It is efficient and can work on appearance-biased datasets, but it cannot explicitly model motion across frames.

• Joint space-time attention compares every patch with every patch in every frame. It is expressive, but its cost grows quickly with spatial resolution and clip length.

• Divided space-time attention factorizes the work into temporal attention followed by spatial attention. This reduces the comparisons per query from NF+1 to N+F+2 while still modeling both motion and appearance.

• In the TimeSformer paper’s comparison, divided space-time attention achieves the best accuracy among the five studied schemes on both Kinetics-400 and Something-Something-V2.

• TimeSformer differs from 3D CNNs by replacing local convolutional kernels with attention over patch tokens. This gives it a direct route to long-range video modeling, especially when enough data and pretraining are available.

• Attention and feature visualizations show that the divided model learns representations that focus on action-relevant regions and separate video classes more clearly than space-only variants.

More Blogs

Vizuara AI Labs

The Transformers

The Transformer Architecture…

Read more

8 months ago · 167 likes · 5 comments · Mayank Pratap Singh

Vizuara AI Labs

Vision Transformers

Table of Contents…

Read more

7 months ago · 71 likes · Mayank Pratap Singh

Vizuara AI Labs

Segment Anything Model (SAM)

Figure 0: Detailed Architecture of the Segment Anything Model (SAM…

Read more

9 months ago · 17 likes · Mayank Pratap Singh and Sreedath Panat

I’m also building Audio Deep Learning projects and LLM projects, sharing and discussing them on LinkedIn and Twitter. If you’re someone curious about these topics, I’d love to connect with you all!

Mayank Pratap Singh

LinkedIn : www.linkedin.com/in/mayankpratapsingh022

Twitter/X : x.com/Mayank_022.

This Blog Covers

• Why earlier generators such as GANs and VAEs left room for a different approach

• The diffusion idea: learn to reverse a noising process

• The small set of equations needed to read a DDPM diagram

• How U-Nets, timestep embeddings, and latent spaces make diffusion practical

• How conditioning and transformer backbones shape modern image generation

1.0 Introduction

Image generation used to be easy to describe and hard to train. A generator would take a random vector and try to turn it into a realistic image. If the image looked wrong, the training signal had to explain what was wrong. That sounds reasonable until you remember how many ways an image can fail. The subject can be distorted, the texture can be weak, the pose can collapse, the background can look fake, or the image can be sharp but repetitive. A single image is not one decision. It is a large collection of decisions that need to agree.

Diffusion models became important because they change the shape of the problem. Instead of asking a neural network to create a full image in one shot, they ask it to solve a much smaller task many times: given a noisy image, predict how to remove a little of the noise. The same model repeats that correction from heavy noise to a clean image. This makes generation feel less like a leap and more like a guided recovery process.

This chapter starts with the older ideas because they explain why diffusion was such a useful turn. GANs showed that learned generators could make sharp, impressive images, but adversarial training could be fragile. VAEs showed that a model could learn a useful latent space, but samples often looked too smooth. Diffusion combines a different training target with a gradual sampling process. It is not free, since sampling can require many network calls, but the training signal is simple, stable, and visually powerful.

We will keep the math focused. You do not need a full derivation to understand why diffusion works. You need to know how clean images are mixed with Gaussian noise, why a timestep matters, what the denoiser predicts, and how repeated denoising turns a random sample into an image.

1.1 Why generative models needed another path

Before diffusion models became the standard tool for image generation, two families shaped most of the conversation: generative adversarial networks and variational autoencoders. They are worth covering first because diffusion did not replace an empty space. It solved pain points that became clear through years of work with older generators.

GANs were exciting because they could produce crisp images. The generator learned to map a random vector to an image, and a discriminator learned to judge whether an image came from the training set or from the generator. The generator improved by trying to fool the discriminator. This setup created a learned realism signal, which was much richer than a hand-written pixel loss.

Figure 8.1. A basic GAN setup. The generator maps random noise to fake images. The discriminator receives both real training images and generated images, then predicts whether each input is real or fake. The generator improves by trying to fool the discriminator, so the training signal comes from a learned critic rather than from a fixed reconstruction loss.

As shown in Figure 8.1, a GAN is a two-network game. This game is the source of its strength and also the source of many practical problems. If the discriminator becomes too good too quickly, the generator may get weak gradients. If the generator finds a small set of outputs that fool the discriminator, it can repeat those outputs and lose diversity. This failure is often called mode collapse. The samples may be sharp, but the training process can feel like balancing two moving targets.

This contrast is useful throughout the chapter. In a GAN, the generator learns from the discriminator’s current opinion. In a diffusion model, the target for a training example is the exact noise that was added by the training code. That target is not another network’s opinion. It is a known value. This does not make diffusion magically easy, but it removes the adversarial game from the core training loop.

VAEs took a different route. A variational autoencoder learns to encode an image into a latent distribution and decode a latent sample back into an image. This gives the model a compressed internal space, which is a useful idea for generation and editing. The tradeoff is that VAE samples and reconstructions often look smoother than GAN samples, especially when the decoder is trained with losses that reward average-looking reconstructions.

Figure 8.2. A variational autoencoder. The encoder maps input data to a latent distribution, a latent sample is drawn, and the decoder reconstructs the output. VAEs are important for this chapter because modern latent diffusion systems borrow the idea of working in a compressed representation, then use diffusion to model that latent space more powerfully.

As shown in Figure 8.2, the part of the VAE story that survived strongly into modern systems is compression. A good latent representation can keep the meaningful structure of an image while using far fewer values than the original pixels. Diffusion models later used this idea to reduce cost. Instead of running the denoising model directly on every pixel, latent diffusion runs it on a smaller latent tensor and decodes the result at the end.

Older generators taught the field two lessons. First, sharp images matter, and a model needs a strong learned signal for realism. Second, the representation matters, because pixel space is expensive. Diffusion uses both lessons but changes the training problem. It turns image generation into the task of learning how to remove noise at many levels of corruption.

That change sounds modest, but it is a major shift in supervision. A GAN has to learn from a moving discriminator. A VAE has to balance reconstruction quality against a regularized latent space. A diffusion model can create its own training pairs. The input is a noisy version of a real image, and the target is the noise that the training code added. This gives the model many cleanly defined tasks from each training image.

It also changes how we should think about creativity in the model. The denoiser is not storing a library of clean images and retrieving one during sampling. It learns the local direction from “less likely under the data” toward “more likely under the data” at many noise levels. When sampling begins from random noise, those local directions are applied again and again. The final image is the result of the full path, not a single lookup.

This is the motivation for the rest of the chapter. Diffusion is useful because it breaks a hard one-shot generation problem into a sequence of denoising problems. The cost is that sampling has to run the model repeatedly. The benefit is that each training example has a simple, known target, and this makes the method easier to scale than adversarial training in many settings.

1.2 The Diffusion Idea

The central idea of diffusion is simple: take real data, slowly corrupt it into noise, and train a model to reverse that corruption. During training, the forward corruption process is not learned. We define it. During generation, the reverse process is learned. The model starts from noise and repeatedly takes small steps toward a cleaner sample.

Figure 8.3. The goal of diffusion modeling. The true image distribution P_data(x) is known only through examples, not through a neat formula. A diffusion model learns a sampler that starts from a simple base distribution and gradually moves toward regions that look like the training data.

As shown in Figure 8.3, the problem can be read from a distribution point of view. The training images are samples from a complicated image distribution. We do not know the full distribution, and we cannot directly sample new images from it. What we can sample from easily is Gaussian noise. Diffusion connects these two worlds by defining a path from data to noise, then learning how to travel back.

This path is easier to understand if you imagine a clean image being damaged step by step. At first, the subject is still visible. Later, only broad structure remains. At the end, the image is essentially noise. Because we control this damaging process, we can create unlimited supervised training examples from every image: pick a timestep, add the right amount of noise, and ask the model to predict the noise that was added.

Figure 8.4. A data manifold being diffused into a simple noise distribution. At t=0, the samples lie on a structured spiral. As the forward process adds noise, the structure spreads out and eventually becomes close to an isotropic Gaussian. The reverse process learns to move in the opposite direction, from the easy endpoint back toward the structured data.

The toy dataset in Figure 8.4 is two dimensional, but the picture is the same for images. The structured data starts on a thin, meaningful shape. Noise gradually washes that structure out until the final distribution is simple. The reverse model learns how to move a sample from the noisy cloud back toward the kind of structure seen in the training data.

This is why diffusion can feel backward at first. We spend training time destroying images. The point is that controlled destruction gives us a clean learning signal. If the training code added a particular noise pattern, the network can be trained to recover that pattern. Generation then uses the trained network in the opposite direction.

1.3 Why the endpoint is Gaussian noise

Every generator needs a starting point. If the starting point is already complicated, sampling is already hard. Diffusion models usually start from a standard Gaussian distribution,

This distribution is easy to sample from, easy to reason about, and flexible enough to act as a neutral blank canvas. It does not prefer faces, dogs, product photos, outdoor scenes, or any other category. The structure comes from the learned reverse process.

Figure 8.5. Why diffusion models use a Gaussian base. The structured data distribution is hard to sample from directly, but a maximum entropy Gaussian endpoint is easy to sample from. The forward process transforms structured data into that endpoint, and the reverse process learns how to undo the transformation.

The distribution-level reason is shown in Figure 8.5. A simple Gaussian endpoint lets us begin generation from something we can produce on demand. If we can learn a reliable reverse path from this endpoint to the data distribution, then sampling becomes possible even though the original image distribution is too complex to write down.

Figure 8.6. Base distributions as blank canvases. A uniform distribution and a Gaussian distribution are both high entropy choices, but Gaussian noise is especially convenient for diffusion because repeated additive Gaussian noise stays Gaussian and leads to a tractable endpoint for the forward process.

Possible blank canvases are compared in Figure 8.6. In practice, Gaussian noise is convenient because adding independent Gaussian noise repeatedly behaves cleanly. The endpoint remains mathematically manageable, and the training code can directly sample noisy versions of an image at any timestep. This is a practical choice, not a mystical one.

The important point is that the base distribution is simple on purpose. The model should not need a hand-designed image prior at the start. It should learn image structure from data and from the conditioning information.

1.4 The Forward Process

The forward process is the noising process we define. At each timestep, the clean image keeps some of its signal and receives some new Gaussian noise. The standard notation uses

for the amount of new noise at step

The matching signal retention is

The cumulative signal after many steps is

You can read

as “how much of the original image is still present after step

Early in the chain, this value is close to 1. Late in the chain, it is close to 0.

Figure 8.7. Forward diffusion on an image. The clean horse image is progressively corrupted with Gaussian noise. The distributional view below it shows the same change in coefficients: the signal coefficient decreases, the noise coefficient increases, and the noisy sample becomes less tied to the original image.

The sequence in Figure 8.7 is one of the most important ideas in the chapter. The model is not trained only on clean images and pure noise. It sees many intermediate states. Some examples are barely corrupted. Some are heavily corrupted. This gives the network a curriculum of denoising problems across the whole path from data to noise.

The math you actually need

The forward process can be sampled in one step:

This equation says that a noisy image is a weighted mixture of the clean image

and pureGaussian noise

The value

Controls the signal part, and

controls the noise part.

This equation is worth remembering because it explains how training examples are made. We do not need to simulate every previous step to get X_t We can sample a random timestep t, sample a random noise tensor

, and mix the clean image with that noise using the formula above. That single line is the workhorse behind DDPM-style training.

The square roots in the equation are there because the formula is balancing variance, not ordinary image brightness. You do not need to derive this to use the equation. A useful reading is that

scales the clean image and

scales the noise. When t is small, the clean coefficient is large and the noise coefficient is small. When t is large, the clean coefficient shrinks and the noise coefficient dominates.

This also explains why a noisy training sample is still tied to one clean image. At early timesteps, the noisy sample clearly resembles its source. At later timesteps, the source becomes harder to see, but the training code still knows exactly which clean image and which noise tensor created it. The denoiser receives the noisy sample and the timestep, then learns how much of what it sees should be treated as removable noise.

Figure 8.8. Alpha-beta scheduling in forward diffusion. The beta schedule controls the new noise injected at each step. The alpha values track per-step signal retention, while the cumulative product

tracks how much of the original image remains after many steps.

The curves in Figure 8.8 explain the schedule.

decides how aggressively noise is added at a particular step.

is the part of the signal kept at that step.

is the accumulated signal retention after many steps. A good schedule makes the denoising tasks neither too easy nor too destructive. If the process destroys structure too quickly, the network sees many examples that are almost pure noise.

If it destroys structure too slowly, the sampling process wastes steps making tiny changes.

Modern systems may express the schedule through sigmas, log signal-to-noise ratios, or continuous time variables, but the beginner intuition is the same. A schedule controls how the problem changes as the sample moves from clean data to noise and back.

1.5 What the denoiser learns

The denoiser is the learned part of the system. It receives a noisy sample

the timestep t and often a condition c.

The condition could be a class label, a text prompt, a mask, an edge map, a depth map, a low resolution image, or another signal. In the basic noise-prediction view, the network outputs

which is its estimate of the noise that was used to create

The training target is known because the training code sampled the noise. That gives a simple supervised loss:

This is not the full probabilistic derivation, but it is the equation that explains the practical training loop. Add known noise, ask the network to predict that noise, and update the network when it is wrong.

The loss is simple, but the task is rich. At low noise levels, predicting the noise can mean finding small corruptions around edges and textures. At high noise levels, predicting the noise is tied to understanding which broad structures are plausible under the data and condition. The same loss therefore teaches both local cleanup and global structure, depending on which timestep was sampled.

This is one reason diffusion training is often easier to reason about than GAN training. The network is given a direct regression target. If the prediction is wrong, the loss says so in a straightforward way. The model may still be large, the data may still be messy, and the sampler may still require careful tuning, but the basic training signal is stable.

Figure 8.9. A DDPM denoising step. The noisy image X_t and timestep t enter the denoiser.

The network predicts the noise component

, and that prediction is used to compute the next, slightly cleaner sample

Generation repeats this step many times.

The denoising step in Figure 8.9 shows why diffusion sampling is iterative. The denoiser is not asked to jump straight from noise to the final image. It is asked to make one correction for the current noise level. When those corrections are repeated across a schedule, structure emerges gradually.

Predicting noise is a good first mental model because the target is exact. Some modern models predict the clean sample, a velocity target, or a flow direction instead. These variants matter in advanced systems, but they keep the same basic teaching idea: the model learns a direction that moves a noisy sample toward the data distribution.

The timestep is essential. A sample at a low noise level should be treated differently from a sample at a high noise level. At low noise, the model should refine edges and small details. At high noise, it must infer broad structure from weak evidence and from the condition. Without timestep information, the same network would not know which denoising problem it is solving.

It is also useful to separate the denoiser from the sampler. The denoiser predicts a target from the current noisy state. The sampler decides how to use that prediction to move to the next state. Two systems can use the same trained denoiser with different samplers and produce different speed-quality tradeoffs. A slow sampler may take many careful steps. A fast sampler tries to take fewer, larger steps without drifting away from the learned path.

This is why diffusion research often talks about both models and samplers. The model learns the direction. The sampler chooses the route. For a beginner, DDPM sampling is the easiest route to understand because it follows the chain step by step. Faster samplers keep the same goal but take a more efficient path through the noise levels.

1.6 Giving the network a sense of time

A timestep such as t=427 is just a number. Neural networks usually work better when important scalar values are expanded into richer representations. Diffusion models commonly use sinusoidal timestep embeddings, which turn a scalar timestep into a vector containing sine and cosine features at several frequencies.