AViTS: Adaptive Spatiotemporal Token Selection for Efficient Dynamic-Resolution Generation

AI summary

Overview

Research area: Generative computer vision — efficient inference for Diffusion Transformers (DiTs), specifically dynamic-resolution sampling for text-to-image generation and instruction-based image editing.

Technical level: Advanced. The paper assumes familiarity with diffusion/flow-matching sampling, Transformer attention, cross-attention between text and image tokens, FLOPs-based efficiency accounting, and metrics such as ImageReward, CLIP Score and the GEdit scoring suite.

Scope: AViTS is a training-free token-prioritization framework that decides which latent tokens get upsampled early during dynamic-resolution denoising, using a fusion of text-attention (spatial) and cross-step feature variance (temporal) signals.

What This Paper Is About

Dynamic-resolution diffusion sampling saves compute by denoising at low resolution early and then moving to higher resolution — but most such methods upsample every latent token at the transition, which wastes high-resolution computation and can hurt fine-detail consistency. Existing partial-upsampling methods (e.g., RALU, Fresco) decide which tokens to upsample using low-level cues such as edges or inter-channel variance, which are only weakly tied to the text instruction and to how a token is actually evolving across denoising steps. AViTS's goal is to prioritize upsampling using two signals that those heuristics ignore: how strongly each token aligns with the text condition, and how much each token's representation is still changing over diffusion steps.

Key Contributions

1. Spatiotemporal heterogeneity analysis. Through attention and step-variance statistics on FLUX.1-dev and Qwen-Image-Edit, the authors show that both signals follow pronounced long-tailed, highly non-uniform distributions — only a small fraction of tokens strongly correlate with the text condition, and only a small subset keep evolving during denoising.

2. Spatiotemporal importance modeling. AViTS formulates upsampling prioritization as an importance-estimation problem, defining spatial importance from latent–text cross-attention and temporal importance from token-wise feature variation across steps (step variance), then fusing them into a single score for selective upsampling.

3. Importance-aware selective upsampling. The fused score drives a three-stage resolution schedule in which high-importance tokens are refined at high resolution earlier while low-importance tokens are deferred.

4. Broad, composable acceleration. AViTS is reported to reduce inference cost across multiple models and tasks (FLUX.1-dev, FLUX.1-Kontext-dev, Qwen-Image-Edit) and to compose with feature caching, step/model distillation and quantization, reaching up to 14.76× acceleration with distilled models.

Main Findings

• Spatial and temporal cues alone each help, and together they help most. Under the same compute budget (NFE = 30) and the same upsampling ratio, AViTS (mix) reaches 5.45× speedup with ImageReward 0.9959 and CLIP Score 32.361 on FLUX.1-dev, above attention-only (ImageReward 0.9875, CLIP 32.251) and step-variance-only (ImageReward 0.9846, CLIP 32.202).

• Low-level heuristics and untrained allocations are worse at equal budget. Random reaches ImageReward 0.9257 and CLIP 31.488; Evenly 0.9689 / 32.012; edge detection 0.9512 / 32.074; Fresco 0.9861 / 31.970 — all below the AViTS mix result.

• Text-to-image results on FLUX.1-dev. At NFE = 30, AViTS reduces latency to 8.21 s (3.14×, 2.65× FLOPs reduction) with ImageReward 1.0104 and CLIP Score 32.476. At NFE = 18, latency is 4.73 s with 5.45× speedup and 4.80× FLOPs reduction, at ImageReward 0.9959 and CLIP Score 32.361. The paper also reports a 9.78× latency acceleration setting on FLUX.1-dev.

• Editing results on FLUX.1-Kontext-dev (GEdit-bench). AViTS at NFE = 30 gives 15.63 s latency (3.21× speed, 3.88× FLOPs) with SC 7.08, PQ 7.12, OS 6.71; at NFE = 18, 11.51 s (4.36×, 5.37× FLOPs) with OS 6.70; at NFE = 11, 7.25 s (6.92×, 8.85× FLOPs) with OS 6.57.

• Editing results on Qwen-Image-Edit (GEdit-bench). AViTS at NFE = 30 gives 87.54 s (3.25× speed, 3.72× FLOPs), with GEdit-CN OS 7.57 and GEdit-EN OS 7.62 (EN SC 7.95). At NFE = 18 it reaches 4.49× speed and 5.24× FLOPs reduction with CN OS 7.58 and EN OS 7.55, and at NFE = 11 it reaches 6.95× speed and 8.65× FLOPs reduction.

• Composability with other acceleration axes. AViTS + feature caching reaches 8.54× latency speedup (9.02× FLOPs) at NFE = 18 with CLIP-IQA 0.8217 and CLIP Score 32.47; AViTS + quantization on FLUX.1-dev-int8 reaches 9.00× speedup at NFE = 18 with CLIP-IQA 0.8723; AViTS + step distillation at NFE = 6 reaches 2.55 s, 14.65× latency speedup and 14.76× FLOPs reduction with CLIP-IQA 0.8709.

• Robust hyperparameters. Appendix ablations report that AViTS performs well across a broad range of the fusion weight α on all three models, that intermediate α values are best or near-best, and that the upsampling ratio ρ shows the expected monotonic quality-versus-speed trade-off, with mid-range values chosen as Pareto-optimal.

• Baseline gap on editing. The paper notes that heuristic spatial schedulers degrade at comparable speeds — for example, TaylorSeer reaching 3.60× speedup with OS 6.17 on FLUX.1-Kontext-dev, versus AViTS at 4.36× with OS 6.70.

