One Model, Many Budgets: Elastic Latent Interfaces for Diffusion Transformers

1. O Problema

Os Transformadores de Difusão (DiTs) alcançaram qualidade de geração de imagem e vídeo superior, mas apresentam limitações fundamentais em termos de eficiência computacional e flexibilidade:

• Rigidez de Custo Computacional: O custo de inferência (FLOPs) em DiTs é rigidamente acoplado à resolução da imagem. Reduzir a resolução para economizar recursos degrada a qualidade, e aumentar a resolução aumenta o custo de forma não linear, sem permitir trade-offs graduais.
• Alocação Uniforme Ineficiente: Os DiTs alocam computação uniformemente em todos os tokens espaciais (patches da imagem). Isso é subótimo, pois regiões visuais variam em complexidade; algumas são fáceis de gerar, enquanto outras exigem mais processamento. O modelo gasta recursos em áreas "fáceis" e não consegue redirecionar esse poder computacional para áreas "difíceis" ou importantes.
• Limitações de Métodos Existentes: Técnicas anteriores, como token dropping (descarte de tokens) ou geradores adaptativos, muitas vezes exigem treinamento complexo, perdem informações irreversíveis na inferência ou não oferecem um mecanismo nativo para ajustar o orçamento de inferência sem re-treinamento.

2. Metodologia: ELIT (Elastic Latent Interface Transformer)

O artigo propõe o ELIT, um mecanismo "plug-and-play" compatível com DiTs que desacopla o tamanho da imagem de entrada do custo computacional. A arquitetura introduz uma Interface Latente elástica composta por duas camadas leves de atenção cruzada (Read e Write) e um conjunto de tokens latentes variáveis.

Componentes Principais:

1. Interface Latente Variável: Em vez de operar diretamente em todos os tokens espaciais, o modelo cria um domínio latente com um número variável de tokens ($K$).
2. Camada Read (Leitura): Uma camada de atenção cruzada leve que "puxa" informações dos tokens espaciais (entrada) para a interface latente. Ela prioriza regiões difíceis ou importantes da imagem, filtrando áreas menos relevantes.
3. Núcleo Latente (Latent Core): A maior parte dos blocos do transformador opera sobre a sequência latente variável, não sobre a imagem completa. Isso reduz drasticamente o custo de atenção quadrática ($O(N^2)$) para $O(K^2)$, onde $K \ll N$.
4. Camada Write (Escrita): Uma camada de atenção cruzada simétrica que "escreve" as atualizações do estado latente de volta para os tokens espaciais, permitindo que o modelo refine os detalhes espaciais com base nas informações processadas no núcleo latente.
5. Treinamento com Tail Dropping: Durante o treinamento, os tokens latentes são divididos em grupos. O modelo é treinado com a remoção aleatória de uma fração dos tokens latentes (os "caudas" ou tails). Isso força o modelo a aprender uma ordenação de importância: os primeiros tokens capturam a estrutura global e informações essenciais, enquanto os tokens subsequentes refinam detalhes.

Inferência Flexível:

No momento da inferência, o usuário pode controlar dinamicamente o número de tokens latentes ($\tilde{J}$) mantidos. Reduzir $\tilde{J}$ diminui o custo computacional (FLOPs) e o tempo de latência, permitindo um trade-off suave entre qualidade e velocidade com um único conjunto de pesos.

3. Contribuições Chave

• Alocação Computacional Adaptativa: O ELIT redistribui a computação de forma não uniforme, concentrando recursos onde a informação é mais densa ou difícil, em vez de tratá-la uniformemente.
• Orçamento de Inferência Variável: Um único modelo suporta um espectro contínuo de pontos de qualidade-latência, ajustando-se a diferentes restrições de hardware sem necessidade de múltiplos modelos ou re-treinamento.
• Autoguidance Nativa (CCFG): A arquitetura permite o uso de Autoguidance (usando uma versão mais fraca do próprio modelo como guia) de forma nativa. Ao usar um orçamento menor ($\tilde{J}_w$) para a parte condicional e um maior para a principal, o modelo consegue melhorar a qualidade de geração reduzindo o custo de inferência em ~33% (chamado de Cheap Classifier-Free Guidance - CCFG).
• Compatibilidade Universal: O método foi validado em diversas arquiteturas de base (DiT, U-ViT, HDiT) e em modelos de grande escala (MM-DiT como o Qwen-Image), mantendo a função objetivo de Rectified Flow inalterada.

4. Resultados Experimentais

Os experimentos foram conduzidos no ImageNet-1K (256px e 512px) e Kinetics-700 (vídeo).

• Desempenho em Imagens (ImageNet-1K 512px):
  • O ELIT-DiT superou consistentemente os baselines (DiT, U-ViT, HDiT).
  • Ganhos Médios: Redução de 35,3% no FID e 39,6% no FDD (Frechet Distance on DINOv2 features) em comparação aos modelos base.
  • Multi-Budget: O modelo treinado com múltiplos orçamentos (ELIT-MB) obteve melhorias ainda maiores, com redução de FID de até 53% no DiT base em 512px.
• Eficiência Computacional:
  • O ELIT oferece um trade-off qualidade-custo superior à simples redução do número de passos de amostragem.
  • Em modelos grandes (Qwen-Image, 20B parâmetros), o ELIT permitiu reduzir os FLOPs de amostragem em até 63% (velocidade 2,7x maior) mantendo a qualidade estrutural e apenas degradando levemente detalhes de alta frequência.
• Convergência: O ELIT acelerou a convergência do treinamento em 3,3x a 4,0x em comparação ao DiT padrão, devido à melhor eficiência no uso dos dados.
• Vídeo: Resultados consistentes no Kinetics-700, superando o DiT base em FID, FDD e FVD.

5. Significado e Impacto

O trabalho do ELIT representa um avanço significativo na eficiência de modelos generativos de última geração:

1. Democratização do Acesso: Permite que modelos de alta qualidade sejam executados em hardware com recursos limitados, ajustando o "orçamento" de inferência em tempo real.
2. Eficiência de Recursos: Resolve o problema do desperdício de computação em regiões de imagem simples, otimizando o uso de energia e tempo de inferência.
3. Simplicidade Arquitetural: Ao contrário de soluções complexas que exigem novos otimizadores ou mudanças profundas no treinamento, o ELIT é uma adição mínima (duas camadas de atenção) que pode ser aplicada a qualquer pipeline de DiT existente.
4. Futuro da Geração Multimodal: A validação em modelos massivos como o Qwen-Image sugere que essa abordagem é escalável e essencial para a viabilidade econômica de modelos generativos multimodais de grande porte.

Em resumo, o ELIT transforma a rigidez dos Transformadores de Difusão em flexibilidade, permitindo que um único modelo se adapte dinamicamente às necessidades computacionais do usuário sem sacrificar a qualidade de geração.