Resumo Técnico: SSDD: Decodificador de Difusão de Passo Único para Tokenização de Imagens Eficiente

Declaração do Problema

Modelos generativos de imagem de última geração dependem de tokenizadores para comprimir sinais visuais de alta dimensão em representações latentes compactas. Embora os tokenizadores atuais padrão, como os Autoencoders Variacionais regularizados por KL (KL-VAE), sejam eficazes, eles enfrentam um compromisso fundamental entre distorção (similaridade ao nível de pixel) e desvio de distribuição (realismo perceptual e diversidade). Decodificadores determinísticos minimizam a distorção, mas frequentemente falham em modelar a distribuição complexa de imagens realistas, levando a saídas borradas. Por outro lado, decodificadores generativos podem modelar essa distribuição, mas tipicamente exigem amostragem iterativa (inferência lenta) ou dependem de perdas adversariais (GANs), que são difíceis de escalar e estabilizar.

Decodificadores baseados em difusão existentes oferecem uma abordagem fundamentada para modelar a distribuição condicional de imagens dados os latentes, mas sofrem de duas limitações principais:

1. Inferência Lenta: Eles exigem múltiplos passos de denoising, criando um gargalo para modelos generativos subsequentes.
2. Instabilidade de Treinamento: Alcançar resultados de última geração frequentemente exige perdas GAN, que são difíceis de escalar e podem levar a execuções de treinamento divergentes.

Metodologia

Os autores introduzem o SSDD (Single-Step Diffusion Decoder), uma arquitetura inovadora projetada para superar essas limitações combinando uma arquitetura híbrida escalável, treinamento livre de GAN e uma estratégia de destilação de passo único.

1. Arquitetura Híbrida Escalável

Em vez de depender apenas de U-Nets convolucionais (comuns em difusão no espaço de pixels) ou puramente Transformers (comuns em modelagem latente), o SSDD utiliza uma arquitetura baseada em U-ViT.

• Estrutura: Emprega um backbone U-Net com quatro níveis de blocos ResNet convolucionais, mas substitui o bloco de atenção central por um Transformer operando sobre tokens que representam patches de pixels de $8 \times 8$.
• Condicionamento: O modelo condiciona-se na representação latente $z$ através do upsampling da grade latente e concatenação com a entrada ruidosa. Ele aprimora ainda mais o condicionamento via Adaptive Group Normalization (AdaGN) e Adaptive Layer Normalization (AdaLN) dentro de ambos os blocos ResNet e Transformer.
• Escalabilidade: A arquitetura escala eficientemente de 13,4M (SSDD-S) para 153,8M (SSDD-XL) de parâmetros aumentando as larguras de canal e a profundidade do transformer.

2. Esquema de Treinamento Livre de GAN

O SSDD elimina a necessidade de treinamento adversarial combinando três funções de perda específicas:

• Perda de Flow Matching: O objetivo primário utiliza optimal transport flow matching para prever o campo de velocidade ($v = x - \epsilon$) entre ruído e dados, treinado com uma perda $L_2$.
• Perda Perceptual (LPIPS): Adicionada para guiar a geração em direção a áreas perceptualmente corretas, alinhando características de alto nível.
• Perda REPA: Uma perda de regularização de característica usando recursos do DINOv2-B para estabilizar as representações do transformer e acelerar o treinamento.
• Treinamento em Múltiplos Estágios: O modelo é treinado em estágios: primeiro em cortes de $128 \times 128$ com aumento de redimensionamento aleatório, seguido de ajuste fino nas resoluções alvo ($256 \times 256$).

3. Destilação de Passo Único

Para abordar o gargalo de velocidade de inferência, o SSDD emprega uma estratégia de destilação:

