PaperlayerView exact revision on arXiv

Source-linked AI summary

Registers Matter for Pixel-Space Diffusion Transformers

Nikita Starodubcev, Ilia Sudakov, Ilya Drobyshevskiy, Artem Babenko, Dmitry Baranchuk

arXiv:2605.16147v2cs.CV

TL;DR

Diffusion Transformers (DiTs) resemble Vision Transformers, prompting questions about patch-token outliers and register-token utility. This paper analyzes registers across pixel- and latent-space DiTs and introduces Register Guidance, finding that registers benefit DiTs despite absent outliers, especially in pixel space, while improving feature-map smoothness and generation quality.

  • Problem

    It is unclear whether DiTs exhibit ViT-like patch-token outliers and whether register tokens improve diffusion models, particularly in pixel-space training.

  • Method

    The paper analyzes register-token behavior across DiTs and introduces Register Guidance, using register-aware predictions to steer generation toward improved visual structure and coherence.

  • Results

    Registers improve DiT generation despite absent patch-token outliers, with the strongest gains in pixel space and smoother intermediate feature maps at high noise levels.

  • Takeaways & Limitations

    Recent pixel-space DiTs and text-conditioning sequences implicitly use register-like tokens, suggesting such mechanisms help account for their strong empirical performance.

  • Takeaways & Limitations

    The study covers only certain aspects of register behavior, leaving considerable scope for further investigation.

Abstract

from arXiv · show

Vision Transformers (ViTs) are known to exhibit high-norm patch-token outliers that degrade feature map quality, a problem effectively mitigated by register tokens. As diffusion models increasingly adopt transformer architectures and move toward pixel-space training, they become closer in form to ViTs, raising the question of whether register tokens are also useful for Diffusion Transformers (DiTs). In this work, we show that DiTs differ from ViTs in a key respect: they do not exhibit patch-token outliers but still benefit from registers. Interestingly, registers are more effective in pixel-space DiTs than in latent-space DiTs. By analyzing intermediate representations, we find that register tokens produce cleaner feature maps at high noise levels, which may contribute to their effectiveness in pixel-space generation. We further observe that recent pixel-space DiT architectures implicitly incorporate register-like mechanisms, which may partially account for their strong empirical performance. Motivated by these observations, we propose Register Guidance, a technique that amplifies the contribution of register tokens responsible for improving visual structure and coherence.

1 Introduction

Diffusion Transformers lack the high-norm patch-token outliers seen in ViTs, yet register tokens still help, especially in pixel-space models. Registers smooth high-noise feature maps, specialize into functional roles, and motivate Register Guidance for improved structure and coherence.

  • Findings: Unlike ViTs, DiTs in latent and pixel spaces lack noticeable high-norm patch-token outliers, while added registers develop high-norm tokens themselves.This indicates that registers benefit DiTs despite the absence of the patch-token pathology observed in ViTs.
  • Findings: Pixel-space DiTs benefit most from registers, whereas latent-space models show moderate gains or slight performance degradation.The effect of registers therefore depends on the model’s training space.
  • Mechanisms: In pixel-space DiTs, registers reduce patch-token feature norms and smooth intermediate feature maps, especially at high noise levels.Different registers specialize as norm sinks or carriers of global semantic information.
  • Architectural implications: Recent pixel-space DiTs’ in-context class-conditioning tokens behave like registers, with some encoding global semantics and others acting as norm sinks.This suggests their performance gains may arise largely from register-like behavior rather than additional class information.
  • Register Guidance: Register Guidance uses a no-register model as the negative direction and, combined with CFG, consistently improves recent pixel-space DiTs.The method is motivated by registers improving object structure and coherence while preserving the main content.

2 Register Tokens for Image Diffusion Transformers

