HiLo-Token : Compression adaptative à l'entrée des jetons haute et basse fréquence pour une édition d'images efficace

Résumé

Les outils de retouche d'image créative, tels que les boutons Supprimer ou Remplissage génératif de Photoshop, sont au cœur de l'utilisation quotidienne des clients et représentent une part majeure du trafic dans Photoshop et Lightroom. Cependant, les modèles d'IA générative actuels sont confrontés à des défis de latence significatifs, qui deviennent encore plus prononcés lors du passage des U-Net basés sur la convolution aux Transformers de Diffusion (DiTs). Dans notre évaluation sur des centaines d'échantillons représentatifs de retouche d'image couvrant une large gamme de ratios de masque, le module DiT à lui seul représente en moyenne 73 % de la latence totale du modèle, même après avoir été distillé de 50 étapes temporelles à 8. Pour relever ce défi, nous proposons HiLo-Token, un cadre de compression de tokens adaptatif à l'entrée qui alloue davantage de budget de tokens aux régions à haute fréquence et à contexte riche, tout en attribuant moins de tokens aux zones à basse fréquence. Plus précisément, pour la région de retouche spécifiée par le masque utilisateur, nous conservons tous les tokens à l'intérieur d'un masque dilaté afin de préserver une forte localité et une pertinence contextuelle. En dehors de la région de retouche, nous introduisons une stratégie simple mais efficace de sélection de tokens à haute fréquence basée sur la fréquence spatiale pour capturer les détails locaux importants, tout en utilisant des tokens provenant d'une image sous-échantillonnée d'un facteur 16 pour représenter les composantes à basse fréquence et préserver la structure globale floue. Des expériences approfondies sur des données d'évaluation de niveau production confirment l'efficacité de la méthode proposée, atteignant des accélérations du DiT de 3,13x, 2,59x et 1,67x sur A100-80GB pour les tâches de retouche d'image avec des catégories de ratio de masque petit, moyen et grand, correspondant à des ratios moyens de 6,38 %, 15,92 % et 35,36 % respectivement, sans aucune régression de la qualité de génération.

English

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 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.