Diffusion Models

Modelos de difusão são modelos de IA generativa que criam dados aprendendo a reverter um processo gradual de adição de ruído. Durante o treinamento, imagens ou outros dados são corrompidos com ruído gaussiano ao longo de muitas etapas pequenas, e uma rede neural aprende a prever o ruído em cada etapa. Para gerar uma amostra, o modelo começa com ruído aleatório puro e o remove passo a passo até surgir uma nova amostra. Jascha Sohl-Dickstein e colegas propuseram a abordagem em 2015 em Aprendizado profundo não supervisionado usando termodinâmica fora do equilíbrio, e Jonathan Ho, Ajay Jain e Pieter Abbeel a tornaram competitiva em 2020 com Modelos probabilísticos de difusão para remoção de ruído (DDPM), que estabeleceram um FID de última geração de 3,17 no CIFAR-10. Em 2021, pesquisadores da OpenAI mostraram que os modelos de difusão superavam as GANs na síntese de imagens da ImageNet, e os modelos de difusão e seus equivalentes de correspondência de fluxo agora impulsionam sistemas de conversão de texto em imagem, como o Stable Diffusion, geradores de vídeo, editores de imagem e modelos científicos.

Como funcionam os modelos de difusão#

O processo direto é fixo, não aprendido. Ele adiciona um pouco de ruído gaussiano a cada uma das T etapas; o DDPM usa T = 1.000, com uma variância de ruído que aumenta linearmente de 0.0001 a 0.02 por etapa. Como a soma de ruídos gaussianos continua sendo gaussiana, uma imagem de treinamento pode avançar diretamente para qualquer etapa t por meio de uma expressão fechada, xt = sqrt(alpha_bar_t) * x0 + sqrt(1 - alpha_bar_t) * noise, em que alpha_bar_t é o produto de 1 - beta ao longo das primeiras t etapas. O treinamento escolhe repetidamente uma imagem, uma etapa aleatória e ruído novo, constrói xt e ensina a rede a prever esse ruído usando o erro quadrático médio, o objetivo simplificado do DDPM. Este código aplica o processo direto a uma foto com o cronograma do DDPM:

import urllib.request

import numpy as np
import torch
from PIL import Image

# DDPM's linear schedule: 1,000 steps whose noise variance rises from 0.0001 to 0.02
betas = torch.linspace(1e-4, 0.02, 1000)
alpha_bar = torch.cumprod(1 - betas, dim=0)

# Load a photo and scale its pixels to [-1, 1]
image = Image.open(urllib.request.urlopen("https://ultralytics.com/images/bus.jpg")).convert("RGB").resize((256, 256))
x0 = torch.tensor(np.array(image), dtype=torch.float32).permute(2, 0, 1) / 127.5 - 1

# Jump straight to step t in closed form: x_t = sqrt(alpha_bar) * x0 + sqrt(1 - alpha_bar) * noise
frames = []
for t in (0, 99, 299, 599, 999):
    xt = alpha_bar[t].sqrt() * x0 + (1 - alpha_bar[t]).sqrt() * torch.randn_like(x0)
    frames.append(((xt.clamp(-1, 1) + 1) * 127.5).byte().permute(1, 2, 0).numpy())
    print(f"step {t + 1}: image weight {alpha_bar[t].sqrt():.3f}, noise weight {(1 - alpha_bar[t]).sqrt():.3f}")
Image.fromarray(np.hstack(frames)).save("forward_diffusion.jpg")

O peso da imagem cai de 1.000 na etapa 1 para 0.947 na etapa 100, 0.630 na etapa 300, 0.161 na etapa 600 e 0.006 na etapa 1.000, enquanto o peso do ruído aumenta para 1.000. A imagem salva forward_diffusion.jpg mostra o ônibus nítido, depois granulado, depois tênue e, por fim, desaparecido.

A geração executa o processo ao contrário. Ela começa com ruído gaussiano puro e, a cada etapa, a rede prevê o ruído na amostra atual, um solucionador remove parte dele e o resultado um pouco mais limpo alimenta a etapa seguinte. Um prompt de texto ou rótulo de classe é fornecido à rede como condição, muitas vezes por meio de atenção cruzada. A orientação sem classificador, introduzida por Jonathan Ho e Tim Salimans, treina uma rede tanto com a condição quanto sem ela e, no momento da amostragem, extrapola a partir da previsão incondicional em direção à previsão condicional; uma escala de orientação maior, como o valor padrão 7,5 nos pipelines do Stable Diffusion 1.x, segue o prompt mais de perto, mas reduz a diversidade.