Register tokens benefit image diffusion transformers despite DiTs lacking the patch-token outliers seen in ViTs. Their gains are largest in pixel-space models, where they smooth noisy intermediate representations and can be implicitly realized by in-context or text tokens.

  • DiTs versus ViTs: Unlike ViTs, DiTs show nearly uniform patch-token norms and no attention anomalies in low-information regions, indicating they lack persistent patch-token outliers.This behavior holds across larger pDiT variants and latent-space architectures.
  • Register effects: Adding registers causes pDiTs to develop high-norm register tokens despite the absence of patch-token outliers, while consistently improving generation quality across model sizes.The high-norm register behavior also extends across timesteps, model sizes, latent-space models, and few-shot settings.
  • Register effects: Registers reduce feature norms for all image tokens and primarily improve feature smoothness at high noise, t ∈[0, 0.2], producing smoother and more coherent intermediate representations.At higher timesteps, the feature-norm ratios gradually approach 1.
  • Dependence on diffusion space: Register gains are largest in pixel space, smaller in VAE space, and negative in RAE-based models; applying the same pDiT backbone preserves this trend across spaces.Pixel-space models also have the least smooth intermediate features and the highest patch-token norms, motivating stronger regularization.
  • Register design: pDiTs benefit from registers in deeper blocks (4–11) and require more tokens than ViTs, with pDiT-B performing best using 32 registers versus typically 4 for ViTs.Using registers throughout all layers or only early layers provides performance comparable to the register-free model.
  • Implicit registers: In-context and text tokens can behave as implicit registers: some encode diverse global or semantic information, others act as norm sinks, and pure registers account for most gains over the baseline.In-context conditioning further improves performance by helping form better initial representations.

3 Register Guidance

Register Guidance (RG) amplifies the contribution of register tokens by contrasting predictions from the same model with and without registers, improving generation quality while preserving image content. A register-drop training procedure aligns both prediction regimes, removing artifacts and enabling RG to complement CFG across scales and image sizes.

  • Motivation: With registers, unconditional JiT-B/16 generations show more coherent object structures and cleaner visual details while preserving overall image content.The comparison removes class information to isolate the effect of registers beyond additional conditioning.
  • Inference: Register Guidance uses the register-free model prediction as a weaker counterpart to guide the register-based prediction, with w_rg controlling guidance strength.This mirrors CFG but replaces the unconditional prediction with the prediction from the model evaluated without registers.
  • Inference: Combining RG with CFG further improves generation quality because registers provide information beyond the class label.RG remains complementary to CFG rather than replacing it.
  • Training: A single model trained with occasional register dropping removes the artifacts seen with two separately trained models, while generation quality continues to improve.Separate-model RG improved quality but introduced artifacts at larger guidance scales; the single-model design improves alignment between predictions.
  • Results: RG improves FID as w_rg increases, reaches its best result at w_rg=3.5, and remains complementary to CFG with the best setting at w_rg=1.8 and w_cfg=2.0.Across model scales and image sizes, RG improves over CFG, while RG+CFG gives the best FID in all settings.

4 Discussion · A Related Work · B Implementation Details

The discussion identifies open questions around exploiting registers in DiTs, including dual-stream architectures and extending Register Guidance to text-to-image models. Related work situates the study among attention sinks, representation outliers, and register tokens across generative and vision transformers, while implementation details specify CFG and RG configurations.

  • 4 Discussion: Register-based DiTs remain a promising direction, with substantial scope for further analytical and practical investigation.The authors frame their findings as a step toward better understanding and exploitation of registers in DiTs.
  • 4 Discussion: Dual-stream DiT architectures could process register and patch tokens more effectively, although preliminary efficient designs showed no major improvements.The authors nevertheless consider this direction worthy of further exploration.
  • 4 Discussion: Extending Register Guidance beyond pixel-space class-conditional models to text-to-image models is identified as an important future direction.The current exploration focuses on pixel-space models in the class-conditional setting.
  • A Related Work: Prior work links attention sinks to initial or fixed tokens in autoregressive LLMs, shifting positions in diffusion language models, and first-frame concentration in video DiTs.These studies variously preserve, add, or analyze sink tokens across generation settings.
  • A Related Work: Research on feature-dimension outliers in DiTs connects massive activations to layer normalization and reports improvements in fine-grained image details when they are manipulated.This line of work concerns channel outliers rather than token-level outliers.
  • A Related Work: Register Guidance differs from concurrent channel-outlier guidance by using a register-free weak model, producing improvements in fine details, global structure, and visual coherence.The concurrent approach mainly improves fine-grained details, whereas the register-based weak counterpart is reported to induce substantially larger changes.
  • A Related Work: Studies of ViT registers, text-to-image DiT tokens, and auxiliary padding tokens motivate this work, which extends unusually high-norm token research to image DiTs and emphasizes pixel-space models.The cited literature examines sink artifacts, feature aggregation, post-hoc register insertion, text-token roles, and padding-token effects.
  • B Implementation Details: Table 8 reports training and inference configurations for CFG, RG, and RG+CFG, including guidance parameters and register-drop probability p for RG-based methods.The table separately reports the relevant guidance settings and register-drop probability.

