Taming Outlier Tokens in Diffusion Transformers
Xiaoyu Wu1*
Yifei Wang1*
Tsu-Jui Fu2
arXiv:2605.05206v1 [cs.CV] 6 May 2026
1
Liang-Chieh Chen2
Rice University
2
Zhe Gan2
Chen Wei1
Apple
Abstract We study outlier tokens in Diffusion Transformers (DiTs) for image generation. Prior work has shown that Vision Transformers (ViTs) can produce a small number of high-norm tokens that attract disproportionate attention while carrying limited local information, but their role in generative models remains underexplored. We show that this phenomenon appears in both the encoder and denoiser of modern Representation Autoencoder (RAE)-DiT pipelines: pretrained ViT encoders can produce outlier representations, and DiTs themselves can develop internal outlier tokens, especially in intermediate layers. Moreover, simply masking high-norm tokens does not improve performance, indicating that the problem is not only caused by a few extreme values, but is more closely related to corrupted local patch semantics. To address this issue, we introduce Dual-Stage Registers (DSR), a register-based intervention for both components: trained registers when available, recursive test-time registers otherwise, and diffusion registers for the denoiser. Across ImageNet and large-scale text-to-image generation, these interventions consistently reduce outlier artifacts and improve generation quality. Our results highlight outlier-token control as an important ingredient in building stronger DiTs.
1
Introduction
Vision Transformers (ViTs) [7] have emerged as a foundational architecture in computer vision. Initially, they demonstrated strong performance on visual understanding tasks, including supervised classification (e.g., ViT [7], DeiT [32]), self-supervised representation learning (e.g., DINO models [3, 18, 25]), and contrastive language-image pretraining (e.g., CLIP [22], SigLIP [38, 34]). These successes gradually shifted the default backbone for visual recognition models from convolutional networks such as ResNets [10] to Vision Transformers [4, 23]. More recently, ViTs have shown remarkable effectiveness in generative modeling, particularly for diffusion models [27, 11, 28]. Starting from DiT [20], Transformer-based generators have increasingly replaced UNet architectures [11, 29, 24] in both latent-space [17, 8] and pixel-space diffusion models [16]. Beyond the generator itself, recent work such as RAE [39] further proposes using pretrained ViT encoders to project images into latents, replacing the UNet-based VAE [24], the last remaining convolutional component in state-of-the-art diffusion pipelines. A commonly cited advantage of Transformers is their ability to model global interactions and their favorable optimization properties, albeit at the cost of weaker inductive biases. Recent studies, however, show that this flexibility can introduce nontrivial artifacts in the feature maps. In particular, Vision Transformers Need Registers [5] identifies the ubiquitous presence of outlier tokens in supervised and self-supervised ViTs (e.g., DeiT-III [33], OpenCLIP [13], DINOv2 [18]). These tokens exhibit abnormally large norms, absorb a disproportionate amount of attention, and yet carry little semantic information, resulting in artifacts in feature maps that hinder downstream utilization, especially for dense visual tasks. This behavior is closely related to the attention sink hypothesis in * Equal Contribution
Preprint.
the context of large language models [35], and can be effectively mitigated by introducing register tokens, i.e., dedicated non-patch tokens that explicitly serve as attention sinks during training. However, investigations of outlier tokens have thus far focused almost exclusively on recognition models, leaving their role in generative models largely unexplored. This gap is particularly striking given the growing prevalence of Transformers in generative modeling and the heightened importance of locality and spatial fidelity in generation compared to recognition. In this work, we show that outlier token effects are pervasive in both ViT-based autoencoder and diffusion Transformers. For ViT-based autoencoders, we focus on RAE [39], which has pretrained representation encoders. We empirically confirm the existence of outlier tokens in the encoders and demonstrate their negative impact on generation quality. Using the official DINOv2 implementation [18, 5], we compare DINOv2 encoders with and without register tokens and observe consistent improvements when registers are enabled. For SigLIP 2 [34], which does not include registers by default, we introduce recursive test-time registers [14] and show that this modification, too, leads to improved generation. For Transformer-based generators, identifying outlier tokens is more subtle. In supervised and self-supervised ViTs, inputs and outputs differ in semantic abstraction, and prior work reports that outlier tokens typically emerge in late layers. In contrast, diffusion models operate on inputs and outputs of the same modality—local pixels or latents. Surprisingly, we find that outlier tokens still arise, but predominantly in intermediate layers of the diffusion Transformer. By introducing register tokens into the generator, we achieve consistent gains across a variety of diffusion architectures, including SiT [17], JiT [16], and RAE-based [39, 31] designs. We refer to this unified use of register tokens in both the representation encoder and the diffusion generator as Dual-Stage Registers (DSR). Overall, our results show that DSR provide a simple yet effective way to stabilize Transformer-based diffusion pipelines by addressing outlier tokens in both the tokenizer and the denoiser. This unified intervention consistently improves generation quality across settings. Concretely, for RAE-DiT with SigLIP2-B, it reduces ImageNet-256 FID from 5.89 to 4.58 and improves GenEval from 0.426 to 0.466 on a large-scale text-to-image task.
2
Related Work
Outlier Tokens in Transformers. In Vision Transformers, Darcet et al. identify a small fraction of patch tokens with unusually large feature norms, referred to as outlier tokens [5]. These tokens often carry limited local patch information while inducing irregular attention patterns. To mitigate this effect, they introduce extra learnable register tokens during training, which absorb the outlier behavior and lead to cleaner attention over image patches [5]. Building on this finding, subsequent work traces these outliers to a sparse set of register neurons and proposes a training-free alternative: at inference time, the corresponding activations can be shifted into an additional untrained token, recovering much of the benefit of trained registers in models released without them [14]. Related sink-like phenomena have also been studied in language models. In this setting, attention sinks attract disproportionate attention and have been linked to large residual-stream activations, motivating neuron-level or activation-level test-time interventions [35, 30, 36]. More recent work explores the connection between sink behavior and outlier-driven dynamics [21]. Together, these studies suggest that outlier or sink tokens are a recurring phenomenon in Transformer architectures, and register-augmented ViT backbones have become increasingly common in practice [5]. By contrast, in Diffusion Transformers (DiTs), it remains unclear whether similar outlier tokens emerge during denoising and whether they affect generation quality. Our work addresses this gap. Diffusion Transformers Across Different Input Spaces. DiTs are originally trained in variational autoencoder (VAE) latent spaces for efficiency, but the tokenizer can bottleneck reconstructed image quality because its latents may not fully preserve semantic information [24, 20]. Subsequent work improves latent DiT pipelines through representation alignment with visual features from vision foundation models, improving training efficiency and generation quality [37]. More recent work further enables end-to-end joint optimization of the VAE tokenizer and diffusion model [15]. Apart from VAE space, RAE-based pipelines replace VAEs with pretrained representation encoders paired with learned decoders, yielding semantically richer latents and faster convergence in DiT training [39]. In parallel, another line of work revisits pixel-space diffusion to remove the need of 2
Layer 2
Input Image
Layer 4
Layer 6
Layer 8
Layer 10
Layer 12
Output
Norm Map
PCA Map
Figure 1: Outlier tokens in ViT-based autoencoders. We visualize token-norm maps across the layers of the SigLIP2-B encoder. Severe high-norm tokens emerge in the last few layers: the penultimate layer shows the strongest outlier pattern, while the final output becomes somewhat more stable, potentially due to the reconstruction-related training objective in SigLIP2. tokenizers, motivated by the desire to avoid lossy encode–decode stages and simplify the modeling stack. These approaches train diffusion transformers directly on images and show competitive performance when paired with appropriate architectures and training recipes [16, 1].
3
Outlier Tokens in RAE-DiT
RAE-DiT consists of two Transformer-based stages: a ViT-based representation encoder that serves as the tokenizer, and a diffusion Transformer that generates in the representation space. We examine outlier tokens in both stages and find distinct outlier patterns in the encoder and the diffusion generator. 3.1
Outlier Tokens in ViT-based Autoencoders
Outlier tokens have been observed in standard ViTs, where they typically appear in the final layers and can be mitigated by register tokens [5]. In RAE-DiT, such outliers in encoders are especially relevant because the encoder features define the representation space used for diffusion training. As shown in Fig. 1, using SigLIP2-B as an example, we observe severe high-norm tokens in the last few encoder layers. The penultimate layer exhibits the strongest outliers, while the final output becomes more stable, potentially due to the reconstruction-related training objective in SigLIP2. These observations suggest that outlier tokens already exist in the encoders before diffusion training. 3.2
Outlier Tokens in Transformer-based Generators
Beyond prior analyses of vision foundation models (standard ViTs), we identify a consistent outliertoken phenomenon in the RAE-DiT pipeline: a small subset of tokens in the representation encoder exhibits abnormally large norms, and this effect is further amplified when the encoder features are used as conditioning for diffusion training. As shown in Fig. 2, for RAE-DiT with SigLIP2-B as the encoder, outliers concentrate in intermediate layers, and their severity decreases as the diffusion noise level increases, suggesting a compounding mechanism in which encoder anomalies are amplified by the denoising objective. This differs from the standard ViT pattern discussed above, where outliers are typically most pronounced in the final layers [5]. To test whether the degradation is merely an extreme-value effect—namely, whether a small number of large-loss tokens dominate optimization—we apply a simple token-level loss masking strategy in representation-space diffusion training. Let z0 ∈ RN ×d denote the clean token representation, and let zt = α(t)z0 + σ(t)ϵ be the noisy tokens, where ϵ ∼ N (0, I) and t ∈ [0, 1]. Following prior work [16], the generator predicts ẑ0 = xθ (zt , t, c), which we convert to a v-prediction as v̂θ = (ẑ0 − zt )/(1 − t); the corresponding target is v = (z0 − zt )/(1 − t). We then mask tokens according to their representation norms, defining m_i = \mathbf {1}\!\left [\|z_{0,i}\|_2 \le \tau \right ]. The masked training objective is \mathcal {L}_{v}^{\mathrm {mask}}(\theta ) = \mathbb {E}_{z_0,\epsilon ,t}\!\left [ \frac {1}{\sum _i m_i}\sum _{i=1}^{N} m_i \left \| \hat {v}_\theta (z_t,t,c)_i - v_i\right \|_2^2 \right ] . \label {eq:v_loss_masked}
3
(1)
Layer 3
Layer 13
Layer 21
Layer 25
Layer 26
Layer 27
Output
Decoded Image
t = 0.1
t = 0.3
t = 0.5
t = 0.7
t = 0.9
Figure 2: Outlier Tokens in Transformer-based Generators. We visualize token-norm maps of RAE-DiT with a SigLIP2-B encoder, across different diffusion noise scales and encoder layers. We find that high-norm outliers concentrate in the intermediate layers, while their severity decreases as the diffusion noise level increases. This pattern differs from prior observations in standard ViTs, where artifact tokens are typically most pronounced in the final layers. Table 1: Effects of loss masking. Training Strategies RAE-DiT-XL (SigLIP2-B) + loss masking (τ = 100)
% of tokens filtered
FID↓
IS↑
Prec.↑
Rec.↑
0% 0.1%
5.89 6.06
156.54 152.72
0.686 0.686
0.562 0.562
In other words, tokens with unusually large representation norms are excluded from the diffusion loss. If the degradation were driven primarily by a small number of extreme-loss tokens, this masking strategy would be expected to substantially mitigate the problem. This method simply discards the training signal for tokens identified as encoder outliers. However, as shown in Tab. 1, masking does not improve generation quality, indicating that the degradation cannot be explained primarily by a few extreme loss values. Instead, we hypothesize that outlier tokens are a symptom of corrupted local patch information: removing supervision at those positions cannot restore the missing patch-level semantics, and may instead further weaken local learning. This view is supported by prior findings that generation quality depends on preserving local spatial structure in the representation space. Analyses in iREPA suggest that patchwise structure, rather than global semantics alone, is closely tied to generative performance [26]. Studies of ViT feature artifacts show that high-norm or outlier tokens can harm dense prediction tasks such as segmentation, which depend on accurate patch semantics [5]. Together, these observations motivate us to focus in the following section on restoring patch-level semantics, rather than suppressing extreme values.
4
Improving Generation with Dual-Stage Registers (DSR)
The results above suggest that the degradation cannot be explained solely by a few extreme token losses. Instead, a more plausible explanation is that outlier tokens reflect degraded local patch semantics and spatial structure in the representation space. This interpretation is consistent with prior findings that local structure is important for generative quality [26], and with analyses of artifact tokens in ViTs that connect such anomalies to failures in dense prediction [5]. Motivated by this view, we introduce Dual-Stage Registers (DSR), a lightweight intervention that absorbs token-level artifacts and stabilizes patch representations across the RAE-DiT pipeline. As shown in Fig. 3, DSR patches both sides of the pipeline: the vision encoder and the diffusion transformer. 4
1
2
...
1
N
2
3
...
5 ... N
4
Register Token
ViT Encoder
Diffusion Transformer
Noisy Image Token Image Token
...
1
2
3
4
5
Dropped Token
...
... N
Figure 3: Framework of our Dual-Stage Registers (DSR) method. Our DSR method patches both the vision encoder and the diffusion model with register tokens. The encoder uses a test-time register token, which is inserted only at inference time, while the diffusion model uses 36 trained register tokens, which are learned during diffusion training. During training, we discard the encoder-side register-token outputs before feeding encoder features into the diffusion model. During inference, we discard register-token outputs from both modules and keep only the image-token outputs. Layer 3
Layer 13
Layer 21
Layer 25
Layer 26
Layer 27
Output
Decoded Image
RAE-DiT-XL (DinoV2-B) w/o reg
RAE-DiT-XL (DinoV2-B) w/ reg
Figure 4: Norm map comparison for validating trained encoder registers. We compare RAEDiT(DINOv2-B) w/o and w/ trained encoder registers, at a fixed timestep t = 0.5. We find that introducing trained registers consistently suppresses high-norm token outliers and improves the quality of patch-level representations, which in turn leads to stronger downstream generation. 4.1
Registers in Vision Encoders
We begin by validating the effect of trained registers in a vision encoder where such registers are available. Specifically, we compare DINOv2 [18] trained with and without register tokens1 , measuring outlier severity and downstream generation quality when a DiT is trained on the resulting representation space. As shown in Fig. 4 and Tab. 2, trained registers reduce high-norm token artifacts and improve the quality of patch-level representations, leading to stronger generative performance. Table 2: Comparison of RAE-DiT (DINOv2-B) on ImageNet, w/ and w/o encoder register tokens. Training Strategies
FID↓
IS↑
Prec.↑
Rec.↑
RAE-DiT-XL (DINOV2-B-no-reg) RAE-DiT-XL (DINOV2-B-with-reg)
4.16 3.95
203.46 216.84
0.787 0.737
0.525 0.547
We next study SigLIP2 [34], a widely used vision-language encoder in recent multimodal models, to test whether our approach remains effective in a more practically important setting. In our experiments, the outlier phenomenon is substantially more severe for SigLIP2, while the off-the-shelf checkpoints do not provide trained registers. We therefore adopt test-time register tokens (TTR), following prior work [14], by appending an additional token to the input sequence at inference time without further encoder training. Empirically, we find that the SigLIP2-So400 pipeline exhibits two distinct sources of outliers, as discussed in Appendix Sec. A. To address this, we apply TTR recursively: we first use TTR to stabilize the encoder output, and then apply it again to the resulting representation when a second-stage outlier pattern is detected. As shown in Tab. 5, this recursive TTR scheme consistently reduces outlier severity and improves generation quality. 4.2
Registers in Diffusion Transformers
While encoder-side registers substantially reduce representation outliers, we still observe persistent outlier tokens inside the diffusion transformer itself, especially in intermediate layers, as shown in 1 For this comparison, we use the released DINOv2 checkpoints with and without register tokens.
5
Layer 3
Layer 13
Layer 21
Layer 25
Layer 26
Layer 27
Output
Decoded Image
RAE-DiT-XL (SigLipV2)
RAE-DiT-XL (SigLipV2) + test-time reg
RAE-DiT-XL (SigLipV2) + diffusion reg
RAE-DiT-XL (SigLipV2) + DSR
Figure 5: Norm map comparison across variants. We compare the baseline with two register-token configurations: adding test-time registers in the encoder only, and further adding trained registers in the diffusion model. We find that outliers in the norm map are suppressed only when both sources of outliers are addressed, i.e., when registers are applied to both the encoder and the diffusion model. Layer 3
Layer 13
Layer 21
Layer 25
Layer 26
Layer 27
Output
Decoded Image
RAE-DiT-XL (SigLipV2)
RAE-DiT-XL (SigLipV2) + test-time reg
RAE-DiT-XL (SigLipV2) + diffusion reg
RAE-DiT-XL (SigLipV2) + DSR
Figure 6: PCA map comparison across variants. We observe that adding test-time registers in the encoder yields a strong and visible improvement in the PCA map. Further adding trained registers in the diffusion model brings some additional improvements. Fig. 5. This suggests that the generator can also develop a small set of high-norm tokens, potentially encouraged by the global aggregation behavior of self-attention and the need to maintain global context throughout the denoising process [35]. Motivated by this observation, we introduce a small number of trainable diffusion registers into the diffusion transformer. These registers are learned jointly with the generator and removed at inference time, following common practice in prior work on registers [5, 14]. With diffusion registers, the internal outlier pattern is largely eliminated, as shown in Fig. 5 and Fig. 9. Intermediate-layer semantics are also improved, as illustrated by the PCA visualization in Fig. 6, and generation quality improves consistently, as reported in Tab. 3. More precisely, we find that diffusion registers yield consistent improvements across a wide range of settings. In particular, they improve performance for every input-space variant we evaluate, including pixel space, VAE latents, and multiple representation encoders. This suggests that the gains are not tied to any particular encoder or representation family. Instead, the broad effectiveness of diffusion registers indicates that they address token-level artifacts arising within the diffusion transformer itself, rather than serving as a modification tailored to a specific representation. 6
Table 3: Effect of diffusion registers across different DiT on ImageNet. Baseline rows are evaluated without diffusion registers; indented rows add diffusion registers on top of the same setting. Method
FID↓
IS↑
Prec.↑
Rec.↑
RAE-DiT-XL (DINOv2-B, w/ encoder reg) + diffusion reg RAE-DiT-XL (SigLIP2-B) + diffusion reg RAE-DiT-XL (SigLIP2-B, w/ test-time reg) + diffusion reg
4.11 3.92 5.89 5.33 4.63 4.58
226.44 226.92 156.54 166.2 177.2 165.99
0.775 0.773 0.686 0.702 0.748 0.725
0.529 0.542 0.562 0.556 0.542 0.56
VAE-SiT-XL + diffusion reg
16.05 14.47
70.11 78.50
0.550 0.554
0.647 0.651
JIT-H + diffusion reg
30.34 23.14
22.34 26.36
0.424 0.475
0.621 0.611
Table 4: Register tokens vs. in-context conditioning across DiTs and input spaces. In-context conditioning yields smaller gains as the input representation becomes more semantic, whereas register tokens provide consistent improvements. Method
FID↓
IS↑
Prec.↑
Rec.↑
RAE-DiT-XL (SigLIP2-B) + diffusion reg + diffusion in-context condition
5.89 5.33 5.79
156.54 166.2 164.99
0.686 0.702 0.696
0.562 0.556 0.546
VAE-SiT-XL + diffusion reg + diffusion in-context condition
16.05 14.47 14.31
70.11 78.50 77.04
0.550 0.554 0.559
0.647 0.651 0.646
JIT-H + diffusion reg + diffusion in-context condition
30.34 23.14 15.51
22.34 26.36 33.05
0.424 0.475 0.519
0.621 0.611 0.618
We also observe a close relationship between diffusion registers and in-context conditioning [16]. Both augment the input with extra tokens that can carry global context: in-context conditioning uses input-dependent tokens, while registers provide dedicated learnable slots. In Tab. 4, we find that in-context conditioning is most competitive when the input space contains less semantic or structured token representations, while diffusion registers remain beneficial as representations become more semantic. These results suggest that registers provide a simpler and more robust mechanism for supplying global capacity, without relying on the semantics or formatting of in-context tokens.
5
Experiments
We evaluate DSR on ImageNet-1K class-conditional generation at 256 × 256 and text-to-image generation. For ImageNet-1K, we follow the RAE setting [39] and keep the same training epochs, learning rate, and model architecture for fair comparison; our main configuration uses SigLIP2-B as the encoder. While our encoder is slightly different due to the test-time register modification, we use the same decoder as in RAE; we find that retraining the decoder has limited impact (Appendix Sec. C). Additionally, we use data prediction with velocity loss to train our model and RAE baseline, which has been proven effective in JiT [16]. For ImageNet-1K, Tab. 6 shows that DSR substantially reduces gFID for RAE (SigLIP2-B). Combining DSR with the DDT head yields further gains and achieves competitive performance. Fig. 7 plots the per-epoch trajectory and shows that DSR reaches comparable quality with 4× fewer epochs. 7
FID-50k ImageNet @ 256
8
RAE-DiT-XL (SigLip2-B) RAE-DiT-XL (SigLip2-B)+DSR
7
Table 5: Test-time registers on RAE-DiT with SigLIP2B and SigLIP2-So400, trained for 80 epochs.
6
Training Strategies
FID↓
5
RAE-DiT-XL (SigLIP2-B) + test-time register
5.89 156.54 0.686 0.562 4.63 177.2 0.748 0.542
4× faster
4 40
80
120 160 200
320
IS↑
Prec.↑ Rec.↑
RAE-DiT-XL (SigLIP2-So400) 7.04 167.01 0.682 0.515 800
Training Epochs Figure 7: FID vs. epochs on IN-1K 2562 .
+ test-time register + test-time register (recursive)
6.66 166.88 0.687 0.527 6.48 163.35 0.684 0.531
Table 6: Class-conditional performance on ImageNet 256×256. Method
Epochs #Params Generation@256 w/o guidance Generation@256 w/ guidance gFID↓
IS↑
Prec.↑ Rec.↑ gFID↓
IS↑
Prec.↑ Rec.↑
Pixel Diffusion ADM [6] JiT-H/16 [16]
400 600
554M 953M
10.94 –
101.0 –
0.69 –
0.63 –
3.94 1.86
215.8 303.4
0.83 –
0.53 –
Latent Diffusion SiT-XL [17] REPA [37] REPA-E [15] RAE-DiTDH -XL (DINOv2-B)
1400 800 800 800
675M 675M 675M 839M
8.61 5.78 1.70 1.51
131.7 158.3 217.3 242.9
0.68 0.70 0.77 0.79
0.67 0.68 0.66 0.63
2.06 1.29 1.15 1.13
270.3 306.3 304.0 262.6
0.82 0.79 0.79 0.78
0.59 0.64 0.66 0.67
5.89 3.85 4.58 3.26 3.74 2.91 3.56 2.72
156.54 179.82 165.99 185.54 179.24 204.73 181.08 207.76
0.686 0.692 0.725 0.704 0.707 0.686 0.722 0.697
0.562 0.613 0.56 0.62 0.599 0.642 0.589 0.646
– 3.58 – 2.97 – 2.77 – 2.62
– 194.67 – 203.95 – 221.76 – 223.85
– 0.691 – 0.709 – 0.685 – 0.697
– 0.619 – 0.621 – 0.649 – 0.651
Latent Diffusion with Multi-modal encoder 80 RAE-DiT-XL (SigLiP2-B) 800 80 RAE-DiT-XL (SigLiP2-B) + DSR 800 80 RAE-DiTDH (SigLiP2-B) 800 80 RAE-DiTDH -XL (SigLiP2-B) + DSR 800
5.1
676M 676M 839M 839M
Scalability
A key motivation for decoupling the Transformer design from the task is to better leverage scaling [16, 20]. We reports ImageNet 256 × 256 results across model sizes (DiT-B/L/XL). As shown in Tab. 7, DSR consistently improves gFID across all parameter scales, while introducing only a minor increase in GFLOPs. We also demonstrate the scalability of DSR on Scale-RAE [31], which leverages the MetaQuery architecture [19] for text-to-image (T2I) generation and unified modeling. We use SigLIP2-B [34] as the encoder and train the model on the Scale-RAE dataset [31], which contains 24.7M synthetic images generated by FLUX.1-schnell [2].
Table 8: Text-to-image perf. on Scale-RAE. Models
GenEval↑
DPG-Bench↑
Baseline Ours
42.6 46.6
74.3 75.4
The training process takes 12 hours on a google cloud v5p-128 TPU for 10k training steps with a batch size of 2048. We then sample images from the trained model using text prompts and decode the generated latents with SigLIP-B pretrained decoder provided by RAE [39]. We evaluate the sampled images on GenEval [9] and DPG-bench [12]. As shown in Tab. 8, DSR achieves better performance than the original RAE baseline on both benchmarks. More details can be found in Appendix Sec. D. 5.2
Ablation Studies
We conduct an ablation study of diffusion registers on RAE-DiT-XL with SigLIP2-B on ImageNet 256×256. The baseline model uses no diffusion registers. We study two factors: the insertion depth, 8
Table 7: Scalability of DSR across DiT model sizes for RAE-DiT with SigLIP2-B on ImageNet 256×256. All models are trained for 100k iterations under the same training and evaluation protocol. For each size, we compare the baseline with DSR. Model
FID↓
IS↑
Prec.↑
Rec.↑
Gflops.↓
RAE-DiT-S (SigLIP2-B) + DSR
28.03 23.93
66.15 63.86
0.277 0.498
0.302 0.470
12.4 13.7+10.4%
RAE-DiT-B (SigLIP2-B) + DSR
20.36 9.81
78.76 110.23
0.539 0.637
0.493 0.543
46.6 51.16+9.9%
RAE-DiT-XL (SigLIP2-B) + DSR
5.89 4.58
156.54 165.99
0.686 0.725
0.562 0.560
238.1 262.9+10.4%
Table 9: Ablation on diffusion registers for RAE-DiT-XL (SigLIP2-B) on ImageNet 256×256. We vary the insertion starting block (with a fixed number of 36 registers) and the number of registers (with a fixed starting block of 8). Starting block
#Regs
FID↓
IS↑
Pre.↑
Rec.↑
Gflops ↓
-
0
5.89
156.54
0.686
0.562
238.1
0 8 16 24
36 36 36 36
5.54 5.33 5.49 5.68
163.45 166.20 168.01 159.04
0.701 0.702 0.700 0.685
0.552 0.556 0.549 0.560
272.8+14.6% 262.9+10.4% 253.0+6.3% 243.1+2.1%
8 8 8 8
1 4 36 100
6.16 5.47 5.33 5.58
153.95 165.55 166.20 165.61
0.685 0.700 0.702 0.701
0.555 0.553 0.556 0.546
238.9+0.3% 240.8+1.1% 262.9+10.4% 307.5+29.1%
defined as the transformer block from which registers are introduced, and the number of registers. In the first set of experiments, we fix the number of registers to 36 and vary the insertion depth. In the second, we fix the insertion depth to block 8 and vary the number of registers. As shown in Tab. 9, diffusion registers provide clear but non-monotonic gains, with a distinct sweet spot in both insertion depth and register count. When fixing the register count at 36, introducing registers too early or too late is less effective, while starting from block 8 gives the best performance. This suggests that registers are most useful when introduced in the early-to-middle part of the generator, where they can influence a substantial portion of the computation without interfering too strongly with the earliest layers. When fixing the starting block to 8, varying the number of registers shows a similar trend: very small numbers do not help, moderate numbers improve performance, the best results are obtained with 36 registers, and performance degrades at 100 registers. Overall, these results suggest that diffusion registers are most effective at an appropriate depth and capacity.
6
Conclusion
We study outlier tokens in Diffusion Transformers and show that they arise not only in pretrained vision encoders, but also within the diffusion model itself. Our analysis suggests that these outliers reflect degraded patch-level semantics rather than merely a few abnormally large token norms. Motivated by this insight, we introduce simple register-based interventions for both the encoder and the diffusion transformer. These interventions consistently improve training stability and generation quality in both ImageNet generation and large-scale text-to-image generation, and generalize across diverse input spaces and model variants. Overall, our results identify outlier-token control as an important ingredient for building stronger and more robust Transformer-based diffusion pipelines.
9
References [1] Alan Baade, Eric Ryan Chan, Kyle Sargent, Changan Chen, Justin Johnson, Ehsan Adeli, and Li Fei-Fei. Latent forcing: Reordering the diffusion trajectory for pixel-space image generation. arXiv preprint arXiv:2602.11401, 2026. [2] Black Forest Labs. Flux. https://github.com/black-forest-labs/flux, 2024. GitHub repository. [3] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 9650–9660, 2021. [4] Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. In International Conference on Learning Representation, 2022. [5] Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. arXiv preprint arXiv:2309.16588, 2023. [6] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780–8794, 2021. [7] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020. [8] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference on machine learning, 2024. [9] Dhruba Ghosh, Hannaneh Hajishirzi, and Ludwig Schmidt. Geneval: An object-focused framework for evaluating text-to-image alignment. Advances in Neural Information Processing Systems, 36:52132–52152, 2023. [10] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016. [11] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020. [12] Xiwei Hu, Rui Wang, Yixiao Fang, Bin Fu, Pei Cheng, and Gang Yu. Ella: Equip diffusion models with llm for enhanced semantic alignment. arXiv preprint arXiv:2403.05135, 2024. [13] Gabriel Ilharco, Mitchell Wortsman, Ross Wightman, Cade Gordon, Nicholas Carlini, Rohan Taori, Achal Dave, Vaishaal Shankar, Hongseok Namkoong, John Miller, Hannaneh Hajishirzi, Ali Farhadi, and Ludwig Schmidt. Openclip. [14] Nick Jiang, Amil Dravid, Alexei Efros, and Yossi Gandelsman. Vision transformers don’t need trained registers. arXiv preprint arXiv:2506.08010, 2025. [15] Xingjian Leng, Jaskirat Singh, Yunzhong Hou, Zhenchang Xing, Saining Xie, and Liang Zheng. Repa-e: Unlocking vae for end-to-end tuning of latent diffusion transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 18262–18272, 2025. [16] Tianhong Li and Kaiming He. Back to basics: Let denoising generative models denoise. arXiv preprint arXiv:2511.13720, 2025. [17] Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In European Conference on Computer Vision, pages 23–40. Springer, 2024. 10
[18] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193, 2023. [19] Xichen Pan, Satya Narayan Shukla, Aashu Singh, Zhuokai Zhao, Shlok Kumar Mishra, Jialiang Wang, Zhiyang Xu, Jiuhai Chen, Kunpeng Li, Felix Juefei-Xu, et al. Transfer between modalities with metaqueries. arXiv preprint arXiv:2504.06256, 2025. [20] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 4195–4205, 2023. [21] Zihan Qiu, Zeyu Huang, Kaiyue Wen, Peng Jin, Bo Zheng, Yuxin Zhou, Haofeng Huang, Zekun Wang, Xiao Li, Huaqing Zhang, et al. A unified view of attention and residual sinks: Outlier-driven rescaling is essential for transformer training. arXiv preprint arXiv:2601.22966, 2026. [22] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748–8763. PmLR, 2021. [23] Maithra Raghu, Thomas Unterthiner, Simon Kornblith, Chiyuan Zhang, and Alexey Dosovitskiy. Do vision transformers see like convolutional neural networks? Advances in neural information processing systems, 34:12116–12128, 2021. [24] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684–10695, 2022. [25] Oriane Siméoni, Huy V Vo, Maximilian Seitzer, Federico Baldassarre, Maxime Oquab, Cijo Jose, Vasil Khalidov, Marc Szafraniec, Seungeun Yi, Michaël Ramamonjisoa, et al. Dinov3. arXiv preprint arXiv:2508.10104, 2025. [26] Jaskirat Singh, Xingjian Leng, Zongze Wu, Liang Zheng, Richard Zhang, Eli Shechtman, and Saining Xie. What matters for representation alignment: Global information or spatial structure? arXiv preprint arXiv:2512.10794, 2025. [27] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256–2265. pmlr, 2015. [28] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019. [29] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020. [30] Mingjie Sun, Xinlei Chen, J Zico Kolter, and Zhuang Liu. Massive activations in large language models. arXiv preprint arXiv:2402.17762, 2024. [31] Shengbang Tong, Boyang Zheng, Ziteng Wang, Bingda Tang, Nanye Ma, Ellis Brown, Jihan Yang, Rob Fergus, Yann LeCun, and Saining Xie. Scaling text-to-image diffusion transformers with representation autoencoders. arXiv preprint arXiv:2601.16208, 2026. [32] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pages 10347–10357. PMLR, 2021. [33] Hugo Touvron, Matthieu Cord, and Hervé Jégou. Deit iii: Revenge of the vit. In European conference on computer vision, pages 516–533. Springer, 2022. 11
[34] Michael Tschannen, Alexey Gritsenko, Xiao Wang, Muhammad Ferjad Naeem, Ibrahim Alabdulmohsin, Nikhil Parthasarathy, Talfan Evans, Lucas Beyer, Ye Xia, Basil Mustafa, et al. Siglip 2: Multilingual vision-language encoders with improved semantic understanding, localization, and dense features. arXiv preprint arXiv:2502.14786, 2025. [35] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, 2024. [36] Itay Yona, Ilia Shumailov, Jamie Hayes, Federico Barbero, and Yossi Gandelsman. Interpreting the repeated token phenomenon in large language models. arXiv preprint arXiv:2503.08908, 2025. [37] Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, and Saining Xie. Representation alignment for generation: Training diffusion transformers is easier than you think. arXiv preprint arXiv:2410.06940, 2024. [38] Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, and Lucas Beyer. Sigmoid loss for language image pre-training. In Proceedings of the IEEE/CVF international conference on computer vision, pages 11975–11986, 2023. [39] Boyang Zheng, Nanye Ma, Shengbang Tong, and Saining Xie. Diffusion transformers with representation autoencoders. arXiv preprint arXiv:2510.11690, 2025.
12
A
Outlier Visualization on S IG LIP2-So400
We further analyze the norm distribution of S IG LIP2-So400 and observe that its outliers do not form a single homogeneous group. Instead, they appear as two clearly separated clusters, indicating that a single-pass filtering strategy is insufficient. As shown in Fig. 8, filtering once removes only the more salient cluster, while the other remains visible. To address this, we adopt a recursive procedure that re-localizes outlier neurons after the first filtering pass and then jointly filters all detected outliers, allowing us to suppress both groups more effectively.
Figure 8: Two outlier sources in the norm distribution of S IG LIP2-So400. We compute the ℓ2 norm of S IG LIP2-So400 output features on 10k randomly selected images from the ImageNet-1K validation set. Left: the original norm distribution shows two separated outlier groups. Middle: applying our filtering once removes only one group, leaving the other largely intact. Right: we therefore use a recursive procedure that re-localizes outlier neurons after the first pass and filters all detected outliers jointly, which effectively suppresses both groups.
B
More Visualization Layer 3
Layer 13
Layer 21
Layer 25
Layer 26
Layer 27
Output
RAE(SigLipV2)
RAE(SigLipV2) + test-time reg
RAE(SigLipV2) + diffusion reg RAE(SigLipV2) + test-time reg + diffusion reg
Figure 9: Quantitative measurements of outliers. We report the fraction of outlier tokens across layers under different setups. An outlier is defined as a token whose ℓ2 norm exceeds 2× the median token norm. To complement the qualitative visualization, we also provide a quantitative view of outlier behavior across layers. Figure 9 reports the fraction of outlier tokens under different setups, where we define an outlier as a token whose ℓ2 norm is greater than 2× the median token norm. These results show that outlier tokens are not uniformly distributed across depth, but instead exhibit clear layer-dependent patterns that vary across settings.
C
Retrained Decoder
Applying test-time registers alters the encoder feature distribution. In principle, this distribution shift may require retraining the Stage-2 decoder. Interestingly, we find that the shift is mild: retraining the decoder yields no visible improvement over using the pretrained decoder. 13
Table 10: Retraining the decoder. Retraining the Stage-2 decoder after applying test-time registers brings negligible gains, indicating that the encoder distribution shift is mild. Training Strategy
Decoder
rFID↓
gFID↓
IS↑
DiT-XL (SigLIP2-B) DiT-XL (SigLIP2-B) + DSR DiT-XL (SigLIP2-B) + DSR DiT-XL (SigLIP2-B) + DSR DiT-XL (SigLIP2-B) + DSR
Pretrained Pretrained Retrained (σ=0.8) Retrained (σ=1.5) Retrained (σ=2.0)
0.82 0.82 0.58 0.67 0.75
5.89 4.58 6.00 5.35 5.28
156.54 165.99 164.20 157.59 157.98
Method component optimizer learning rate schedule global batch size Training epoch Dataset max learning rate optimizer betas loss model Register num Register adding block Vision encoder Vision decoder Encoder test time register Training device TPU memory usage per chip Training speed
Scale-RAE LLM
DiT
Scale-RAE+DSR LLM DiT
AdamW cosine w/ warmup ratio 0.0134 2048 1 Scale-RAE dataset (24.7M) 5e-5 5e-4 5e-5 5e-4 (0.9, 0.999) (0.9, 0.95) (0.9, 0.999) (0.9, 0.95) autoregressive loss diffusion loss autoregressive loss diffusion loss Qwen2.5 1.5B DiT 2.4B Qwen2.5 1.5B DiT 2.4B 0 36 0 8 google/siglip2-base-patch16-256 nyu-visionx/RAE-siglip2-base-p16-i256-ViTXL-n08 No Yes v5p-128 22.34GiB 23.68GiB 4.00s/iter 4.09s/iter
Table 11: Training configuration comparison between Scale-RAE and DSR.
For SigLIP2-B, we search over augmented noise scales σ, a standard knob for balancing reconstruction quality and gFID in prior work [39]. As shown in Tab. 10, We find no setting that consistently outperforms the baseline, retraining under different σ yields results comparable to the pretrained decoder. It suggests that test-time registers only mildly shift the encoder feature distribution seen by the decoder.
D
More Results on Text-to-Image Experiments
Training configuration. Tab. 11 summarizes the computational resources required and the training hyperparameter configuration for the scaling-up experiment. Following Scale-RAE [31], We use SPMD sharding together with TorchXLA to train the LLM and DiT models. Due to limited computational resources, we use only the Scale-RAE dataset, which is one quarter the size of the data used in the original Scale-RAE [31], and train the model for one epoch. Step-to-step comparison. We also perform a step-to-step comparison between DSR and ScaleRAE baseline. The detailed comparison is shown in Fig. 11. We observe consistent improvements of DSR over the Scale-RAE [31] baseline on GenEval [9]. On DPG-Bench [12], DSR shows a steadier improvement. We further compare the training losses of the baseline and DSR. As shown in Fig. 10, the green curve corresponding to DSR is consistently lower and more stable overall, and the same trend is also visible in the diffusion loss. In contrast, the language-model loss of DSR exhibits more spikes during training, which may help explain why it sometimes underperforms the baseline on DPG-Bench, a benchmark that emphasizes following longer and more complex instructions. This 14
suggests that additional improvements on the language side, such as mitigating the effect of outlier tokens, may further enhance performance. We leave this direction for future work.
Figure 10: Training loss comparison between Scale-RAE and DSR. Top: total loss. Bottom left: image diffusion loss. Bottom right: language loss. RAE-DiT (SigLip2-B) RAE-DiT (SigLip2-B)+Our Method
74
DPG-Bench Score
GenEval Score
76
RAE-DiT (SigLip2-B) RAE-DiT (SigLip2-B)+Our Method
45 40 35 30 25
72 70 68 66 64 62
2000
4000
6000
Training Steps
8000
10000
2000
4000
6000
Training Steps
8000
10000
Figure 11: Step-to-step comparison between Scale-RAE and DSR on Geneval and DPG-Bench. Extended samples. We provide extended samples from Scale-RAE [31] and DSR on DPGBench [12] and Geneval [9] (Fig. 12, Fig. 13, Fig. 14, Fig. 15, Fig. 16, Fig 17, Fig 18,Fig 19). The prompts for generated images are: • During the twilight hour, an individual can be seen extending an arm towards the sky, pointing at a trio of wild birds gliding through the rich deep blue of the early evening sky. The birds’ silhouettes contrast distinctly against the fading light, their wings spread wide as they soar. The person is silhouetted against the dusky sky, creating a peaceful scene of human connection with nature. 15
• An expansive palace constructed from iridescent materials that shimmer with hues reminiscent of a vivid, Slime-like substance, majestically stands at the heart of a fantastical realm. Its towers twist skyward, defying conventional architecture with their organic, flowing forms. In the foreground, a field of exotic flowers blooms, each petal displaying an array of otherworldly colors that could have been plucked from a Lovecraftian spectrum, while overhead, a radiant sun bathes the surreal landscape in brilliant light. • a highly intricate and vibrant cityscape that reflects a fusion of Moebius’s imaginative design and Makoto Shinkai’s detailed animation style. The streets are aglow with neon signs in a kaleidoscope of colors, casting reflections on the glossy, rain-slicked pavements. Towering skyscrapers with glowing windows rise towards a starless night sky, as the artwork garners significant attention and praise on ArtStation. • Adjacent to each other in a room, a large rectangular bed draped in a navy-blue comforter sits parallel to a square-shaped nightstand with a matte finish. The nightstand holds an angular lamp and a small stack of hardcover books. The two pieces of furniture are positioned on a plush beige carpet that covers the majority of the floor space. • a photo of a cow. • a photo of a purple potted plant. • a photo of an elephant below a surfboard. • a photo of a white toilet and a red apple.
Figure 12: Scale-RAE baseline (top) vs. DSR (bottom) on DPG-Bench [12].
Figure 13: Scale-RAE baseline (top) vs. DSR (bottom) on DPG-Bench [12].
16
Figure 14: Scale-RAE baseline (top) vs. DSR (bottom) on DPG-Bench [12].
Figure 15: Scale-RAE baseline (top) vs. DSR (bottom) on DPG-Bench [12].
Figure 16: Scale-RAE baseline (top) vs. DSR (bottom) on GenEval [9].
17
Figure 17: Scale-RAE baseline (top) vs. DSR (bottom) on GenEval [9].
Figure 18: Scale-RAE baseline (top) vs. DSR (bottom) on GenEval [9].
Figure 19: Scale-RAE baseline (top) vs. DSR (bottom) on GenEval [9].
18