Taming Outlier Tokens in Diffusion Transformers
- URL: http://arxiv.org/abs/2605.05206v1
- Date: Wed, 06 May 2026 17:59:42 GMT
- Title: Taming Outlier Tokens in Diffusion Transformers
- Abstract summary: We study outlier tokens in Diffusion Transformers (DiTs) for image generation.<n>We show that this phenomenon appears in both the encoder and denoiser of modern Representation Autoencoder (RAE)-DiT pipelines.<n>We introduce Dual-Stage Registers (DSR), a register-based intervention for both components.
- Score: 55.42341508886889
- License: http://creativecommons.org/licenses/by/4.0/
- 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.
Related papers
- Where Does Generative Difficulty Reside? An Empirical Study of Target Representations [62.54876287800644]
The target representation defines the distribution an image generator must learn, yet it is often treated as an interchangeable interface.<n>We study raw pixels, SD-VAE latents and DINOv2 as well as MAE representation-autoencoder features within a unified masked autoregressive rectified-flow model.<n>Our results show that compression, reconstruction fidelity, token dimensionality, and visible semantic clustering do not individually predict generative behavior.
arXiv Detail & Related papers (2026-08-01T12:32:31Z) - Registers Matter for Pixel-Space Diffusion Transformers [22.836387612443037]
Vision Transformers (ViTs) are known to exhibit high-norm patch-token outliers that degrade feature map quality.<n>We show that DiTs differ from ViTs in a key respect: they do not exhibit patch-token outliers.<n>We find that register tokens significantly improve convergence and generation quality of pixel-space DiTs.
arXiv Detail & Related papers (2026-05-15T16:27:10Z) - ViT Registers and Fractal ViT [7.970267053011483]
We invent and test a variant of large vision transformers (ViTs) called fractal ViT.<n>These models do not improve upon ViT with registers, highlighting the fact that these findings may be scale, domain, or application-specific.
arXiv Detail & Related papers (2026-01-21T22:24:53Z) - Vision Transformers Don't Need Trained Registers [19.41083885307305]
A sparse set of neurons is responsible for concentrating high-norm activations on tokens.<n>We create a training-free approach to mitigate these artifacts.<n>Our results suggest that test-time registers effectively take on the role of register tokens at test-time.
arXiv Detail & Related papers (2025-06-09T17:59:57Z) - Vision Transformers with Self-Distilled Registers [11.649023403110528]
Post Hoc Registers (PH-Reg) is an efficient self-distillation method that integrates registers into an existing ViT without requiring additional labeled data and full retraining.<n>We show that our approach can effectively reduce the number of artifact tokens, improving the segmentation and depth prediction of the student ViT under zero-shot and linear probing.
arXiv Detail & Related papers (2025-05-27T17:59:41Z) - Leveraging Registers in Vision Transformers for Robust Adaptation [36.30125926550378]
Vision Transformers (ViTs) have shown success across a variety of tasks due to their ability to capture global image representations.<n>Recent studies have identified the existence of high-norm tokens in ViTs, which can interfere with unsupervised object discovery.<n>To address this, the use of "registers" which are additional tokens that isolate high norm patch tokens while capturing global image-level information has been proposed.
arXiv Detail & Related papers (2025-01-08T19:02:32Z) - Locality-Aware Generalizable Implicit Neural Representation [54.93702310461174]
Generalizable implicit neural representation (INR) enables a single continuous function to represent multiple data instances.
We propose a novel framework for generalizable INR that combines a transformer encoder with a locality-aware INR decoder.
Our framework significantly outperforms previous generalizable INRs and validates the usefulness of the locality-aware latents for downstream tasks.
arXiv Detail & Related papers (2023-10-09T11:26:58Z) - Accurate Image Restoration with Attention Retractable Transformer [50.05204240159985]
We propose Attention Retractable Transformer (ART) for image restoration.
ART presents both dense and sparse attention modules in the network.
We conduct extensive experiments on image super-resolution, denoising, and JPEG compression artifact reduction tasks.
arXiv Detail & Related papers (2022-10-04T07:35:01Z) - Less is More: Pay Less Attention in Vision Transformers [61.05787583247392]
Less attention vIsion Transformer builds upon the fact that convolutions, fully-connected layers, and self-attentions have almost equivalent mathematical expressions for processing image patch sequences.
The proposed LIT achieves promising performance on image recognition tasks, including image classification, object detection and instance segmentation.
arXiv Detail & Related papers (2021-05-29T05:26:07Z) - TFill: Image Completion via a Transformer-Based Architecture [69.62228639870114]
We propose treating image completion as a directionless sequence-to-sequence prediction task.
We employ a restrictive CNN with small and non-overlapping RF for token representation.
In a second phase, to improve appearance consistency between visible and generated regions, a novel attention-aware layer (AAL) is introduced.
arXiv Detail & Related papers (2021-04-02T01:42:01Z) - Addressing Some Limitations of Transformers with Feedback Memory [51.94640029417114]
Transformers have been successfully applied to sequential, auto-regressive tasks despite being feedforward networks.
We propose the Feedback Transformer architecture that exposes all previous representations to all future representations.
We demonstrate on a variety of benchmarks in language modeling, machine translation, and reinforcement learning that the increased representation capacity can create small, shallow models with much stronger performance than comparable Transformers.
arXiv Detail & Related papers (2020-02-21T16:37:57Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.