B.1 Analysis Implementation Details

The implementation follows established DiT and JiT configurations while introducing trainable register tokens without additional layers. Experiments cover multiple pixel-space and latent-space model sizes, resolutions, conditioning settings, and training regimes.

  • pDiT implementation: pDiTs follow JiT’s flow-matching setup with x-prediction and trainable register tokens replacing in-context conditioning.The forward process is x_t = tx + (1 − t)ϵ, with t = 1 denoting clean data; registers are trainable parameters without additional layers.
  • pDiT implementation: pDiTs include B (131M), L (459M), and H (953M) models trained on ImageNet at 256 × 256 and 512 × 512 resolutions.B and L use batch size 1024, while the H model at 512 × 512 uses batch size 512 because of computational limits.
  • PixelDiT implementation: PixelDiT experiments use the XL model (800M) across 120-epoch runs with registers, without registers, and in-context conditioning.All configurations are matched except gradient clipping, reduced from 1.0 to 0.5 for stability.
  • Latent-space implementation: Latent-space models follow their original setups, with RAE-XL trained for 80 epochs using batch size 256 and SiTs trained for 300 epochs using the original batch size.Register tokens use the same implementation as in the pixel-space models.

B.2 Register Guidance Implementation Details · C Additional Analysis Results

The implementation evaluates Register Guidance against CFG and RG+CFG in JiT models trained from scratch, while also testing fine-tuning of pretrained models. Fine-tuning uses 50 epochs, a 2e−5 learning rate, and p = 0.3, and yields the best performance according to the supplied passage.

  • B.2 Register Guidance Implementation Details: JiT models are trained from scratch while registers are periodically dropped with probability p.The implementation compares CFG, RG, and RG+CFG settings.
  • B.2 Register Guidance Implementation Details: The three compared settings are CFG as the baseline, RG, and RG+CFG.All models are trained for 600 epochs.
  • B.2 Register Guidance Implementation Details: All settings use the interval guidance strategy.The optimal hyperparameters are reported in Table 8.
  • B.2 Register Guidance Implementation Details: Pretrained JiT models that cannot generate without registers are also evaluated through fine-tuning.This setup is reported in Tables 6 and 7.
  • B.2 Register Guidance Implementation Details: 50 epochs and a lower learning rate of 2e−5 are used in the fine-tuning setup.Other parameters remain the same as in Table 8.
  • B.2 Register Guidance Implementation Details: p = 0.3 is used for fine-tuning, and this setup gives the best performance while applying Register Guidance to already trained models with little addition.The larger register-drop probability is used while keeping other parameters unchanged.

C.1 Analysis on ImageNet