Aceleração da amostragem#

A amostragem DDPM exige 1.000 avaliações da rede por imagem, então boa parte do progresso do campo se concentrou em reduzir esse número:

• Amostradores melhores: o DDIM reutiliza um DDPM treinado com um amostrador determinístico e produz amostras de alta qualidade de 10 a 50 vezes mais rápido em tempo real. Solucionadores ODE posteriores reduzem a geração típica para cerca de 20 a 50 etapas.
• Difusão latente: um modelo de difusão latente remove o ruído de uma representação comprimida por um autoencoder, em vez de fazê-lo diretamente nos pixels. O autoencoder do Stable Diffusion transforma uma imagem de 512×512×3 em uma representação latente de 64×64×4, com 48 vezes menos valores.
• Destilação: os modelos de consistência aprendem a mapear diretamente qualquer amostra ruidosa para uma amostra limpa, possibilitando a geração em uma etapa; os modelos de consistência latente aplicam a mesma ideia ao Stable Diffusion em 2 a 4 etapas.
• Caminhos mais retos: a correspondência de fluxo e o fluxo retificado treinam ao longo de caminhos retos do ruído aos dados, que os solucionadores conseguem percorrer em menos etapas.

A própria rede de remoção de ruído passou da U-Net convolucional do DDPM e do Stable Diffusion 1.x ao transformador de difusão (DiT) usado pela maioria dos geradores atuais de imagens e vídeos.

Modelos de difusão versus outros modelos generativos#

• Redes adversariais generativas (GANs): uma GAN gera uma amostra em uma única passagem direta, mas seu treinamento adversarial entre gerador e discriminador é propenso à instabilidade e ao colapso de modos, em que o gerador cobre apenas parte dos dados. O treinamento por difusão é uma regressão estável e cobre a distribuição dos dados de forma mais completa, ao custo de muitas etapas de amostragem.
• Autoencoders variacionais (VAEs): um VAE decodifica um vetor latente em uma amostra em uma única passagem, e suas amostras tendem a ser mais borradas. Modelos de difusão latente usam um autoencoder no estilo VAE apenas para comprimir imagens e deixam a geração a cargo do modelo de difusão.
• Modelos autorregressivos: modelos de linguagem e modelos visuais autorregressivos geram um token ou uma escala de cada vez, condicionados ao que veio antes. Um modelo de difusão refina todas as partes da amostra em paralelo a cada etapa, e os modelos de linguagem por difusão aplicam o mesmo processo de remoção de ruído ao texto.

Aplicações no mundo real#

• Geração e edição de imagens: modelos de texto para imagem criam imagens a partir de prompts, o preenchimento de imagens regenera uma região mascarada para combinar com o entorno, o ControlNet restringe o layout com mapas de bordas, profundidade ou pose, e os ampliadores por difusão realizam super-resolução.
• Geração de vídeo: modelos de geração de vídeo removem o ruído de representações latentes comprimidas que abrangem muitos quadros de uma só vez, de modo que cada quadro é gerado em conjunto com os vizinhos.
• Dados sintéticos de treinamento: imagens geradas preenchem lacunas em conjuntos de dados de visão, como defeitos raros, condições meteorológicas incomuns ou pontos de vista pouco comuns. Em um conjunto de dados da Ultralytics Platform, a ação Gerar imagens semelhantes cria de 1 a 16 imagens novas por lote, inspiradas em uma imagem existente, e as imagens que você mantém são adicionadas ao conjunto de dados sem rótulos, prontas para anotação.
• Robótica: as políticas de difusão geram sequências de ações de robôs removendo o ruído delas, com condicionamento baseado em imagens de câmeras.
• Ciência: o AlphaFold 3 prevê as estruturas de complexos de proteínas, ácidos nucleicos e moléculas pequenas com um módulo de difusão que remove o ruído das coordenadas atômicas brutas, substituindo o módulo de estrutura do AlphaFold 2.