Resumo Técnico: Sparse Context – Geração Condicionada por Referência Eficiente via Queda de Tokens

Declaração do Problema

Modelos de difusão baseados em referência tornaram-se fundamentais para a geração de imagens controlável, permitindo tarefas como edição de imagem e síntese dirigida por sujeito, condicionando a geração a imagens de entrada. No entanto, esses modelos enfrentam gargalos computacionais significativos. Em modelos de difusão baseados em Transformers modernos (DiTs), as imagens de referência são tipicamente codificadas em grades espaciais densas de tokens via VAE e concatenadas com os tokens ruidosos de entrada. Como as operações de autoatenção escalam quadraticamente com o número de tokens, a inclusão de uma ou mais imagens de referência de alta resolução aumenta drasticamente o tempo de inferência e o uso de memória. Esse custo torna-se proibitivo em aplicações do mundo real que envolvem múltiplas referências. Embora existam técnicas de eficiência para geração dirigida por prompt, elas permanecem subexploradas para modelos condicionais de referência, que apresentam desafios únicos devido à representação densa de entradas visuais.

Metodologia

O artigo introduz o Sparse Context, um framework projetado para construir representações de referência esparsas retendo apenas um subconjunto reduzido de tokens de referência. A metodologia aborda o problema através de três componentes principais:

1. Observação de Redundância

Os autores demonstram primeiro que os tokens de referência contêm redundância significativa. Através de um experimento de "Poda de Contexto Ingênua" (Naïve Context Pruning) em um modelo FLUX.2-Klien pré-treinado, eles mostram que descartar aleatoriamente até 80% dos tokens de referência durante a inferência preserva amplamente o layout grosseiro da cena, apesar de uma perda nos detalhes finos. Isso sugere que os modelos condicionados por referência são inerentemente robustos a entradas esparsas, mas a lacuna de desempenho entre referências completas e parciais é exacerbada por um descompasso entre a distribuição de treino e teste.

2. Estratégia de Treinamento: Queda de Tokens Aleatória

Para mitigar o descompasso de distribuição e permitir uma geração esparsa robusta, os autores realizam o ajuste fino (fine-tuning) de um gerador condicionado por referência de última geração usando queda de tokens aleatória durante o treinamento.

• Processo: Durante o treinamento, uma fração aleatória de tokens de referência é descartada (mantendo uma fração $f \in [0.05, 0.25]$) antes da concatenação com os tokens alvo ruidosos.
• Objetivo: Isso incentiva o modelo a aprender representações robustas que funcionem efetivamente com informações de referência parciais, desacoplando o modelo de qualquer regra específica de seleção de tokens.
• Implementação: Um LoRA (Low-Rank Adaptation) leve é aplicado ao modelo base FLUX-2-Klein-9B em um conjunto de dados curado que abrange edição dirigida por instrução e personalização.

3. Estratégia de Inferência: Seleção de Tokens Consciente da Tarefa

Enquanto o treinamento utiliza a queda aleatória para garantir robustez, a inferência emprega estratégias de seleção de tokens conscientes da tarefa para maximizar a qualidade dentro de um orçamento de computação fixo:

• Edição de Imagem Dirigida por Instrução: Para preservar a estrutura espacial, o método utiliza uma estratégia de amostragem sensível a bordas. Um mapa de bordas Canny é extraído da imagem de referência, e os tokens são amostrados com base nesta distribuição de probabilidade, priorizando limites estruturais sobre regiões uniformes.
• Personalização: Para preservar a identidade do sujeito, o método utiliza amostragem baseada em saliência. Um mapa de saliência aproxima a máscara do sujeito, permitindo que o modelo concentre os tokens nas regiões mais informativas da imagem de referência.

Principais Contribuições

1. Framework de Contexto Esparso: Uma abordagem inovadora que constrói representações de referência esparsas ao descartar tokens, reduzindo significativamente o custo quadrático da autoatenção sem modificar a arquitetura do modelo subjacente.
2. Estratégia de Ajuste Fino Robusta: Um regime de treinamento usando queda de tokens aleatória que permite ao modelo generalizar através de diferentes níveis de esparsidade, permitindo o controle flexível do orçamento de computação em tempo de inferência.
3. Inferência Consciente da Tarefa: A proposta de heurísticas específicas de seleção de tokens (baseadas em bordas para edição, baseadas em saliência para personalização) que superam a queda aleatória ao priorizar informações críticas para a tarefa.
4. Ortogonalidade a Métodos Existentes: Demonstração de que o Sparse Context é compatível e complementar a outras técnicas de eficiência, como Token Merging (ToMe) e KV-caching, permitindo ganhos de aceleração compostos.

Resultados Experimentais

O método foi avaliado em edição de imagem dirigida por instrução (PIE-Bench) e personalização baseada em referência (Subject-200K, CustomDiffusion-105) usando o modelo FLUX-2-Klein-9B.

• Eficiência:
  • Referência Única: Alcança até 2× de aceleração na inferência (ex: 1.83× com retenção de 10% de tokens).
  • Múltiplas Referências: Alcança até 4× de aceleração na inferência para cenários com múltiplas referências (ex: 6 referências), uma vez que os tokens de referência dominam o cálculo de atenção.
  • Eficiência Combinada: Quando integrado com KV-caching para geração de múltiplas referências, o método alcança acelerações de até 6.72×.
• Qualidade:
  • Edição de Imagem: O método preserva a estrutura da cena e a identidade significativamente melhor do que a linha de base de Poda de Contexto Ingênua. Métricas quantitativas (LPIPS, DINO, CLIP-I) mostram que o Sparse Context mantém alta fidelidade mesmo em frações baixas de tokens (ex: $f=0.1$), enquanto a linha de base sofre com corrupção de identidade.
  • Personalização: A amostragem baseada em saliência preserva com sucesso detalhes finos do sujeito (ex: cores específicas, padrões) que são perdidos com a queda aleatória.
• Integração: A combinação do Sparse Context com Token Merging melhora ainda mais o desempenho, demonstrando que o modelo ajustado aproveita efetivamente outras técnicas de redução de tokens.

Significância e Alegações

O artigo afirma que o Sparse Context desbloqueia o potencial de eficiência inexplorado inerente aos modelos de difusão baseados em referência. Ao reconhecer que grades de tokens densas são frequentemente desnecessárias, os autores demonstram que a geração baseada em referência pode ser tornada significativamente mais eficiente sem comprometer a controlabilidade ou a fidelidade visual.

Os autores enfatizam que sua abordagem é simples e modular, exigindo apenas um ajuste fino leve e oferecendo controle flexível em tempo de inferência. Eles afirmam que este método aborda um gargalo crítico na implantação de modelos generativos de alta resolução e múltiplas referências, tornando-os mais práticos para aplicações do mundo real. O trabalho destaca que a redundância nas representações de referência é uma propriedade fundamental que pode ser explorada para reduzir custos computacionais, um achado que se estende além dos modelos específicos testados para a classe mais ampla de sistemas de geração de imagem e vídeo baseados em DiT.