Across ImageNet analyses, registers create high-norm register-token outliers while reducing patch-token norms, improving feature quality and performance—especially for pixel-space DiTs. Their benefits depend on delayed injection and are strongest where intermediate representations are noisier.

  • Outliers and registers in DiTs: DiTs lack ViT-like patch-token outliers, but introducing registers consistently produces high-norm tokens within registers across timesteps, model sizes, resolutions, and diffusion spaces.The effect also appears in few-step latent- and pixel-space models, including MeanFlows and pMF.
  • Pixel-space versus latent-space: Registers improve ImageNet 512×512 performance, with the largest gains in pixel space, moderate gains in VAE space, and degraded performance in DINOv2 space.Additional backbone results show degradation for RAE-based models and similar performance with or without registers in VAE space.
  • Improving Feature Map Quality: Register-token outliers are formed mainly after the MLP layer, rather than after attention, while registers consistently reduce patch-token feature norms.This reduction is observed in larger pixel-space pDiT variants but not in SSL ViTs such as DINOv2.
  • Improving Feature Map Quality: Register tokens improve internal representations at high noise levels, beginning at block 4, where registers are introduced, according to both TV ratio and correlation-decay analyses.Lower correlation-decay slope indicates stronger spatial organization.
  • Effect of Register Injection Layer and Number of Register Tokens: Unlike ViTs, pDiT-B/16 benefits primarily from delayed register introduction; injecting registers across layers 0–11 performs similarly to a model without registers.Registers introduced from layer 0 produce substantially less informative tokens and lower probing accuracy than those introduced from layer 4.
  • Pixel-space versus latent-space: Pixel-space DiTs have larger feature norms and noisier intermediate representations than latent-space counterparts, potentially explaining their stronger gains from registers.Patch-token norms in pDiTs are substantially higher across all timesteps.

C.2 Analysis on Text-to-Image Models · D Additional Experiments

Text-to-image models exhibit high-norm outliers primarily in auxiliary text tokens, supporting their register-like role, while image-token outliers are transient or weaker. Additional fine-tuning experiments identify register-drop probability p = 0.3 as the effective setting for Register Guidance.

  • C.2 Analysis on Text-to-Image Models: The analyzed text-to-image models are pretrained open-source FLUX, SD3.5 Large, Qwen-Image, and Z-Image, rather than models trained with registers.These models append an auxiliary textual sequence to image tokens, and the sequence does not directly participate in the diffusion loss.
  • C.2 Analysis on Text-to-Image Models: The auxiliary sequence is appended to image tokens and propagates textual information, rather than directly participating in diffusion loss as in the classical DiT cross-attention regime.This architectural distinction motivates examining whether auxiliary text tokens can function like registers.
  • D Additional Experiments: p = 0.3 is selected for JiT-B/16 register-drop fine-tuning because p = 0.1 requires longer convergence and p = 0.7 degrades original performance.p = 0 does not improve RG, so the intermediate value p = 0.3 is used.
  • D Additional Experiments: Register-drop fine-tuning with p = 0.3 enables effective RG and RG+CFG, outperforming the original p = 0.0 fine-tuning setup using CFG across JiT models.The table reports the best FID across fine-tuning epochs for each setting; p = 0.0 does not enable effective RG.
  • C.2 Analysis on Text-to-Image Models: SD3.5, Qwen-Image, and Z-Image show high-norm outliers in text tokens while image-token norms remain comparatively uniform.Auxiliary text tokens may implicitly serve as repositories for high-norm representations, resembling registers in ImageNet-based DiTs.
  • C.2 Analysis on Text-to-Image Models: FLUX and PixArt-α also show image-token outliers, but these disappear in later layers, indicating that patch tokens cannot sustain high norms under the loss constraint.In FLUX, image-token outliers are smaller than text-token outliers; PixArt-α has one intermediate-layer outlier at block 16.
  • C.2 Analysis on Text-to-Image Models: The observed patterns suggest models need dedicated space for outliers because patch tokens are inconvenient slots, reinforcing the importance of register tokens.This interpretation follows the contrast between persistent text-token outliers and transient or weaker image-token outliers.

D.1 Register Guidance

