ImageRAGTurbo: Towards One-step Text-to-Image Generation with Retrieval-Augmented Diffusion Models

Resumo Técnico: ImageRAGTurbo

1. O Problema

Os modelos de difusão atuais são líderes na geração de imagens a partir de texto, mas sofrem de duas limitações principais que impedem sua aplicação em tempo real:

• Latência de Amostragem Iterativa: Os modelos padrão exigem 25 a 50 passos de amostragem (denoising) sequenciais para gerar uma imagem de alta qualidade, o que cria um gargalo computacional.
• Compromisso Qualidade vs. Velocidade: Modelos de "poucos passos" (few-step) ou de "um passo" (one-step), como os baseados em distillation (ex: Stable Diffusion Turbo), reduzem drasticamente a latência, mas frequentemente sacrificam a fidelidade ao prompt (a capacidade de gerar exatamente o que foi solicitado) e a qualidade da imagem.
• Custo de Treinamento: Métodos existentes para melhorar modelos de poucos passos (como adversarial distillation) exigem recursos computacionais massivos para treinamento.

O objetivo do trabalho é desenvolver um método que permita a geração de imagens em um único passo (one-step) mantendo alta fidelidade ao prompt e baixa latência, sem exigir um custo de treinamento proibitivo.

2. Metodologia: ImageRAGTurbo

Os autores propõem o ImageRAGTurbo, uma abordagem que integra Geração Aumentada por Recuperação (RAG) a modelos de difusão de poucos passos. A ideia central é que a recuperação de exemplos relevantes (pares texto-imagem) de um banco de dados fornece contexto rico que simplifica a tarefa do modelo de mapear ruído aleatório para a distribuição alvo.

A arquitetura funciona em duas etapas principais:

A. Recuperação e Injeção Direta (Investigação Preliminar):

• Dado um prompt de entrada ($p_{tgt}$), o sistema recupera pares texto-imagem relevantes de um banco de dados.
• Os autores exploraram inicialmente uma injeção direta no espaço H (H-space), que corresponde aos mapas de características da camada mais profunda do UNet (responsáveis por conceitos de alto nível como objetos e atributos).
• Eles fundem as características do espaço H do prompt alvo com as características do conteúdo recuperado usando interpolação esférica. Embora isso tenha melhorado a fidelidade, exigia uma busca exaustiva de hiperparâmetros durante a inferência, o que aumentava a latência.

B. Ajuste Eficiente via Adaptador H (Solução Proposta):
Para resolver o problema da busca de hiperparâmetros, o método final introduz um adaptador treinável no espaço H:

1. Mecanismo de Cross-Attention: Um adaptador leve ($g_\phi$) é inserido no espaço H do denoiser. Ele utiliza um mecanismo de cross-attention para aprender automaticamente as correlações entre as características recuperadas e as características do prompt alvo.
2. Fusão Automática: O adaptador funde o conteúdo recuperado com o prompt alvo, determinando automaticamente a força da contribuição baseada na similaridade semântica (usando embeddings CLIP).
3. Treinamento: O modelo é ajustado (finetuned) usando treinamento adversário latente (Latent Adversarial Training).
   • Um discriminador (com encoder do professor congelado e camadas projetáveis treináveis) diferencia as latências geradas pelo aluno das latências do professor.
   • São utilizadas funções de perda auxiliares: Score Distillation Loss e Latent LPIPS Loss para estabilizar o treinamento e garantir qualidade visual.
   • Apenas o adaptador e o decodificador do aluno são ajustados, tornando o processo eficiente.

3. Principais Contribuições

• Novo Framework RAG para Few-Step: Introdução de um framework de recuperação aumentada especificamente otimizado para modelos de difusão de poucos passos, melhorando significativamente a fidelidade ao prompt sem sacrificar a velocidade.
• Ajuste Eficiente com Adaptador Leve: Desenvolvimento de uma técnica de finetuning que utiliza um adaptador de rede leve no espaço H, reduzindo drasticamente os requisitos computacionais em comparação com métodos tradicionais de treinamento de modelos de poucos passos.
• Desempenho Superior: Demonstração experimental de que o método supera abordagens existentes em qualidade de geração e alinhamento com o prompt, mantendo uma velocidade de inferência comparável.

4. Resultados Experimentais

O método foi avaliado nos benchmarks MS-COCO e TIFA, comparado com modelos como Stable Diffusion (50 passos), Latent Consistency Models (4 passos) e Stable Diffusion Turbo (1 passo).

• Qualidade e Fidelidade (TIFA e CLIP):
  • O ImageRAGTurbo (1 passo) superou o Stable Diffusion Turbo (1 passo, sem RAG) em 2,2% no score TIFA e 1,37% no score CLIP.
  • Em certas categorias (objetos, atividades, materiais), o modelo de 1 passo do ImageRAGTurbo atingiu resultados comparáveis ou até ligeiramente superiores ao modelo Stable Diffusion padrão de 50 passos.
  • Superou significativamente o modelo RDM (outro método RAG) em TIFA (+7,6%) e qualidade estética, enquanto usava apenas 1 passo em vez de 50.
• Eficiência (Latência):
  • O tempo total de inferência foi de 116,7 ms/imagem, praticamente idêntico ao Stable Diffusion Turbo (113,8 ms).
  • O overhead introduzido pelo mecanismo de recuperação e pelo adaptador foi de apenas 2,5% no tempo de latência.
  • O método é aproximadamente 25 vezes mais rápido que o modelo professor de 50 passos, com uma queda mínima de 1,2% no score TIFA.

5. Significância e Conclusão

O ImageRAGTurbo representa um avanço significativo na viabilidade de modelos de geração de imagem em tempo real. Ao demonstrar que a recuperação de contexto (RAG) pode compensar a perda de qualidade inerente à redução drástica de passos de amostragem, o trabalho oferece uma solução prática para aplicações interativas que exigem baixa latência e alta precisão.

A abordagem prova que é possível alinhar modelos de difusão de "um passo" com prompts complexos sem o custo computacional massivo de treinamento de modelos distilled tradicionais. O trabalho abre caminho para futuras pesquisas em recuperação composicional mais fina e a aplicação de RAG em outros domínios de geração multimodal.