HiLo-Token: Compressão de Tokens de Alta e Baixa Frequência Adaptativa à Entrada para Edição Eficiente de Imagens

Resumo

Ferramentas criativas de edição de imagem, como os botões Remover ou Preenchimento Generativo do Photoshop, são centrais para o uso cotidiano dos clientes e representam uma parcela significativa do tráfego no Photoshop e no Lightroom. No entanto, os modelos generativos de IA atuais enfrentam desafios significativos de latência, que se tornam ainda mais pronunciados na transição de U-Nets baseadas em convolução para Transformers de Difusão (DiTs). Em nossa avaliação com centenas de amostras representativas de edição de imagem abrangendo uma ampla gama de proporções de máscara, o módulo DiT sozinho responde por uma média de 73% da latência total do modelo, mesmo após ser destilado de 50 etapas de tempo para 8. Para enfrentar esse desafio, propomos o HiLo-Token, uma estrutura de compressão de tokens adaptável à entrada que aloca mais orçamento de tokens para regiões de alta frequência e contexto rico, enquanto atribui menos tokens a áreas de baixa frequência. Especificamente, para a região de edição especificada pela máscara do usuário, mantemos todos os tokens dentro de uma máscara dilatada para preservar forte localidade e relevância contextual. Fora da região de edição, introduzimos uma estratégia simples, porém eficaz, de seleção de tokens de alta frequência baseada na frequência espacial para capturar detalhes locais importantes, enquanto utilizamos tokens de uma imagem subamostrada em 16x para representar componentes de baixa frequência e preservar a estrutura global desfocada. Experimentos extensivos em dados de avaliação em nível de produção validam a eficácia do método proposto, alcançando acelerações DiT de 3,13x, 2,59x e 1,67x no A100-80GB para tarefas de edição de imagem nas categorias de proporção de máscara pequena, média e grande, com proporções médias de 6,38%, 15,92% e 35,36%, respectivamente, sem qualquer regressão na qualidade da geração.

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.