Methodology in Plain English

The method runs inference in three stages:

1. Low-resolution denoising plus signal collection. A sample is initialized at the target resolution, spatially downsampled by a factor of two, and denoised for N₁ steps at the reduced token count. Over the last N_T steps at that reduced resolution, the method records two things: the latent snapshot at each step, and the cross-modal attention between image tokens and text tokens.

2. Selective upsampling. From those records it computes two scores per token. The spatial score averages image-to-text attention over heads and over the collection steps — a token that the text condition strongly "looks at" is considered semantically important. The temporal score is the mean over feature channels of the variance of each token's features across the recorded snapshots — a token whose features are still moving is considered under construction and worth refining early. Both scores are min–max normalized to [0, 1] and combined linearly with a weight α. The top K tokens (K = ⌊ρ·M′⌋) are upsampled via orthogonal upsampling, mixed-resolution noise is re-injected, and N₂ denoising steps are run.

3. Full-resolution refinement. The remaining tokens are expanded to complete the sequence, noise is re-injected with coordinate binding, tokens are reordered by spatial position, and N₃ denoising steps recover fine details.

A practical detail: FLUX-family models use fused attention kernels that do not expose attention weights, so the authors attach a forward hook to each double-stream block, capture Q and K, apply the model's own QK normalization, and reconstruct the image-to-text attention submatrix. MLLM-based models like Qwen-Image-Edit expose the image-to-text attention submatrix directly. A tiny Gaussian perturbation (order 10⁻⁶) is added before sorting to break ties.

Why This Matters

Impact on research. The paper reframes partial upsampling in dynamic-resolution diffusion as an importance-estimation problem rather than a heuristic one, and provides a concrete demonstration that semantic alignment and cross-step dynamics carry complementary, measurable information. Because AViTS is training-free and orthogonal to distillation, quantization and caching, it is a plug-in for existing acceleration stacks rather than a competing axis — a useful pattern for the efficiency literature.

Real-world applications.

• Interactive image editing. The GEdit results target instruction-driven editing (e.g., Qwen-Image-Edit, FLUX.1-Kontext-dev), where users expect fast turnaround on edit requests.
• High-resolution text-to-image generation. Cost in DiTs grows with token count, so latency reductions at high resolution matter for consumer-facing generation tools.
• Resource-constrained deployment. The paper explicitly frames latency and memory as barriers to real-time deployment on constrained hardware, and the quantization combination (FLUX.1-dev-int8) speaks to that setting.
• Creative iteration workflows. Faster sampling at stable quality supports prompt-and-retry loops rather than single-shot generation.

Industry relevance. The work is partially sponsored by the Terminal Intelligent Computing Division of Alibaba Cloud, and evaluations span NVIDIA A800 and NVIDIA H20 GPUs, indicating relevance to production inference hardware and to serving stacks built around Qwen-Image-Edit.

Future Directions

• Learned importance predictors. The current scores are computed from attention statistics and step variance with a fixed linear fusion; whether a learned or adaptive predictor would outperform the hand-designed fusion at the same budget is not resolved.

• Extending beyond the evaluated models and tasks. All experiments cover FLUX.1-dev, FLUX.1-Kontext-dev and Qwen-Image-Edit on DrawBench and GEdit; generalizing the same signals to video generation or other DiT backbones is untested here.

• Understanding the attention-extraction overhead. Capturing Q and K via forward hooks and reconstructing attention adds work; how this scales as backbones grow is not reported.

• Theory for the quality–efficiency frontier. The paper demonstrates a Pareto trade-off in ρ empirically but does not derive conditions under which deferring low-importance tokens is provably safe, or characterize when a deferred token's later refinement can recover lost detail.

Target Audience

Researchers and engineers working on diffusion-model inference efficiency — particularly those focused on dynamic-resolution sampling, token sparsification, and serving Diffusion Transformers under latency or FLOPs budgets. It is also relevant to practitioners combining acceleration techniques (caching, distillation, quantization) who want a training-free component that composes with their existing pipeline. Readers without a background in diffusion sampling and Transformer attention will find the method details dense, though the core intuition (prioritize tokens that the text cares about and that are still changing) is accessible.

Authors’ abstract

Diffusion Transformers (DiTs) achieve high-quality generation but are costly due to iterative sampling. Dynamic-resolution sampling reduces early-stage cost by denoising at low resolution; however, uniformly upsampling all latent tokens at resolution transitions incurs redundant computation and may degrade fine-detail consistency. Existing partial upsampling strategies typically rely on local latent structure cues or single-step statistics, making it difficult to jointly capture token-text semantic relevance and token-wise representation dynamics across diffusion steps. We propose AViTS, an adaptive spatiotemporal token selection framework for dynamic-resolution DiTs. AViTS models spatial importance via latent-text attention and temporal importance via token-level feature variation across diffusion timesteps, and fuses them to enable spatiotemporal importance-aware selective upsampling: it prioritizes resolution refinement for critical tokens while deferring less important ones, thereby reducing redundant high-resolution computation and improving the quality-efficiency trade-off. AViTS achieves up to 6.34x on FLUX and nearly 9x FLOPs reduction on Qwen-Image-Edit and FLUX.1-Kontext-dev, orthogonal to distillation, quantization, and feature caching, and reaching 14.76x with distilled models. Code: https://github.com/QHR69/AViTS