Source-linked AI summary
HiLo-Token: Input-Adaptive High-Low Frequency Token Compression for Efficient Image Editing
Haoran You, Yotam Nitzan, Lingzhi Zhang, Yifan Gong, Mang-Tik Chiu, Connelly Barnes, Yan Kang, Yuqian Zhou, Eli Shechtman, Sohrab Amirghodsi
TL;DR
Diffusion Transformers make generative image editing costly to serve, creating a latency challenge for large-scale creative workflows. HiLo-Token adaptively compresses tokens by preserving high-frequency and masked regions while representing low-frequency context compactly, achieving up to 3.13× DiT speedup without generation-quality regression.
Problem
Diffusion Transformers create substantial serving-cost and latency challenges for large-scale generative image editing.
Method
HiLo-Token adaptively allocates more tokens to high-frequency, context-rich regions while using masked-region retention and downsampled tokens for broader low-frequency structure.
Results
3.13×, 2.59×, and 1.67× DiT speedups were achieved for small, medium, and large mask-ratio groups, respectively, without regression in generation quality.
Takeaways & Limitations
HiLo-Token provides a practical framework for accelerating DiTs in deployed generative image-editing workflows while preserving editing quality.
Takeaways & Limitations
The work primarily targets latency reduction for cloud serving, while memory reduction is described as a by-product rather than an edge-deployment objective.
Abstract
from arXiv · showhide
Creative image editing tools, such as Photoshop's Remove or Generative Fill buttons, are central to everyday customer use and account for a major share of traffic in Photoshop and Lightroom. However, current generative AI models face significant latency challenges, which become even more pronounced when transitioning from convolution-based U-Nets to Diffusion Transformers (DiTs). In our evaluation on hundreds of representative image editing samples spanning a wide range of mask ratios, the DiT module alone accounts for an average of 73% of the total model latency, even after being distilled from 50 timesteps down to 8 timesteps. To tackle this challenge, we propose $\textbf{HiLo-Token}$, an input-adaptive token compression framework that allocates more token budget to high-frequency, rich-context regions while assigning fewer tokens to low-frequency areas. Specifically, for the editing region specified by the user mask, we retain all tokens within a dilated mask to preserve strong locality and contextual relevance. Outside the editing region, we introduce a simple yet effective high-frequency token selection strategy based on spatial frequency to capture important local details, while using tokens from a 16x downsampled image to represent low-frequency components and preserve the blurry but global structure. Extensive experiments on production-level evaluation data validate the effectiveness of the proposed method, achieving 3.13x, 2.59x, and 1.67x DiT speedups on A100-80GB for image editing tasks across small, medium, and large mask ratio categories with average ratios of 6.38%, 15.92%, and 35.36%, respectively, without any regression in generation quality.
1. Introduction
HiLo-Token addresses the latency and serving-cost challenges of DiT-based image editing with input-adaptive token compression that preserves locality while allocating tokens according to spatial frequency. On production-level evaluations, it accelerates editing without generation-quality regression.
- Motivation: DiT-based image editing creates significant latency and serving challenges as generative models transition from convolutional U-Nets to Diffusion Transformers.The paper notes that DiT-based models require hardware upgrades to achieve acceptable latency.
- Method: HiLo-Token allocates more tokens to high-frequency, context-rich regions and fewer tokens to low-frequency areas.It retains all tokens within a dilated user mask to preserve locality and contextual relevance.
- Method: Frequency-aware selection preserves local details with high-frequency tokens and global structure with tokens from a 16× downsampled image.This design addresses context loss caused by retaining only tokens within the dilated mask.
- Method: Sobel-based edge detection guides token selection through a spatial-frequency map without relying on semantic priors or cross-region relations.The method applies pooling and regionalization to the detected frequency map.
- Results: 3.13×, 2.59×, and 1.67× DiT speedups were achieved on A100-80GB for small, medium, and large mask ratios of 6.38%, 15.92%, and 35.36%, respectively, without generation-quality regression.HiLo-Token also reduced the Amazon AWS p5.48xlarge nodes required to serve Remove by 33%.
2. The Proposed Method
HiLo-Token adapts token allocation for image editing by combining low-frequency global structure with high-frequency local detail and mask-centered context. It is integrated into ME-based specialist models, targeting the DiT latency bottleneck through supervised fine-tuning and few-step distillation.
- Post-Training Integration: HiLo-Token is applied during supervised fine-tuning to adapt the ME generalist into task-specific specialists, followed by timestep distillation for few-step inference.The student can use 8-step inference while preserving high-fidelity editing quality.
- ME Model: The ME architecture combines a VAE encoder, VAE decoder, DiT backbone, refiner, and optional text encoder for conditioning.The VAE maps images and masks into latent space, while the refiner operates in pixel space.
- Profiling and Motivation: Approximately 70% of total runtime is consistently attributed to the DiT across 512^2, 768^2, 1024^2, and 2048^2 resolutions.VAE and refiner increasingly dominate memory consumption at large resolutions, but this work primarily targets latency reduction for cloud serving.
- Mask Distribution: More than 50% of editing requests use masks covering less than 10% of the image, and 90% involve editing no more than 50%.These statistics motivate processing masked regions and relevant surrounding context rather than the full image in most scenarios.
- HiLo-Token Selection: HiLo-Token extracts low-frequency tokens from a 16× downsampled image and selects high-frequency tokens using Sobel-based spatial-frequency magnitudes.The frequency map is pooled 16× for coherent region-level selection, while the user mask is dilated to retain nearby context before concatenating both token types.
3. Experiments
Experiments apply HiLo-Token to the ME model and evaluate quality and latency on curated editing data and representative user cases. The framework preserves comparable editing quality while substantially accelerating DiT inference and remaining compatible with additional efficiency techniques.
- Experimental setup: HiLo-Token is finetuned on approximately 407,630 image–mask pairs spanning multiple editing categories with synthetic and real data.
- Evaluation: Quality evaluation uses a dedicated QE team across diverse editing categories, while latency is measured on A100-80GB GPUs.
- Quality results: The ME generalist achieves an overall ImgEdit score of 4.3, comparable to FLUX Kontext, Nano Banana, and GPT Image, while surpassing competitors on soft effect removal.
- Quality results: Token pruning yields comparable quality in most user-study cases, with tie rates of 48%, 70%, and 81% for removal, generative fill, and generative expand.
- Efficiency results: DiT speedups reach 1.67×, 2.59×, and 3.13× across small, medium, and large mask-ratio groups, while FP8 adds up to 40% latency reduction and five-step distillation adds 37.5%.The speedup analysis covers 92 representative user editing cases; five-step distillation causes minor quality degradation affecting fewer than 5% of images.
4. Related Works
Related work on efficient Diffusion Transformers addresses their token-quadratic computational cost through token, layer, and timestep optimization. HiLo-Token extends this direction with input-adaptive token pruning for image editing, operating on top of timestep distillation without quality regression.
- Efficient DiTs: DiTs remain computationally expensive because their complexity grows quadratically with the number of tokens.This computational burden has motivated substantial research into deployment-efficient DiTs.
- Efficient DiTs: Existing DiT-efficiency methods span three orthogonal dimensions: token, layer, and timestep optimization.These categories organize approaches aimed at improving DiT deployment efficiency.
- Efficient DiTs: HiLo-Token performs input-adaptive token pruning for image editing on top of timestep distillation, targeting efficiency without quality regression.The method is positioned as a principled token-level approach within the broader efficient-DiT literature.
- Dynamic Inference: Dynamic inference methods adapt computation to the input or intermediate states, including early exits and conditional layer skipping.Prior work contrasts these dynamic strategies with static model compression and commonly modulates execution along network depth.
5. Conclusion
HiLo-Token is an input-adaptive token compression framework that allocates more computation to high-frequency, context-rich regions and fewer tokens to low-frequency areas. It enables efficient generative image editing models deployed in Adobe applications without compromising editing quality.
- 5. Conclusion: HiLo-Token allocates more computation to high-frequency, context-rich regions while using fewer tokens for low-frequency areas.The framework performs dynamic token allocation based on input characteristics.
- 5. Conclusion: The framework supports efficient generative image editing models deployed in Adobe applications, including Generative Fill and Remove.These applications are identified as deployment examples.
- 5. Conclusion: HiLo-Token enables dynamic token allocation without compromising editing quality.The passage presents this as a practical property of the framework.