Register Guidance is artifact-free when one model supports generation with and without registers, whereas using two separate models can produce artifacts. The best results combine CFG with RG and use fine-tuning with an appropriately chosen register-drop probability.

  • Artifact behavior: Using two separate models can cause Register Guidance artifacts despite similar generated images, but a single model eliminates this mismatch without introducing artifacts.The paper provides visual comparisons of both approaches in Figures 26 and 27.
  • Guidance configuration: The best configuration combines classifier-free guidance (CFG) and Register Guidance (RG).Visual results illustrating this combination are provided in Figure 28.
  • Fine-tuning setup: Register Guidance performs best in a fine-tuning regime using pretrained JiT checkpoints and a larger register-drop probability than training from scratch.The experiments fine-tune pretrained JiT checkpoints for several epochs and ablate p for JiT-B/16 using FID across fine-tuning epochs.

D.2 Decoupled Processing of Register and Patch Tokens

Because register and patch tokens play distinct roles, the paper replaces fully shared processing with a parameter-efficient dual-stream design. The resulting architecture selectively decouples components, introduces registers only in later layers, and improves generation quality with limited overhead.

  • Dual-stream architecture: Register and patch tokens have distinct roles, motivating specialized dual-stream processing instead of fully shared parameters.Interactions between token types are retained through attention while their processing parameters can differ.
  • Dual-stream architecture: Duplicating all transformer components performs best but substantially increases parameters; adaLN and MLP provide the largest individual gains, whereas Attention and RMSNorm contribute limited gains.The ablation evaluates separately duplicated components for register and patch tokens.
  • Compact dual-stream design: Register tokens are introduced only in later stages, after initial single-stream processing, because early-layer registers are ineffective.The compact architecture transitions from single-stream to dual-stream layers after the early stages.
  • Results: The parameter-efficient architecture consistently improves generation quality across model scales and resolutions with approximately 14% additional parameters and no additional GFLOPs runtime overhead.Naive duplication would increase parameters by approximately 77%.
  • Compact dual-stream design: The final compact design dualizes MLP, adaLN, and RMSNorm while sharing Attention, achieving the best trade-off between parameter count and generation quality.Sharing Attention improves performance and reduces parameters, while sharing MLP or removing dual adaLN or RMSNorm degrades results.

D.3 Other results

Register tokens improve REPA-enhanced models, remain most effective when introduced only in deeper layers with more registers, and support strong qualitative generation results when combined with CFG and RG. The best ablation configuration uses delayed register introduction and 32 registers.

  • Representation alignment: Register tokens consistently improve the performance of REPA-enhanced models, while the compact dual-stream design remains effective with REPA.The evaluation uses the REPA adaptation proposed for JiT-based architectures.
  • Register-token ablations: Delayed introduction in layers 4–11 outperforms both no registers and registers enabled from layers 0–11, with 32 registers performing best.Increasing the number of registers improves performance; comparison with layers 4–9 further informs the preferred configuration, though the supplied passage is truncated.
  • Qualitative generation: Using CFG and RG together for the JiT-H/16 model achieves an FID of 1.80 in qualitative generation results.The result is presented in Figures 29 and 30.

E Limitations

The study’s conclusions are limited by its evaluation scope and incomplete explanation of why registers help. Register Guidance also adds tuning requirements and extra model evaluations that can reduce sampling speed.

  • Evaluation scope: The evaluation centers on ImageNet DiTs and pretrained text-to-image models, so broader training recipes and datasets are needed to strengthen the conclusions.The observed behavior is consistent across architectures, model sizes, resolutions, and few-step models, but broader evaluation remains absent.
  • Mechanistic understanding: The paper links registers to improved intermediate representations and generation quality but does not comprehensively explain why they help.
  • Register Guidance: Register Guidance introduces hyperparameters p and w_rg that may require tuning for new tasks, datasets, or architectures.The cited hyperparameters are the register-drop probability p and guidance scale w_rg.
  • Register Guidance: The best-performing RG+CFG setup requires additional model evaluations, increasing function evaluations and limiting sampling speed versus the original model.
Loading 2605.16147v2…