Inteligência artificial
Compreendendo modelos de difusão: um mergulho profundo na IA generativa
Modelos de difusão surgiram como uma abordagem poderosa em IA generativa, produzindo resultados de última geração em geração de imagem, áudio e vídeo. Neste artigo técnico aprofundado, exploraremos como os modelos de difusão funcionam, suas principais inovações e por que se tornaram tão bem-sucedidos. Abordaremos os fundamentos matemáticos, o processo de treinamento, os algoritmos de amostragem e as aplicações de ponta dessa nova e empolgante tecnologia.
Introdução aos Modelos de Difusão
Modelos de difusão são uma classe de modelos generativos que aprendem a gradualmente reduzir o ruído dos dados revertendo um processo de difusão. A ideia central é começar com ruído puro e refiná-lo iterativamente em uma amostra de alta qualidade da distribuição alvo.
Essa abordagem foi inspirada pela termodinâmica de não equilíbrio – especificamente, o processo de reversão da difusão para recuperar a estrutura. No contexto do aprendizado de máquina, podemos pensar nisso como aprender a reverter a adição gradual de ruído aos dados.
Algumas vantagens principais dos modelos de difusão incluem:
- Qualidade de imagem de última geração, superando GANs em muitos casos
- Treinamento estável sem dinâmicas adversas
- Altamente paralelizável
- Arquitetura flexível – qualquer modelo que mapeie entradas para saídas da mesma dimensionalidade pode ser usado
- Forte embasamento teórico
Vamos nos aprofundar em como os modelos de difusão funcionam.
Fonte: Canção et al.
Equações Diferenciais Estocásticas governam os processos direto e reverso em modelos de difusão. A SDE direta adiciona ruído aos dados, transformando-os gradualmente em uma distribuição de ruído. A SDE reversa, guiada por uma função de pontuação aprendida, remove progressivamente o ruído, levando à geração de imagens realistas a partir de ruído aleatório. Essa abordagem é essencial para atingir desempenho generativo de alta qualidade em espaços de estado contínuos
O Processo de Difusão Direta
O processo de difusão direta começa com um ponto de dados x₀ amostrado da distribuição de dados real e adiciona gradualmente ruído gaussiano ao longo de T intervalos de tempo para produzir versões cada vez mais ruidosas x₁, x₂, …, xT.
Em cada passo de tempo t, adicionamos uma pequena quantidade de ruído de acordo com:
x_t = √(1 - β_t) * x_{t-1} + √(β_t) * ε
Onde:
- β_t é um cronograma de variação que controla quanto ruído é adicionado em cada etapa
- ε é ruído gaussiano aleatório
Esse processo continua até que xT seja quase ruído gaussiano puro.
Matematicamente, podemos descrever isso como uma cadeia de Markov:
q(x_t | x_{t-1}) = N(x_t; √(1 - β_t) * x_{t-1}, β_t * I)
Onde N denota uma distribuição gaussiana.
O cronograma β_t é tipicamente escolhido para ser pequeno para os primeiros passos de tempo e aumentar ao longo do tempo. Escolhas comuns incluem cronogramas lineares, cosseno ou sigmoides.
O Processo de Difusão Reversa
O objetivo de um modelo de difusão é aprender o inverso desse processo: começar com ruído puro xT e progressivamente reduzi-lo para recuperar uma amostra limpa x₀.
Modelamos esse processo reverso como:
p_θ(x_{t-1} | x_t) = N(x_{t-1}; μ_θ(x_t, t), σ_θ^2(x_t, t))
Onde μ_θ e σ_θ^2 são funções aprendidas (normalmente redes neurais) parametrizadas por θ.
A principal inovação é que não precisamos modelar explicitamente a distribuição reversa completa. Em vez disso, podemos parametrizá-la em termos do processo direto, que conhecemos.
Especificamente, podemos mostrar que a média ótima do processo reverso μ* é:
μ* = 1/√(1 - β_t) * (x_t - β_t/√(1 - α_t) * ε_θ(x_t, t))
Onde:
- α_t = 1 – β_t
- ε_θ é uma rede de previsão de ruído aprendida
Isso nos dá um objetivo simples: treinar uma rede neural ε_θ para prever o ruído que foi adicionado em cada etapa.
Objetivo de treinamento
O objetivo de treinamento para modelos de difusão pode ser derivado de inferência variacional. Após alguma simplificação, chegamos a uma perda L2 simples:
L = E_t,x₀,ε [ ||ε - ε_θ(x_t, t)||² ]
Onde:
- t é amostrado uniformemente de 1 a T
- x₀ é amostrado a partir dos dados de treinamento
- ε é ruído gaussiano amostrado
- x_t é construído adicionando ruído a x₀ de acordo com o processo direto
Em outras palavras, estamos treinando o modelo para prever o ruído que foi adicionado em cada passo de tempo.
Arquitetura Modelo
Fonte: Ronneberger e outros.
A arquitetura U-Net é central para a etapa de denoising no modelo de difusão. Ela apresenta uma estrutura codificador-decodificador com conexões de salto que ajudam a preservar detalhes de granulação fina durante o processo de reconstrução. O codificador reduz progressivamente a amostragem da imagem de entrada enquanto captura recursos de alto nível, e o decodificador aumenta a amostragem dos recursos codificados para reconstruir a imagem. Essa arquitetura é particularmente eficaz em tarefas que exigem localização precisa, como segmentação de imagem.
A rede de previsão de ruído ε_θ pode usar qualquer arquitetura que mapeie entradas para saídas da mesma dimensionalidade. Arquiteturas de estilo U-Net são uma escolha popular, especialmente para tarefas de geração de imagens.
Uma arquitetura típica pode se parecer com:
class DiffusionUNet(nn.Module):
def __init__(self):
super().__init__()
# Downsampling
self.down1 = UNetBlock(3, 64)
self.down2 = UNetBlock(64, 128)
self.down3 = UNetBlock(128, 256)
# Bottleneck
self.bottleneck = UNetBlock(256, 512)
# Upsampling
self.up3 = UNetBlock(512, 256)
self.up2 = UNetBlock(256, 128)
self.up1 = UNetBlock(128, 64)
# Output
self.out = nn.Conv2d(64, 3, 1)
def forward(self, x, t):
# Embed timestep
t_emb = self.time_embedding(t)
# Downsample
d1 = self.down1(x, t_emb)
d2 = self.down2(d1, t_emb)
d3 = self.down3(d2, t_emb)
# Bottleneck
bottleneck = self.bottleneck(d3, t_emb)
# Upsample
u3 = self.up3(torch.cat([bottleneck, d3], dim=1), t_emb)
u2 = self.up2(torch.cat([u3, d2], dim=1), t_emb)
u1 = self.up1(torch.cat([u2, d1], dim=1), t_emb)
# Output
return self.out(u1)
Os principais componentes são:
- Arquitetura estilo U-Net com conexões de salto
- Incorporação de tempo para condicionar o passo de tempo
- Profundidade e largura flexíveis
Algoritmo de Amostragem
Depois de treinar nossa rede de previsão de ruído ε_θ, podemos usá-la para gerar novas amostras. O algoritmo básico de amostragem é:
- Comece com ruído gaussiano puro xT
- Para t = T para 1:
- Prever ruído:
ε_θ(x_t, t) - Calcular média:
μ = 1/√(1-β_t) * (x_t - β_t/√(1-α_t) * ε_θ(x_t, t)) - Amostra:
x_{t-1} ~ N(μ, σ_t^2 * I)
- Prever ruído:
- Retornar x₀
Esse processo elimina gradualmente o ruído da amostra, guiado pela nossa rede de previsão de ruído aprendida.
Na prática, existem várias técnicas de amostragem que podem melhorar a qualidade ou a velocidade:
- Amostragem DDIM: Uma variante determinística que permite menos etapas de amostragem
- Amostragem ancestral: Incorpora a variância aprendida σ_θ^2
- Amostragem truncada: Para cedo para geração mais rápida
Aqui está uma implementação básica do algoritmo de amostragem:
def sample(model, n_samples, device):
# Start with pure noise
x = torch.randn(n_samples, 3, 32, 32).to(device)
for t in reversed(range(1000)):
# Add noise to create x_t
t_batch = torch.full((n_samples,), t, device=device)
noise = torch.randn_like(x)
x_t = add_noise(x, noise, t)
# Predict and remove noise
pred_noise = model(x_t, t_batch)
x = remove_noise(x_t, pred_noise, t)
# Add noise for next step (except at t=0)
if t > 0:
noise = torch.randn_like(x)
x = add_noise(x, noise, t-1)
return x
A matemática por trás dos modelos de difusão
Para realmente entender os modelos de difusão, é crucial se aprofundar na matemática que os sustenta. Vamos explorar alguns conceitos-chave com mais detalhes:
Cadeia de Markov e equações diferenciais estocásticas
O processo de difusão para frente em modelos de difusão pode ser visto como uma cadeia de Markov ou, no limite contínuo, como uma equação diferencial estocástica (SDE). A formulação SDE fornece uma estrutura teórica poderosa para analisar e estender modelos de difusão.
O SDE avançado pode ser escrito como:
dx = f(x,t)dt + g(t)dw
Onde:
- f(x,t) é o termo de deriva
- g(t) é o coeficiente de difusão
- dw é um processo de Wiener (movimento browniano)
Diferentes escolhas de f e g levam a diferentes tipos de processos de difusão. Por exemplo:
- Explosão de Variância (VE)
SDE: dx = √(d/dt σ²(t)) dw - Preservação de Variância (VP)
SDE: dx = -0.5 β(t)xdt + √(β(t)) dw
Entender essas EDS nos permite derivar estratégias de amostragem ideais e estender modelos de difusão para novos domínios.
Correspondência de pontuação e redução de ruído Correspondência de pontuação
A conexão entre modelos de difusão e correspondência de pontuação fornece outra perspectiva valiosa. A função de pontuação é definida como o gradiente da densidade de probabilidade logarítmica:
s(x) = ∇x log p(x)
A correspondência de pontuação de denoising visa estimar essa função de pontuação treinando um modelo para denoise pontos de dados levemente perturbados. Esse objetivo acaba sendo equivalente ao objetivo de treinamento do modelo de difusão no limite contínuo.
Essa conexão nos permite alavancar técnicas de modelagem generativa baseada em pontuação, como a dinâmica de Langevin recozida para amostragem.
Técnicas de treinamento avançado
Amostragem de Importância
O treinamento do modelo de difusão padrão amostra os passos de tempo uniformemente. No entanto, nem todos os passos de tempo são igualmente importantes para o aprendizado. Técnicas de amostragem de importância podem ser usadas para focar o treinamento nos passos de tempo mais informativos.
Uma abordagem é usar uma distribuição não uniforme ao longo de intervalos de tempo, ponderada pela norma L2 esperada da pontuação:
p(t) ∝ E[||s(x_t, t)||²]
Isso pode levar a um treinamento mais rápido e melhor qualidade da amostra.
Destilação Progressiva
Destilação progressiva é uma técnica para criar modelos de amostragem mais rápidos sem sacrificar a qualidade. O processo funciona da seguinte forma:
- Treine um modelo de difusão base com muitos intervalos de tempo (por exemplo, 1000)
- Crie um modelo de aluno com menos intervalos de tempo (por exemplo, 100)
- Treine o aluno para corresponder ao processo de redução de ruído do modelo base
- Repita os passos 2-3, reduzindo progressivamente os intervalos de tempo
Isso permite uma geração de alta qualidade com significativamente menos etapas de redução de ruído.
Inovações Arquitetônicas
Modelos de difusão baseados em transformadores
Embora as arquiteturas U-Net tenham sido populares para modelos de difusão de imagem, trabalhos recentes exploraram o uso de arquiteturas de transformadores. Os transformadores oferecem várias vantagens potenciais:
- Melhor tratamento de dependências de longo alcance
- Mecanismos de condicionamento mais flexíveis
- Escala mais fácil para tamanhos de modelos maiores
Modelos como DiT (Transformadores de Difusão) mostraram resultados promissores, potencialmente oferecendo um caminho para uma geração de qualidade ainda maior.
Modelos de difusão hierárquica
Modelos de difusão hierárquica geram dados em múltiplas escalas, permitindo tanto coerência global quanto detalhes refinados. O processo normalmente envolve:
- Gerando uma saída de baixa resolução
- Upsampling e refinamento progressivos
Essa abordagem pode ser particularmente eficaz para geração de imagens de alta resolução ou geração de conteúdo de formato longo.
Tópicos avançados
Orientação sem classificador
Orientação sem classificador é uma técnica para melhorar a qualidade e a controlabilidade da amostra. A ideia-chave é treinar dois modelos de difusão:
- Um modelo incondicional p(x_t)
- Um modelo condicional p(x_t | y) onde y é alguma informação de condicionamento (por exemplo, prompt de texto)
Durante a amostragem, interpolamos entre estes modelos:
ε_θ = (1 + w) * ε_θ(x_t | y) - w * ε_θ(x_t)
Onde w > 0 é uma escala de orientação que controla o quanto enfatizar o modelo condicional.
Isso permite um condicionamento mais forte sem a necessidade de retreinar o modelo. Isso tem sido crucial para o sucesso de modelos de texto para imagem, como o DALL-E 2 e o Stable Diffusion.
Difusão latente
Fonte: Rombach e outros.
Modelo de difusão latente (LDM) O processo envolve a codificação de dados de entrada em um espaço latente onde o processo de difusão ocorre. O modelo adiciona progressivamente ruído à representação latente da imagem, levando à geração de uma versão ruidosa, que é então desnoisada usando uma arquitetura U-Net. A U-Net, guiada por mecanismos de atenção cruzada, integra informações de várias fontes de condicionamento, como mapas semânticos, texto e representações de imagem, reconstruindo, em última análise, a imagem no espaço de pixels. Este processo é fundamental na geração de imagens de alta qualidade com uma estrutura controlada e atributos desejados.
Isso oferece várias vantagens:
- Treinamento e amostragem mais rápidos
- Melhor manuseio de imagens de alta resolução
- Mais fácil de incorporar condicionamento
O processo funciona da seguinte maneira:
- Treine um autocodificador para compactar imagens em um espaço latente
- Treine um modelo de difusão neste espaço latente
- Para geração, amostra no espaço latente e decodificação em pixels
Essa abordagem tem sido altamente bem-sucedida, impulsionando modelos como o Stable Diffusion.
Modelos de consistência
Modelos de consistência são uma inovação recente que visa melhorar a velocidade e a qualidade dos modelos de difusão. A ideia-chave é treinar um único modelo que possa mapear de qualquer nível de ruído diretamente para a saída final, em vez de exigir denoising iterativo.
Isso é alcançado por meio de uma função de perda cuidadosamente projetada que impõe consistência entre previsões em diferentes níveis de ruído. O resultado é um modelo que pode gerar amostras de alta qualidade em uma única passagem para frente, acelerando drasticamente a inferência.
Dicas práticas para treinar modelos de difusão
Treinar modelos de difusão de alta qualidade pode ser desafiador. Aqui estão algumas dicas práticas para melhorar a estabilidade e os resultados do treinamento:
- Recorte de gradiente: Use recorte de gradiente para evitar gradientes explosivos, especialmente no início do treinamento.
- EMA de pesos do modelo: Mantenha uma média móvel exponencial (MME) dos pesos do modelo para amostragem, o que pode levar a uma geração mais estável e de maior qualidade.
- Aumento de dados:Para modelos de imagem, aumentos simples como inversões horizontais aleatórias podem melhorar a generalização.
- Programação de ruído: Experimente diferentes programações de ruído (linear, cosseno, sigmoide) para descobrir o que funciona melhor para seus dados.
- Treinamento de precisão mista: Use treinamento de precisão mista para reduzir o uso de memória e acelerar o treinamento, especialmente para modelos grandes.
- Geração condicional:Mesmo que seu objetivo final seja a geração incondicional, o treinamento com condicionamento (por exemplo, em classes de imagem) pode melhorar a qualidade geral da amostra.
Avaliando modelos de difusão
Avaliar adequadamente modelos generativos é crucial, mas desafiador. Aqui estão algumas métricas e abordagens comuns:
Distância inicial de Fréchet (FID)
FID é uma métrica amplamente usada para avaliar a qualidade e a diversidade de imagens geradas. Ela compara as estatísticas de amostras geradas com dados reais no espaço de recursos de um classificador pré-treinado (tipicamente InceptionV3).
Pontuações FID mais baixas indicam distribuições de melhor qualidade e mais realistas. No entanto, o FID tem limitações e não deve ser a única métrica utilizada.
Pontuação inicial (IS)
Pontuação Inicial mede tanto a qualidade quanto a diversidade das imagens geradas. Ele usa uma rede Inception pré-treinada para calcular:
IS = exp(E[KL(p(y|x) || p(y))])
Onde p(y|x) é a distribuição de classe condicional para a imagem gerada x.
Um IS mais alto indica melhor qualidade e diversidade, mas tem limitações conhecidas, especialmente para conjuntos de dados muito diferentes do ImageNet.
Log-verossimilhança negativa (NLL)
Para modelos de difusão, podemos calcular a log-verossimilhança negativa de dados retidos. Isso fornece uma medida direta de quão bem o modelo se ajusta à distribuição de dados verdadeira.
No entanto, pode ser computacionalmente caro estimar o NLL com precisão para dados de alta dimensão.
Avaliação Humana
Para muitas aplicações, especialmente as criativas, a avaliação humana continua crucial. Isso pode envolver:
- Comparações lado a lado com outros modelos
- Avaliações no estilo de teste de Turing
- Avaliações específicas de tarefas (por exemplo, legendas de imagens para modelos de texto para imagem)
Embora subjetiva, a avaliação humana pode capturar aspectos de qualidade que as métricas automatizadas não percebem.
Modelos de Difusão em Produção
A implantação de modelos de difusão em ambientes de produção apresenta desafios únicos. Aqui estão algumas considerações e práticas recomendadas:
Otimização para Inferência
- Exportação ONNX: Converta modelos para o formato ONNX para inferência mais rápida em diferentes hardwares.
- Quantização: Use técnicas como a quantização INT8 para reduzir o tamanho do modelo e melhorar a velocidade de inferência.
- Cache: Para modelos condicionais, armazene em cache os resultados intermediários do modelo incondicional para acelerar a orientação sem classificador.
- Processamento em lote: Aproveite o processamento em lote para fazer uso eficiente dos recursos da GPU.
Escala
- Inferência distribuída: Para aplicações de alto rendimento, implemente inferência distribuída em várias GPUs ou máquinas.
- Amostragem adaptável: Ajuste dinamicamente o número de etapas de amostragem com base na compensação desejada entre qualidade e velocidade.
- Geração progressiva: Para saídas grandes (por exemplo, imagens de alta resolução), gere progressivamente de baixa para alta resolução para fornecer resultados iniciais mais rápidos.
Segurança e Filtragem
- A filtragem de conteúdo: Implementar sistemas robustos de filtragem de conteúdo para evitar a geração de conteúdo prejudicial ou inapropriado.
- watermarking: Considere incorporar marcas d'água invisíveis no conteúdo gerado para rastreabilidade.
Aplicações
Os modelos de difusão obtiveram sucesso em uma ampla gama de tarefas generativas:
Geração de Imagem
A geração de imagens é onde os modelos de difusão ganharam destaque pela primeira vez. Alguns exemplos notáveis incluem:
- DALL-E3: Modelo de texto para imagem da OpenAI, combinando um codificador de texto CLIP com um decodificador de imagem de difusão
- Difusão Estável: Um modelo de difusão latente de código aberto para geração de texto para imagem
- Imagem: Modelo de difusão de texto para imagem do Google
Esses modelos podem gerar imagens altamente realistas e criativas a partir de descrições de texto, superando abordagens anteriores baseadas em GAN.
Geração de vídeo
Modelos de difusão também foram aplicados à geração de vídeo:
- Modelos de difusão de vídeo: Gerando vídeo tratando o tempo como uma dimensão adicional no processo de difusão
- Faz um video: Modelo de difusão de texto para vídeo da Meta
- Vídeo de imagem: Modelo de difusão de texto para vídeo do Google
Esses modelos podem gerar videoclipes curtos a partir de descrições de texto, abrindo novas possibilidades para a criação de conteúdo.
Geração 3D
Trabalhos recentes estenderam modelos de difusão para geração 3D:
- DreamFusion: Geração de texto para 3D usando modelos de difusão 2D
- Ponto: Modelo de difusão de nuvem de pontos da OpenAI para geração de objetos 3D
Essas abordagens permitem a criação de ativos 3D a partir de descrições de texto, com aplicações em jogos, VR/AR e design de produtos.
Desafios e Direções Futuras
Embora os modelos de difusão tenham demonstrado um sucesso notável, ainda existem vários desafios e áreas para pesquisas futuras:
Eficiência Computacional
O processo de amostragem iterativa de modelos de difusão pode ser lento, especialmente para saídas de alta resolução. Abordagens como difusão latente e modelos de consistência visam abordar isso, mas melhorias adicionais em eficiência são uma área ativa de pesquisa.
Controlabilidade
Embora técnicas como a orientação sem classificador tenham melhorado a controlabilidade, ainda há trabalho a ser feito para permitir um controle mais preciso sobre as saídas geradas. Isso é especialmente importante para aplicações criativas.
Geração Multimodal
Os modelos de difusão atuais se destacam na geração de modalidade única (por exemplo, imagens ou áudio). Desenvolver modelos de difusão verdadeiramente multimodais que podem gerar perfeitamente entre modalidades é uma direção empolgante para o trabalho futuro.
Compreensão Teórica
Embora os modelos de difusão apresentem resultados empíricos sólidos, ainda há muito a ser compreendido sobre por que eles funcionam tão bem. Desenvolver uma compreensão teórica mais aprofundada pode levar a melhorias adicionais e novas aplicações.
Conclusão
Os modelos de difusão representam um passo à frente na IA generativa, oferecendo resultados de alta qualidade em uma variedade de modalidades. Ao aprender a reverter um processo de adição de ruído, eles fornecem uma abordagem flexível e teoricamente fundamentada para a geração.
De ferramentas criativas a simulações científicas, a capacidade de gerar dados complexos e de alta dimensão tem o potencial de transformar muitos campos. No entanto, é importante abordar essas tecnologias poderosas com cuidado, considerando tanto seu imenso potencial quanto os desafios éticos que elas representam.
