Efficient Diffusion Transformer with Step-wise Dynamic Attention Mediators

TL;DR

This paper introduces a diffusion transformer with mediator tokens and dynamic attention to reduce computational costs and improve image quality during denoising, especially in early diffusion stages.

Contribution

We propose a novel diffusion transformer framework with mediator tokens and dynamic adjustment mechanisms, enhancing efficiency and image quality over prior models.

Findings

Reduces inference FLOPs significantly.

Achieves state-of-the-art FID score of 2.01.

Improves image quality with lower computational cost.

Abstract

This paper identifies significant redundancy in the query-key interactions within self-attention mechanisms of diffusion transformer models, particularly during the early stages of denoising diffusion steps. In response to this observation, we present a novel diffusion transformer framework incorporating an additional set of mediator tokens to engage with queries and keys separately. By modulating the number of mediator tokens during the denoising generation phases, our model initiates the denoising process with a precise, non-ambiguous stage and gradually transitions to a phase enriched with detail. Concurrently, integrating mediator tokens simplifies the attention module's complexity to a linear scale, enhancing the efficiency of global attention processes. Additionally, we propose a time-step dynamic mediator token adjustment mechanism that further decreases the required computational FLOPs for generation, simultaneously facilitating the generation of high-quality images within the constraints of varied inference budgets. Extensive experiments demonstrate that the proposed method can improve the generated image quality while also reducing the inference cost of diffusion transformers. When integrated with the recent work SiT, our method achieves a state-of-the-art FID score of 2.01. The source code is available at https://github.com/LeapLabTHU/Attention-Mediators.