Papers
High-Resolution Image Synthesis with Latent Diffusion Models
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, Björn Ommer · CVPR 2022, 2022
O problema: difusão no espaço de pixels é cara de treinar e de amostrar
Modelos de difusão (DMs) constroem uma imagem aplicando em sequência vários autoencoders de remoção de ruído e, na época do paper, já davam resultados de ponta em síntese de imagens. O problema apontado pelos autores é o custo: como esses modelos operam diretamente no espaço de pixels, treinar os mais poderosos costuma consumir centenas de dias de GPU (o paper cita 150 a 1000 dias de V100 em um trabalho anterior), e gerar 50 mil amostras levaria cerca de 5 dias em uma única A100 naquele mesmo trabalho. A inferência também é cara porque a mesma rede precisa rodar de forma sequencial por muitos passos (de 25 a 1000, no trabalho citado). O paper explica a raiz do desperdício: por serem modelos baseados em verossimilhança, eles tendem a gastar capacidade modelando detalhes imperceptíveis da imagem, e, mesmo que a função de perda possa ser reponderada para ignorar parte desses detalhes, a rede e os gradientes ainda precisam ser avaliados sobre todos os pixels. Os autores veem duas consequências: o treinamento fica restrito a quem tem enorme poder computacional, e usar um modelo já treinado também é caro em tempo e memória. A meta declarada é reduzir a complexidade computacional de treino e de amostragem sem prejudicar o desempenho, o que tornaria essa classe de modelos mais acessível.
Primeiro estágio: um autoencoder perceptual que comprime sem destruir detalhe
A ideia dos autores é separar a compressão da geração. Eles observam que o aprendizado de um modelo de verossimilhança tem dois momentos: uma compressão perceptual, que remove detalhes de alta frequência mas aprende pouca variação semântica, e uma compressão semântica, em que o modelo generativo aprende a composição conceitual dos dados. O primeiro momento é entregue a um autoencoder treinado à parte. Dado uma imagem x, o codificador E produz um latente z = E(x), com altura e largura reduzidas por um fator de downsampling f, e o decodificador D reconstrói a imagem a partir de z. O autoencoder é treinado combinando uma perda perceptual com um objetivo adversarial baseado em patches, o que mantém as reconstruções realistas localmente e evita o borramento causado por usar só perdas de pixel como L1 ou L2. Para impedir latentes com variância arbitrariamente alta, o paper testa duas regularizações: uma leve penalidade KL em direção a uma normal padrão (KL-reg.) e uma camada de quantização vetorial dentro do decodificador (VQ-reg.). Como o modelo de difusão seguinte aproveita a estrutura bidimensional do latente, a compressão pode ser branda, ao contrário de trabalhos anteriores que usavam uma ordenação 1D do espaço latente e compressões bem mais agressivas. Esse estágio é treinado uma única vez e pode ser reutilizado em vários treinamentos de difusão.
Segundo estágio: difundir no espaço latente
Com E e D prontos, o modelo de difusão passa a trabalhar no latente em vez de nos pixels. O processo é o mesmo de um DM comum, uma sequência de autoencoders de remoção de ruído treinados para prever uma versão sem ruído da entrada, mas agora a entrada é zt, a versão ruidosa de z = E(x). A rede de remoção de ruído é uma UNet condicionada ao passo de tempo, e a perda é o erro quadrático entre o ruído verdadeiro e o ruído previsto pela rede. Como o processo direto de adição de ruído é fixo, zt pode ser obtido de forma eficiente a partir do codificador durante o treino. Na geração, uma amostra do latente é decodificada para o espaço de imagem com uma única passada pelo decodificador D. Segundo os autores, esse espaço é mais adequado a modelos de verossimilhança porque eles podem se concentrar nos bits semânticos importantes e treinar numa dimensão menor, computacionalmente mais barata. Como o latente mantém estrutura espacial, o modelo ainda aproveita os vieses indutivos de imagens da UNet, construída principalmente com convoluções 2D. Em contraste com trabalhos que usavam transformers autorregressivos sobre um espaço discreto muito comprimido, isso evita a necessidade de compressão agressiva. O paper também diz que, como o modelo é convolucional, em tarefas densamente condicionadas como super-resolução e inpainting ele pode gerar imagens grandes e consistentes, de cerca de 1024 por 1024 pixels.
Condicionamento por cross-attention: do texto à imagem
Modelos de difusão podem, em princípio, modelar distribuições condicionais, em que a rede de remoção de ruído recebe também uma entrada y, como texto, mapas semânticos ou outra imagem. Os autores observam que combinar DMs com condicionamentos além de rótulos de classe ou versões borradas da entrada era uma área pouco explorada. A solução proposta é adicionar mecanismos de cross-attention à UNet. Um codificador específico do domínio, chamado τθ, projeta y para uma representação intermediária. Essa representação fornece as chaves e os valores da atenção, enquanto as consultas vêm de uma representação intermediária achatada da própria UNet, e o resultado é injetado em camadas intermediárias da rede. O codificador τθ e a UNet são otimizados em conjunto com a perda condicional. Para texto, o paper usa um tokenizador BERT e implementa τθ como um transformer, treinando um modelo KL-regularizado de 1,45 bilhão de parâmetros no LAION-400M. O desenho é flexível: trocar de modalidade significa trocar o codificador τθ, e o paper também treina, com o mesmo mecanismo, modelos condicionados a classe e a layout. Para alguns casos de condicionamento denso, como mapas semânticos, os autores usam em vez disso a concatenação das entradas com o latente.
Trade-offs de compressão e limitações
O paper analisa como o fator de downsampling f afeta o resultado, comparando LDM-f com f em 1, 2, 4, 8, 16 e 32, onde LDM-1 é difusão em pixels. Todos os modelos desta análise usam uma única A100, o mesmo número de passos e o mesmo número de parâmetros. Os resultados mostram dois extremos ruins. Com f pequeno (LDM-1 e LDM-2), o treino progride devagar, porque grande parte da compressão perceptual fica a cargo do modelo de difusão. Com f muito grande, a fidelidade estagna depois de poucos passos, porque o primeiro estágio comprime demais, perde informação e limita a qualidade atingível. Os valores de LDM-4 a LDM-16 equilibram eficiência e fidelidade; depois de 2 milhões de passos, o paper relata uma diferença de 38 em FID entre LDM-1 e LDM-8. Nos testes com amostrador DDIM, LDM-4 e LDM-8 têm FID muito menor que o LDM-1 em pixels e, ao mesmo tempo, maior vazão de amostras, e o paper conclui que LDM-4 e LDM-8 oferecem as melhores condições para síntese de alta qualidade. Conjuntos complexos como o ImageNet exigem taxas de compressão menores. Nas limitações, os autores reconhecem que a amostragem sequencial ainda é mais lenta que a de GANs e que a capacidade de reconstrução do autoencoder pode virar gargalo quando se exige precisão fina em pixels, e supõem que seus modelos de super-resolução já sejam um pouco limitados nesse aspecto.
Veja as ideias em ação
Pratique na escola de System Design os desenhos que usam este paper.