• Configuração Professor-Aluno: Um decodificador professor de múltiplos passos (treinado com 8 passos) gera reconstruções de alta qualidade e diversas. Um decodificador aluno é então ajustado para reproduzir essas saídas em um único passo.
• Preservação de Perda: Ao contrário da destilação padrão que pode depender apenas de regressão $L_2$, a destilação do SSDD preserva o objetivo de treinamento completo. O aluno minimiza as perdas de flow-matching e LPIPS computadas contra as saídas de múltiplos passos do professor.
• Resultado: Isso permite que o modelo herde as propriedades perceptuais e generativas do professor de múltiplos passos enquanto opera com o throughput de um modelo de passo único.

4. Encoders Compartilhados

Os autores demonstram que os decodificadores SSDD podem ser treinados com um encoder compartilhado (especificamente o encoder KL-regularizado padrão de Rombach et al., 2022). Isso permite que o SSDD sirva como um substituto direto para decodificadores VAE existentes sem a necessidade de retreinar o encoder ou criar um novo espaço latente para cada tamanho de decodificador.

Principais Contribuições

1. Primeiro Decodificador de Difusão de Passo Único Livre de GAN: O SSDD é a primeira arquitetura a alcançar tokenização contínua de última geração sem perdas adversariais, garantindo estabilidade de treinamento em todas as escalas de modelo.
2. Arquitetura Híbrida: A integração de vieses indutivos convolucionais com capacidades generativas de transformer em um decodificador de espaço de pixel produz propriedades de escalonamento superiores em comparação com U-Nets puras.
3. Destilação Eficiente: O método de destilação proposto comprime com sucesso o comportamento de difusão de múltiplos passos em um único passo com mínima degradação de qualidade.
4. Compatibilidade de Espaço Latente Compartilhado: A capacidade de usar encoders pré-treinados torna o SSDD imediatamente aplicável a pipelines generativos existentes.

Resultados Experimentais

Avaliado no ImageNet-1k (256$\times$256) e COCO, o SSDD demonstra melhorias significativas sobre o KL-VAE e outros decodificadores de difusão:

• Qualidade de Reconstrução (configuração f8c4):
  • SSDD-S (13,4M params) melhora o rFID (reconstruction FID) de 0,87 (KL-VAE) para 0,46.
  • SSDD-M (48M params, tamanho similar ao KL-VAE) alcança um rFID de 0,39.
  • O SSDD supera todos os baselines em métricas perceptuais (LPIPS, DreamSim) enquanto mantém métricas de baixo nível competitivas (PSNR, SSIM), embora priorize a qualidade perceptual sobre a distorção perfeita de pixels.
• Throughput:
  • O SSDD-S alcança um throughput 1,4$\times$ maior que o KL-VAE, oferecendo melhor qualidade.
  • Comparado a decodificadores de difusão de múltiplos passos (ex: $\epsilon$-VAE), o SSDD oferece até 10$\times$ mais velocidade de decodificação em níveis de rFID similares.
• Geração Downstream:
  • Ao ser usado para treinar Diffusion Transformers (DiT), o SSDD melhora o gFID (generation FID) em comparação ao KL-VAE.
  • Usando uma razão de compressão maior (f16c4), o SSDD permite uma amostragem 3,8$\times$ mais rápida para o modelo DiT enquanto preserva a qualidade de geração.
• Generalização: Modelos SSDD generalizam bem para COCO e resoluções mais altas (até 1024$\times$1024) sem retreinar do zero, e podem decodificar latentes de encoders congelados como DiTo e DC-AE.

Significância

O artigo afirma que o SSDD estabelece um novo padrão para tokenização de imagem ao resolver a tensão entre velocidade, estabilidade e fidelidade generativa. Ao remover a dependência de GANs e amostragem iterativa, o SSDD fornece um substituto direto e escalável para o KL-VAE que permite o treinamento de modelos generativos de maior qualidade e mais rápidos. O trabalho demonstra que a tokenização baseada em difusão pode ser tanto eficiente quanto estável, pavimentando o caminho para futuros modelos generativos de larga escala sem o overhead computacional da inferência de múltiplos passos ou a instabilidade do treinamento adversarial.