Registers Matter for Pixel-Space Diffusion Transformers

2026-07-08 · Source: cs.CV updates on arXiv.org · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Computer Vision · Depth: Expert, extended

Summary

A study on Diffusion Transformers (DiTs) reveals that while they do not exhibit the high-norm patch-token outliers common in Vision Transformers (ViTs), register tokens significantly enhance the convergence and generation quality of pixel-space DiTs. Unlike ViTs where registers absorb existing outliers, in DiTs, registers themselves develop high-norm tokens. This mechanism leads to cleaner, smoother intermediate feature maps, particularly at high noise levels, and reduces patch-token feature norms. Register tokens also specialize, acting as both norm sinks and encoders of global semantic information. The research notes that existing pixel-space DiT architectures, such as JiT, implicitly incorporate register-like behaviors through in-context conditioning. Motivated by these findings, a parameter-efficient dual-stream architecture is proposed, which processes register and patch tokens distinctly, improving generation quality on ImageNet 256x256 (e.g., FID from 3.71 to 3.48) with only a ~14% increase in model size and no runtime overhead.

Key takeaway

For Machine Learning Engineers optimizing pixel-space Diffusion Transformers, integrating register tokens is crucial for enhancing generation quality and convergence. Your models will benefit from cleaner internal representations and reduced patch-token feature norms. Consider implementing a parameter-efficient dual-stream architecture, specializing RMSNorm, adaLN, and MLP for register tokens, and introducing registers in deeper layers (e.g., 4-11) to achieve significant performance gains with minimal overhead. This approach can improve FID scores, as demonstrated by the 3.71 to 3.48 improvement.

Key insights

Register tokens, despite DiTs lacking ViT-like outliers, significantly improve pixel-space Diffusion Transformer performance by regularizing internal representations.

Principles

• Register tokens reduce patch-token feature norms.
• Registers improve feature smoothness at high noise levels.
• Register tokens specialize as norm sinks or semantic encoders.

Method

A parameter-efficient dual-stream architecture processes register and patch tokens distinctly, sharing attention but specializing RMSNorm, adaLN, and MLP layers. Registers are introduced in deeper blocks (e.g., layers 4-11) for optimal effect.

In practice

• Introduce register tokens in deeper DiT layers (e.g., 4-11).
• Use more register tokens for DiTs (e.g., 32 for pDiT-B).
• Decouple processing for register and patch tokens.

Topics

• Diffusion Transformers
• Register Tokens
• Pixel-Space Models
• Vision Transformers
• Generative AI
• Neural Architecture
• Image Generation

Code references

• black-forest-labs/flux

Best for: Research Scientist, AI Scientist, Machine Learning Engineer, Computer Vision Engineer

Related on AIssential

• Rethinking Token Reduction for Diffusion Models via Output-Similarity-Awareness — Output-centric token reduction in Diffusion Models, guided by prior-step output similarity, minimizes recovery error and boosts generation quality.
• RiT: Vanilla Diffusion Transformers Suffice in Representation Space — DINOv2's favorable geometric properties enable vanilla Diffusion Transformers with "x"-prediction to achieve superior, efficient image generation.
• DiffusionBench: On Holistic Evaluation of Diffusion Transformers — ImageNet-centric DiT evaluation is insufficient; T2I performance does not correlate, necessitating holistic benchmarks like DiffusionBench.
• Show the Signal, Hide the Noise: Spectral Forcing for Pixel-Space Diffusion — Spectral Forcing improves pixel-space diffusion by explicitly filtering high-frequency noise, optimizing model capacity for signal processing.
• Distilling Drifting Transformers with Representation Autoencoders — Drift-RAE enables stable, high-performance distillation of flow models in semantically rich Representation Autoencoder latent spaces.
• Variable-Length Tokenization via Learnable Global Merging for Diffusion Transformers — Merging similar tokens enables variable-length tokenization while preserving cross-length representation alignment in Diffusion Transformers.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by cs.CV updates on arXiv.org.