LowDiff: Efficient Diffusion Sampling with Low-Resolution Condition

1. O Problema

Os modelos de difusão alcançaram sucesso notável na geração de imagens de alta fidelidade, mas sua aplicação prática é frequentemente limitada pela lentidão no processo de amostragem (inferência). As abordagens anteriores para melhorar a eficiência focaram principalmente em:

• Compressão de modelos (pruning, quantização).
• Redução do número total de passos de remoção de ruído (NFE - Number of Function Evaluations).
• Geração em espaços latentes de baixa dimensão.

No entanto, essas técnicas muitas vezes negligenciam a possibilidade de aproveitar múltiplas resoluções de entrada durante o processo de geração. Gerar imagens de alta resolução diretamente a partir do ruído exige um custo computacional quadrático em relação à resolução, tornando o processo ineficiente, especialmente nos estágios iniciais onde a estrutura da imagem ainda é grosseira e não requer alta definição.

2. Metodologia (LowDiff)

O LowDiff propõe um novo framework de difusão baseado em uma abordagem cascata (em cascata) e unificada, que gera saídas de resolução progressivamente mais alta.

Principais Componentes Técnicos:

• Geração Condicional de Baixa para Alta Resolução:

  • Em vez de gerar diretamente uma imagem de alta resolução (ex: 32x32 ou 64x64), o modelo começa gerando uma imagem de muito baixa resolução (ex: 8x8) incondicionalmente.
  • Imagens de resoluções subsequentes são geradas condicionadas às imagens de baixa resolução previamente geradas (que foram aumentadas de escala/upsampled).
  • Isso permite reduzir drasticamente o número de passos de remoção de ruído necessários para as resoluções mais altas, pois a estrutura global já foi estabelecida na fase de baixa resolução.

• Arquitetura Unificada (Reutilização de Blocos):

  • Diferente de métodos anteriores (como CDM) que treinam modelos separados para cada resolução, o LowDiff utiliza um único modelo unificado (baseado em U-Net ou Transformer).
  • Reutilização de Sub-redes: Os blocos internos de baixa resolução da rede são reutilizados tanto para gerar imagens de baixa resolução quanto como sub-blocos para processar características downsampled na geração de alta resolução.
  • Adaptadores Leves: Para suportar múltiplas resoluções no mesmo modelo, são introduzidas camadas de convolução de entrada/saída específicas para cada resolução e embeddings de resolução (que informam ao modelo qual resolução está sendo processada), minimizando a sobrecarga arquitetural.

• Treinamento e Amostragem:

  • Treinamento End-to-End: O modelo é treinado do zero sem necessidade de modelos auxiliares ou funções de perda complexas.
  • Augmentation por Truncamento: Durante o treinamento, a geração de alta resolução é condicionada não à imagem limpa de baixa resolução, mas a uma versão parcialmente desnoizada (com ruído injetado independentemente) da imagem de baixa resolução. Isso aumenta a robustez do modelo durante a inferência.
  • Pipeline de Inferência: A amostragem ocorre de baixo para cima (bottom-up). Começa-se com ruído puro na menor resolução, gera-se uma imagem parcialmente desnoizada, faz-se upsample, e usa-se isso como condição para a próxima resolução, repetindo-se até a resolução final.

3. Contribuições Chave

1. DM Eficiente Multi-Resolução: Um framework cascata que inicia a geração em baixa resolução, reduzindo significativamente o número de passos necessários para a geração de alta resolução sem comprometer a qualidade final.
2. Modelo Unificado End-to-End: Uma arquitetura que reutiliza os mesmos blocos de rede para todas as resoluções, reduzindo o tamanho do modelo e o uso de memória em comparação com abordagens que usam múltiplos modelos independentes.
3. Melhor Trade-off Qualidade-Eficiência: O método alcança qualidade de geração comparável ou superior aos modelos base, mas com um ganho de throughput (vazão) de 50% a 80%.

4. Resultados Experimentais

Os autores avaliaram o LowDiff em conjuntos de dados condicionais e incondicionais (CIFAR-10, FFHQ, ImageNet), comparando com modelos de referência como EDM e LightningDiT.

• CIFAR-10 (Incondicional):
  • FID: 2.11 (vs. 2.04 do EDM, que usa 35 passos).
  • IS: 9.87.
  • Eficiência: Redução de passos efetivos para 22 (aproximadamente 60% menos esforço computacional).
• CIFAR-10 (Condicionado):
  • FID: 1.94, IS: 10.03.
• FFHQ 64x64:
  • FID: 2.43 (melhor que o baseline de 2.47).
  • Aceleração: Redução de latência de 12.06s para 7.95s (52.8% de speedup).
• ImageNet 256x256:
  • Baseado em LightningDiT-B/1: FID 4.00 (vs. 4.38 do baseline).
  • Baseado em LightningDiT-XL/1: FID 1.59.
  • Aceleração: Speedup de até 83.3% no modelo XL.

Estudos de Ablação:

• Mostrar que adicionar mais estágios de baixa resolução (ex: 3 estágios vs. 2 estágios) melhora a eficiência e a vazão sem perder qualidade.
• Confirmar que o modelo unificado (LowDiff) supera ou iguala a qualidade de modelos cascata separados (CDM) com menos parâmetros e menor uso de memória.

5. Significância e Conclusão

O LowDiff representa um avanço significativo na eficiência de modelos de difusão ao demonstrar que a geração de alta resolução não precisa ser feita inteiramente em alta resolução desde o início. Ao deslocar a carga computacional para resoluções mais baixas (onde o custo é menor) e usar essas saídas como condições robustas para refinamento posterior, o método:

• Elimina a redundância nos estágios iniciais de alto ruído.
• Oferece uma solução escalável que funciona tanto no espaço de pixels quanto no espaço latente.
• Simplifica o pipeline de treinamento (end-to-end, sem distilação complexa).

Este trabalho abre caminho para a aplicação prática de modelos de difusão em cenários com recursos computacionais limitados, mantendo a alta qualidade de geração exigida por aplicações reais. O código e os modelos estão disponíveis publicamente no GitHub.