Token Pruning for In-Context Generation in Diffusion Transformers

2026-02-02

AI Summary

This work addresses the computational bottleneck in context-aware generation with Diffusion Transformers (DiT), where concatenating reference images leads to excessively long input sequences. Existing token compression methods overlook the inherent asymmetry between reference and target tokens in spatial, temporal, and functional roles. To resolve this, the authors propose ToPi, a training-free token pruning framework that identifies quality-sensitive attention layers via offline calibration, constructs a context token influence evaluation mechanism, and employs a temporally adaptive strategy to dynamically adjust pruning decisions. As the first non-uniform pruning approach tailored for DiT-based context generation, ToPi explicitly models the role disparity between reference and target tokens, achieving over 30% inference speedup while preserving structural fidelity and visual consistency.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Generation
• Machine Learning: Deep Generative Models & Autoencoders

Application Category

• Search and Retrieval-Augmented AI: Personalized, context-aware and across-device search
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

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.

Problem

• token pruning
• in-context generation
• Diffusion Transformers
• computational bottleneck
• role asymmetry

Innovation

• Token Pruning
• In-Context Generation
• Diffusion Transformers
• Training-Free
• Sensitivity Analysis