Token Pruning for In-Context Generation in Diffusion Transformers

TL;DR

ToPi tackles the efficiency bottleneck of in-context generation in Diffusion Transformers by pruning reference tokens in a principled, training-free manner. It combines offline Representative Layer Anchoring with an attention-and-value-based token influence metric and a fidelity-driven pruning scheme that adapts over the diffusion trajectory. Empirical results on Flux.1-Kontext and Qwen-Image-Edit show substantial inference speedups (exceeding 30% in some tasks) while preserving structural fidelity and visual coherence; orthogonal to temporal acceleration methods, it achieves additional gains when combined. The work advances practical deployment of in-context diffusion models by enabling context-aware computation with minimal runtime overhead.

Abstract

In-context generation significantly enhances Diffusion Transformers (DiTs) by enabling controllable image-to-image generation through reference examples. However, the resulting input concatenation drastically increases sequence length, creating a substantial computational bottleneck. Existing token reduction techniques, primarily tailored for text-to-image synthesis, fall short in this paradigm as they apply uniform reduction strategies, overlooking the inherent role asymmetry between reference contexts and target latents across spatial, temporal, and functional dimensions. To bridge this gap, we introduce ToPi, a training-free token pruning framework tailored for in-context generation in DiTs. Specifically, ToPi utilizes offline calibration-driven sensitivity analysis to identify pivotal attention layers, serving as a robust proxy for redundancy estimation. Leveraging these layers, we derive a novel influence metric to quantify the contribution of each context token for selective pruning, coupled with a temporal update strategy that adapts to the evolving diffusion trajectory. Empirical evaluations demonstrate that ToPi can achieve over 30\% speedup in inference while maintaining structural fidelity and visual consistency across complex image generation tasks.

Figures (11)

• Figure 1: Layer-wise distribution of the total target-aware attention score ($\sum S_j$), averaged across timesteps.
• Figure 2: Evolution of individual target-aware attention scores ($S_j$) for reference tokens across timesteps, averaged across layers.
• Figure 3: Temporal trajectory of the total attention score ($\sum S_j$) under different task conditions, averaged across layers.
• Figure 4: Overview of ToPi. The upper panel depicts the denoising trajectory, which contains multiple denoising intervals. Each interval comprises a mask updating step and multiple token pruning steps. This process is controlled by three stratified mechanisms detailed in the lower panels: (1) Offline Calibration: One-time identification of representative layers $\mathcal{L}$ of a DiT to filter out insensitive layers. (2) Periodic Scoring: Importance scoring and mask updates every $\Delta T$ steps to adapt to semantic shifts. (3) Step-wise Pruning: Instant token pruning and context realignment applied at each denoising step.
• Figure 5: Qualitative comparison on the AnyEdit benchmark. We compare Flux.1-Kontext (left) against Qwen-Image-Edit (right). Best viewed zoomed in. Additional quantitative analyses are provided in Appendix \ref{['app:qualitative']}.

Theorems & Definitions (1)

• Definition 4.1: Token Influence Score