The J-space of a diffusion language model

Summary

The same workspace, on a different schedule. Both models carry the workspace's defining marks — a bounded set of verbalizable concepts, a minority of the residual, and a motor band near the output. It is the same object in both; the diffusion model only forms it later in depth and in sharper phases.

A shared sketch, not a running commitment. At the first denoising step the diffusion model's positions hold one common set of concepts — overlapping about 3 times as much as the autoregressive model's — but none is yet committed to the word it will become. It spreads a rough plan across the whole span and sharpens it over the denoising steps, where the autoregressive model fixes each token locally as it goes.

Spatial and transient, not sequential and persistent. The autoregressive model carries its workspace forward along the sequence; the diffusion model rebuilds the residual each denoising step and shares it across all positions at once. The same lens reads a workspace spread over space in one case and carried through time in the other.

Part 1 · A workspace in a diffusion model

The workspace is defined by a few properties. Its contents are verbalizable — the J-lens reads them as vocabulary tokens. It is limited in capacity — a bounded, small part of the residual, with most of the model's features outside it. And near the output the averaged Jacobian approaches the identity, so the plain logit lens already reads the state; the paper calls this the motor band. The J-lens itself multiplies a layer's activation by that Jacobian before unembedding and decomposes the result into a sparse, non-negative set of vocabulary directions — the J-space.

We measure these on both 8B models. Capacity is bounded and the verbalizable part is a minority of the residual in both (22 (Llama), 11 (LLaDA) concepts above a tenth of the top weight; 22% and 9% of the norm). The J-lens and logit lens agree only near the output (0.00→0.40 for Llama, 0.00→0.38 for LLaDA); in the autoregressive model that agreement builds through the middle layers, while in the diffusion model it appears only at the final layers.

The workspace is present in the diffusion model, with the same bounded, output-aligned character as in the autoregressive one. The two differ in where in depth it forms, not in whether it forms — the first of the three differences the rest of the study fills in.

Part 2 · The workspace across depth

A masked diffusion model predicts only at a [MASK], so we read its workspace there. The first question is how the workspace forms through the network's layers: where it becomes readable, and how its structure changes with depth.

J-lens sharpening

On the prompt “The capital of France is”, the J-lens reads the answer out of the residual stream earlier than the logit lens, because it multiplies in the Jacobian the logit lens omits. In Llama the answer appears in the mid-upper layers; in LLaDA it appears only at the late layers, read at the [MASK] — the workspace becomes verbalizable later in the diffusion model.

Similarity between layers

To read the structure directly, we measure linear CKA between the J-space representations (Jℓ hℓ) at every pair of layers. High similarity means two layers hold the same kind of representation; this is the grid the paper uses to read off the sensory / workspace / motor bands. The diffusion model shows these bands more sharply than the autoregressive one.

The diffusion model forms its readable workspace later in depth and separates that depth into sharper phases than the autoregressive model, whose representation blends smoothly from layer to layer.

Part 3 · The workspace across denoising steps

A diffusion model generates over denoising steps, so its workspace has a second axis the autoregressive model lacks. We read the same J-space at each step of the denoising process, and then intervene on it step by step.

The J-space across layers and steps

This figure reads the pooled J-space at a chosen layer and, for the diffusion model, a chosen denoising step, on the prompt “Scientists think the dinosaurs went extinct because a massive”. At the late layer both models hold concepts related to the answer — asteroid, collision, earth — and the diffusion model's reading shifts across steps as tokens commit, from vaguer to sharper. Earlier layers are less word-like, matching the motor band forming late. Beneath the two columns, LLaDA's generation fills in step by step, so each step's workspace sits beside the tokens it commits; the autoregressive column has only the depth axis.

Similarity between steps

Measuring CKA between the J-space at every pair of denoising steps shows how the workspace moves through time. It is neither fixed nor random: the first step is the most distinct, and after it neighbouring steps are alike while similarity falls off with distance.

Steering the scaffold

If the workspace holds a plan, injecting a concept into it should redirect the generation — the paper's directed modulation, here on the step axis. We add a set of J-lens concept directions to the residual at a middle layer during denoising, but only from a chosen step onward, and read off the result on a longer prompt.

The workspace is malleable early and fixed late. Injecting a concept from the first step rewrites the theme; the same injection applied only in the final steps barely moves the committed text, because the scaffold has already set the plan. Together with the step-to-step gradient, this says the diffusion model's residual is efficient only within a step — a transient workspace it rebuilds and commits as it goes, where the autoregressive model carries one forward along the whole sequence.

Part 4 · The workspace across positions

The last axis is space. Broadcast is the defining property of a global workspace: its contents are available at once to many consumers. In an autoregressive model that sharing runs along the sequence, one active position at a time; in a diffusion model a single denoising step reads every position together, so the workspace can be shared across the whole span at once.

Broadcast across positions

We test this directly. At the first denoising step every generated slot is still masked; we read the top concepts of each slot's workspace, then let the model finish so we know the word each slot becomes. For the autoregressive model we read, for each generated token, the workspace at the position that decides it — one step before the token.

The two use the workspace in opposite ways. In the diffusion model the positions hold a shared sketch — their top concepts overlap 0.21 against the autoregressive model's 0.07, about 3 times as much — but the sketch is not yet localized: a slot lists the specific word it will become only 8% of the time. The autoregressive model is the opposite, listing the coming token at its deciding position 58% of the time while sharing little between positions. The diffusion model spreads one rough plan across every position at once and resolves the individual words out of it over the denoising steps; the autoregressive model fixes each word locally as it reaches it. In the terms we began with, the autoregressive model carries a persistent workspace forward through time, while the diffusion model spreads a transient one across space — the same lens onto a slightly different object.

Part 5 · The two workspaces on reasoning

The earlier parts describe the shape of the workspace; the last question is what it holds on harder content, and whether the two architectures arrive at the same place. We generate each model's answer and read its J-space concepts across depth on two sets of prompts — multi-step reasoning and long-range dependencies — and compare them side by side. A shared read-layer control sets the depth for both models at once, which makes the legibility schedule concrete: the autoregressive model resolves the answer's field in the mid-stack, the diffusion model only in its final layers.

Reasoning

Ten reasoning prompts — categories, opposites, a syllogism, an analogy, a conditional — each read at the position that produces the answer. Both models answer all 10 correctly. At a shared read layer the concepts favour whichever model is legible there — Llama through the mid layers, LLaDA at the end — and where each is legible the field names the same content.

Long-range dependencies

Long-range prompts test whether the workspace binds information across a distance: a coreference, a bound attribute, a distant fact, a list index, a long subject, and a negated reference. Both models resolve all 6, and the J-space at the resolving position names the correct referent in each.

Read the same way, the two workspaces reach the same answers. The architectures differ in where the workspace forms, how it moves across denoising steps, and how it is shared across positions — the earlier parts — but on this content the concepts they surface and the answers they commit to line up, autoregressive and diffusion alike.

Appendix · Exact conditions

Every test above, with its read site, prompts, and layers or steps.

Averaged Jacobian by forward-mode JVP (checked against reverse-mode to ~1e-8); norm is each model's final RMSNorm; concepts are non-negative gradient pursuit over word-token candidates. Differences from the paper: dozens of prompts rather than ~1000, and the diffusion model is read at a [MASK] over masked prompts.