Registers Matter for Pixel-Space Diffusion Transformers

TL;DR

Register tokens were introduced for ViTs to address outliers in patch tokens. We notice that diffusion transformers don't have such outliers, suggesting that registers may not be needed there. But they do, and adding them creates high-norm tokens inside the registers themselves. Pixel-space DiTs benefit the most, while latent-space models barely change their performance.

In pixel-space DiTs, registers lower patch-token norms and smooth the features, with some tokens acting as norm sinks and others carrying global information. This also explains recent pixel-space DiTs with in-context class tokens: their gains likely come from this register-like behavior, not from extra class information.

Building on this, we propose Register Guidance, which uses the model's prediction without registers as a negative direction (instead of the unconditional prediction used in CFG) for consistent quality gains.

Background: outliers in ViTs

In ViTs (e.g. DINOv2), a few patch tokens attain anomalously high norms, usually in low-information regions like the background. This produces artifacts in the attention maps. We visualize these outliers for different models below, across depth.

These outlier tokens are colored red: they are static artifacts pinned to the same background locations across layers. For example, in DeiT they are present throughout, while in DINOv2 they emerge only at deeper blocks. Register tokens were proposed to absorb this behavior, and the bottom row shows that the artifacts disappear once registers are added.

Without registers

With registers

Depth step 23

Attention is computed as [CLS] → patches, averaged over heads.
The high-norm outliers are static artifacts that persist in the same background location across layers.

The puzzle: DiTs have no outliers
but benefit from registers

We find that DiTs, in both latent and pixel space, do not exhibit these high-norm patch-token outliers; their token norms are nearly uniform and their attention maps stay focused on the main objects. So registers shouldn't help, but they do.

Layer 5

Noise (t) more noise less noise t=0.1 (more noise)

We show attention over the main object only (the head most focused on it), across layers and noise levels.

Attention maps can be noisy and head-dependent. Token feature norms tell the story more cleanly: in ViTs a few patch tokens attain anomalously high norms, the outliers that motivate register tokens, whereas our diffusion transformers keep nearly uniform token norms, with no such spikes.

DINOv2 ViT-g token norms — Per-token feature norm vs. token index. **Top row:** ViTs show isolated high-norm outlier tokens. **Bottom rows:** diffusion transformers (pDiT, SiT, DiT, RAE) have nearly uniform norms, with no outliers.

DINOv2 ViT-B token norms — Per-token feature norm vs. token index. **Top row:** ViTs show isolated high-norm outlier tokens. **Bottom rows:** diffusion transformers (pDiT, SiT, DiT, RAE) have nearly uniform norms, with no outliers.

So registers seem unnecessary, yet they help. We find that register tokens show the largest improvements in pixel space, provide smaller gains in VAE space, and, interestingly, degrade performance in RAE-based models.

	RAE-space DiT^DH backbone	VAE-space SiT backbone	Pixel-space pDiT backbone
Base size
w/ reg.	7.48	9.40	5.30
w/o reg.	6.58	10.40	7.39
Large size
w/ reg.	4.44	2.38	2.69
w/o reg.	3.91	2.53	3.52

FID (↓ lower is better) with and without register tokens. Registers help most in pixel space, help less in VAE space, and hurt in RAE space. Bold marks the better setting per column.

Registers make intermediate features cleaner

We saw that pixel-space DiTs have no patch-token outliers, yet registers still help. What do they change?

Adding register tokens makes high-norm outliers appear, but they emerge within the register tokens themselves, while the patch tokens stay clean. Below, pDiT-B/16 without registers has uniform token norms. With registers, the first 32 (register) tokens spike, leaving the patches uniform.

pDiT-B/16 token norms without registers — Per-token feature norm vs. token index for pDiT-B/16. **Left:** without registers, norms are uniform. **Right:** with registers, the register tokens (shaded) become high-norm outliers while patch tokens remain uniform.

pDiT-B/16 token norms with registers — Per-token feature norm vs. token index for pDiT-B/16. **Left:** without registers, norms are uniform. **Right:** with registers, the register tokens (shaded) become high-norm outliers while patch tokens remain uniform.

And this matters: register tokens make the intermediate feature maps cleaner at high noise. Below, PCA of the patch features for pDiT-B/16: with registers the map is far less noisy than without. Drag the sliders to step through blocks and noise levels.

Image cat

Block 7

Noise (t) more noise less noise t=0.05

In-context tokens play the role of registers

Recent DiTs incorporate conditioning signals (e.g., text or class labels) by appending additional tokens to the sequence of image patches and processing them jointly through shared transformer layers. For example, JiT employs conditioning by adding duplicated class embeddings to the input sequence, leading to notable improvements in generation quality. This raises the question of whether such in-context tokens implicitly function as register tokens.

Interestingly, we find that in-context tokens become outliers, as registers did.

JiT-B/16 token norms with incontext conditioning — Per-token feature norm vs. token index for pDiT-B/16 and JiT-B/16. **Left:** pDiT-B/16 with registers. **Right:** JiT-B/16 with in-context class tokens, the register and in-context tokens (shaded) become high-norm outliers while patch tokens remain uniform.

We dive deeper. The figure below shows classification accuracy for in-context tokens. It is highly diverse: some tokens achieve high accuracy (≈ 0.9), others moderate accuracy (≈ 0.4), and some very low accuracy (≈ 0.02). These results suggest the following: (a) tokens with the highest norms act as norm sinks, yielding the lowest accuracy, and (b) tokens with moderate norms encode diverse information about the image beyond class-specific features.

**Linear probing accuracy for in-context tokens.** Bar height shows classification accuracy, while color indicates token norm magnitude. High-norm tokens yield low accuracy, suggesting they act as norm sinks.

To confirm that in-context tokens capture broad information about the image, rather than only class-specific features as intended, we visualize attention maps for the in-context tokens below. We observe that these tokens attend to distinct semantic regions of the image. In this example, some tokens focus on background elements, such as the jungle, while others attend to object parts, such as the tiger’s head and back, or foreground objects, such as the stick.

In-context token 1 attention map — Attention maps for selected in-context tokens. Different in-context tokens attend to different semantic regions, including the background, object parts, and foreground elements.

Register guidance

Registers help the model: they give it an extra internal signal that boosts generation quality, rather than simply acting as additional class information. Our idea is simple: let's amplify this signal.

We propose Register Guidance, a guidance method that uses the model's prediction without registers as a negative direction, instead of relying on an unconditional prediction as in classifier-free guidance (CFG). Pushing away from the register-free prediction strengthens exactly the signal the registers contribute, and it combines naturally with CFG.

Pseudocode for register guidance: training with register dropout, RG at inference, and combined RG + CFG at inference — During training we randomly drop the registers (probability p) so the model learns both register and register-free predictions. At inference, register guidance (RG) amplifies the difference between the two, and combines naturally with classifier-free guidance (CFG).

Below we present a generation example of CFG + RG. Use the sliders to vary the subject, the CFG scale, and the RG scale: CFG sharpens the class identity, while RG improves visual details and object structure.

Generated sample for the selected subject, CFG and RG scales

Image Macaw

CFG 1.0

RG 0.0

JiT-L/16. CFG sharpens class identity, RG strengthens the register signal.

BibTeX

@article{starodubcev2026registers,
  title   = {Registers Matter for Pixel-Space Diffusion Transformers},
  author  = {Starodubcev, Nikita and Sudakov, Ilia and Drobyshevskiy, Ilya
             and Babenko, Artem and Baranchuk, Dmitry},
  journal = {arXiv preprint arXiv:2605.16147},
  year    = {2026}
}