Introdução aos modelos de difusão para aprendizado de máquina

Este artigo foi originalmente publicado em Montagem AI e republicado para TOPBOTS com permissão do autor.

Modelos de difusão são modelos generativos que vêm ganhando popularidade significativa nos últimos anos, e por boas razões. Um punhado de artigos seminais lançados na década de 2020 sozinho mostraram ao mundo do que os modelos de difusão são capazes, como vencer GANs^[6] na síntese de imagens. Mais recentemente, os praticantes terão visto Modelos de Difusão usados ​​em DALL-E2, modelo de geração de imagens da OpenAI lançado no mês passado.

Dada a recente onda de sucesso dos Modelos de Difusão, muitos praticantes de Aprendizado de Máquina certamente estão interessados ​​em seu funcionamento interno. Neste artigo, examinaremos a fundamentos teóricos para Modelos de Difusãoe, em seguida, demonstrar como gerar imagens com um Modelo de difusão no PyTorch. Vamos mergulhar!

Se este conteúdo educacional aprofundado for útil para você, inscreva-se em nossa lista de discussão de IA para ser alertado quando lançarmos novo material. 

Modelos de Difusão - Introdução

Os modelos de difusão são generativo modelos, o que significa que eles são usados ​​para gerar dados semelhantes aos dados nos quais são treinados. Fundamentalmente, os Modelos de Difusão funcionam por destruindo dados de treinamento através da adição sucessiva de ruído gaussiano, e então aprendendo a recuperar os dados por revertendo este processo de ruído. Após o treinamento, podemos usar o Modelo de Difusão para gerar dados simplesmente passando o ruído amostrado aleatoriamente através do processo de remoção de ruído aprendido.

Mais especificamente, um modelo de difusão é um modelo de variável latente que mapeia para o espaço latente usando uma cadeia de Markov fixa. Essa cadeia gradualmente adiciona ruído aos dados para obter o q(x posterior aproximado)_1:T|x₀), onde x₁,…, X_T são as variáveis ​​latentes com a mesma dimensionalidade que x₀. Na figura abaixo, vemos tal cadeia de Markov manifestada para dados de imagem.

Em última análise, a imagem é transformada assintoticamente em ruído gaussiano puro. o meta de treinar um modelo de difusão é aprender a reverso processo - ou seja, treinamento p_θ(x_t-1|x_t). Ao percorrer para trás ao longo dessa cadeia, podemos gerar novos dados.

Benefícios dos Modelos de Difusão

Como mencionado acima, a pesquisa sobre Modelos de Difusão explodiu nos últimos anos. Inspirado na termodinâmica de não equilíbrio^[1], Modelos de Difusão produzem atualmente Qualidade de imagem de última geração, cujos exemplos podem ser vistos abaixo:

Além da qualidade de imagem de ponta, os modelos de difusão vêm com uma série de outros benefícios, incluindo não requer treinamento adversário. As dificuldades do treinamento adversário estão bem documentadas; e, nos casos em que existem alternativas não adversas com desempenho e eficiência de treinamento comparáveis, geralmente é melhor utilizá-las. No tópico da eficiência do treinamento, os Modelos de Difusão também têm os benefícios adicionais de escalabilidade e paralelização.

Embora os Modelos de Difusão quase pareçam estar produzindo resultados do nada, há muitas escolhas matemáticas cuidadosas e interessantes e detalhes que fornecem a base para esses resultados, e as melhores práticas ainda estão evoluindo na literatura. Vamos dar uma olhada na teoria matemática que sustenta os Modelos de Difusão com mais detalhes agora.

Modelos de difusão - um mergulho profundo

Como mencionado acima, um Modelo de Difusão consiste em um processo de encaminhamento (ou processo de difusão), em que um dado (geralmente uma imagem) é progressivamente ruidoso, e um processo reverso (ou processo de difusão reversa), em que o ruído é transformado novamente em uma amostra da distribuição alvo.

As transições da cadeia de amostragem no processo direto podem ser definidas como gaussianas condicionais quando o nível de ruído é suficientemente baixo. Combinando este fato com a suposição de Markov leva a uma parametrização simples do processo direto:

Nota Matemática

Temos falado sobre corromper os dados por acrescentando Ruído gaussiano, mas a princípio pode não estar claro onde estamos realizando essa adição. De acordo com a equação acima, em cada etapa da cadeia estamos simplesmente amostrando de uma distribuição gaussiana cuja média é o valor anterior (ou seja, imagem) na cadeia.

Essas duas afirmações são equivalentes. Aquilo é

Para entender o porquê, utilizaremos um pequeno abuso de notação ao afirmar que

Onde a implicação final decorre da equivalência matemática entre uma soma de variáveis ​​aleatórias e a convolução de suas distribuições – veja esta página da Wikipedia para obter mais informações.

Em outras palavras, mostramos que afirmar a distribuição de um passo de tempo condicionada ao anterior através da média de uma distribuição gaussiana equivale a afirmar que a distribuição de um dado passo de tempo é a do anterior com a adição de ruído gaussiano. Omitimos os escalares introduzidos pelo esquema de variância e mostramos isso para uma dimensão por simplicidade, mas uma prova semelhante vale para Gaussianos multivariados.

Onde β₁,…,β_T é um esquema de variação (seja aprendido ou fixo) que, se bem comportado, Assegura que x_T é quase uma gaussiana isotrópica para T suficientemente grande.

Como mencionado anteriormente, a “mágica” dos modelos de difusão vem na processo reverso. Durante o treinamento, o modelo aprende a reverter esse processo de difusão para gerar novos dados. Começando com o ruído gaussiano puro p(x_T):=N(x_T,0,I) o modelo aprende a distribuição conjunta p_θ(x_0:T) Como

onde os parâmetros dependentes do tempo das transições gaussianas são aprendidos. Observe em particular que a formulação de Markov afirma que uma determinada distribuição de transição de difusão reversa depende apenas do passo de tempo anterior (ou do passo de tempo seguinte, dependendo de como você o vê):

Training

Um Modelo de Difusão é treinado por encontrar as transições reversas de Markov que maximizam a probabilidade dos dados de treinamento. Na prática, o treinamento consiste de forma equivalente em minimizar o limite superior variacional da probabilidade logarítmica negativa.

Procuramos reescrever o L_VLB em termos de Divergências de Kullback-Leibler (KL). A Divergência KL é uma medida de distância estatística assimétrica de quanto uma distribuição de probabilidade P difere de uma distribuição de referência Q. Estamos interessados ​​em formular L_VLB em termos de divergências KL porque as distribuições de transição em nossa cadeia de Markov são gaussianas, e a divergência KL entre as Gaussianas tem uma forma fechada.

O que é a Divergência KL?

A forma matemática da divergência KL para distribuições contínuas é

Abaixo você pode ver a divergência KL de uma distribuição variável P (azul) de uma distribuição de referência Q (vermelho). A curva verde indica a função dentro da integral na definição para a divergência KL acima, e a área total sob a curva representa o valor da divergência KL de P da Q a qualquer momento, um valor que também é exibido numericamente.

Formação do elenco  em termos de divergências KL

Como mencionado anteriormente, é possível [1] para reescrever L_VLB quase completamente em termos de divergências KL:

onde

Detalhes da Derivação

O limite variacional é igual a

Substituindo as distribuições por suas definições, dada nossa suposição de Markov, obtemos

Usamos regras de log para transformar a expressão em uma soma de logs e, em seguida, extraímos o primeiro termo

Usando o Teorema de Bayes e nossa suposição de Markov, esta expressão se torna

Em seguida, dividimos o termo médio usando regras de log

Isolando o segundo termo, vemos

Colocando isso de volta em nossa equação para L_VLB, temos

Usando regras de log, reorganizamos

Em seguida, notamos a seguinte equivalência para a divergência KL para quaisquer duas distribuições:

Finalmente, aplicando esta equivalência à expressão anterior, chegamos a

Condicionando o processo para frente posterior em x₀ em L_t-1 resulta em uma forma tratável que leva a todas as divergências KL são comparações entre gaussianas. Isso significa que as divergências podem ser calculadas exatamente com expressões de forma fechada em vez de estimativas de Monte Carlo^[3].

Opções de modelo

Com a base matemática para nossa função objetivo estabelecida, agora precisamos fazer várias escolhas sobre como nosso Modelo de Difusão será implementado. Para o processo de encaminhamento, a única opção necessária é definir o cronograma de variação, cujos valores geralmente aumentam durante o processo de encaminhamento.

Para o processo reverso, escolhemos a(s) arquitetura(s) de parametrização/modelo de distribuição gaussiana(s). Note o alto grau de flexibilidade que os Modelos de Difusão permitem - o só O requisito em nossa arquitetura é que sua entrada e saída tenham a mesma dimensionalidade.

Vamos explorar os detalhes dessas escolhas com mais detalhes abaixo.

Processo de encaminhamento e L_T

Conforme observado acima, em relação ao processo forward, devemos definir o cronograma de variância. Em particular, nós os definimos para serem constantes dependentes do tempo, ignorando o fato de que eles podem ser aprendidos. Por exemplo^[3], um esquema linear de β₁= 10^-4 para β_T=0.2 pode ser usado, ou talvez uma série geométrica.

Independentemente dos valores particulares escolhidos, o fato de o esquema de variância ser fixo resulta em L_T tornando-se uma constante em relação ao nosso conjunto de parâmetros apreensíveis, permitindo-nos ignorá-lo no que diz respeito ao treinamento.

Processo reverso e L_1:T-1

Agora discutimos as escolhas necessárias para definir o processo inverso. Lembre-se de que definimos as transições reversas de Markov como gaussianas:

Devemos agora definir as formas funcionais de μ_θ ou Σ_θ. Embora existam maneiras mais complicadas de parametrizar Σ_θ^[5], simplesmente definimos

Ou seja, assumimos que a gaussiana multivariada é um produto de gaussianas independentes com variância idêntica, um valor de variância que pode mudar com o tempo. Nós defina essas variações para serem equivalentes ao nosso cronograma de variação de processo para frente.

Dada esta nova formulação de Σ_θ, temos

que nos permite transformar

para

onde o primeiro termo da diferença é uma combinação linear de x_t e x₀ que depende do esquema de variância β_t. A forma exata desta função não é relevante para nossos propósitos, mas pode ser encontrada em [3].

O significado da proporção acima é que a parametrização mais direta de μ_θ simplesmente prevê a média posterior de difusão. É importante ressaltar que os autores de [3] realmente descobriu que o treinamento μ_θ prever o barulho componente em qualquer passo de tempo produz melhores resultados. Em particular, deixe

onde

Isso leva à seguinte função de perda alternativa, que os autores de [3] encontrados para levar a um treinamento mais estável e melhores resultados:

Os autores de [3] também observam conexões desta formulação de Modelos de Difusão com modelos generativos de correspondência de pontuação baseados na dinâmica de Langevin. De fato, parece que os modelos de difusão e os modelos baseados em pontuação podem ser dois lados da mesma moeda, semelhante ao desenvolvimento independente e simultâneo da mecânica quântica baseada em ondas e da mecânica quântica baseada em matriz, revelando duas formulações equivalentes dos mesmos fenômenos^[2].

Arquitetura de rede

Enquanto nossa função de perda simplificada busca treinar um modelo ϵ_θ, ainda não definimos a arquitetura deste modelo. Observe que o só O requisito para o modelo é que sua dimensionalidade de entrada e saída sejam idênticas.

Dada essa restrição, talvez não seja surpreendente que os Modelos de Difusão de imagem sejam comumente implementados com arquiteturas do tipo U-Net.

Decodificador de Processo Reverso e L₀

O caminho ao longo do processo reverso consiste em muitas transformações sob distribuições gaussianas condicionais contínuas. Ao final do processo inverso, lembre-se de que estamos tentando produzir um imagem, que é composto por valores de pixel inteiros. Portanto, devemos criar uma maneira de obter probabilidades discretas (log) para cada valor de pixel possível em todos os pixels.

A maneira como isso é feito é definindo a última transição na cadeia de difusão reversa para um decodificador discreto independente. Para determinar a probabilidade de uma dada imagem x₀ dado x₁, primeiro impomos independência entre as dimensões de dados:

onde D é a dimensionalidade dos dados e o sobrescrito i indica a extração de uma coordenada. O objetivo agora é determinar a probabilidade de cada valor inteiro para um determinado pixel dado a distribuição entre os valores possíveis para o pixel correspondente na imagem com pouco ruído no tempo t=1:

onde as distribuições de pixels para t = 1 são derivadas da gaussiana multivariada abaixo, cuja matriz de covariância diagonal nos permite dividir a distribuição em um produto de gaussianas univariadas, uma para cada dimensão dos dados:

Assumimos que as imagens consistem em inteiros em 0,1,…,255 (como as imagens RGB padrão) que foram dimensionadas linearmente para [−1,1]. Em seguida, dividimos a linha real em pequenos “baldes”, onde, para um determinado valor de pixel dimensionado x, o bucket para esse intervalo é [x−1/255, x+1/255]. A probabilidade de um valor de pixel x, dada a distribuição gaussiana univariada do pixel correspondente em x₁, é o área sob essa distribuição gaussiana univariada dentro do balde centrado em x.

Abaixo você pode ver a área de cada um desses baldes com suas probabilidades de uma média 0 Gaussiana que, neste contexto, corresponde a uma distribuição com um valor médio de pixel de 255/2 (meio brilho). A curva vermelha representa a distribuição de um pixel específico no t = 1 imagem, e as áreas dão a probabilidade do valor de pixel correspondente no t = 0 imagem.

Dado um t = 0 valor de pixel para cada pixel, o valor de p_θ(x₀|x₁) é simplesmente seu produto. Este processo é sucintamente encapsulado pela seguinte equação:

onde

e

Dada esta equação para p_θ(x₀|x₁), podemos calcular o termo final de L_VLB que não é formulada como uma Divergência KL:

Objetivo final

Como mencionado na última seção, os autores de [3] descobriram que prever o componente de ruído de uma imagem em um determinado intervalo de tempo produziu os melhores resultados. Em última análise, eles usam o seguinte objetivo:

Os algoritmos de treinamento e amostragem para nosso Modelo de Difusão, portanto, podem ser capturados de forma sucinta na figura abaixo:

Resumo da Teoria do Modelo de Difusão

Nesta seção, mergulhamos detalhadamente na teoria dos Modelos de Difusão. Pode ser fácil ficar preso em detalhes matemáticos, então observamos os pontos mais importantes nesta seção abaixo para nos mantermos orientados a partir de uma perspectiva panorâmica:

1. Nosso Modelo de Difusão é parametrizado como um Cadeia de Markov, o que significa que nossas variáveis ​​latentes x₁,…, X_T dependem apenas do passo de tempo anterior (ou seguinte).
2. A distribuições de transição na cadeia de Markov são Gaussian, onde o processo direto requer uma programação de variação e os parâmetros do processo reverso são aprendidos.
3. O processo de difusão garante que x_T is distribuída assintoticamente como uma gaussiana isotrópica para T suficientemente grande.
4. No nosso caso, o cronograma de variação foi corrigido, mas também pode ser aprendido. Para horários fixos, seguir uma progressão geométrica pode proporcionar melhores resultados do que uma progressão linear. Em ambos os casos, as variâncias geralmente aumentam com o tempo na série (ou seja, β_i<β_j para mim
5. Os modelos de difusão são altamente flexível e permitir qualquer arquitetura cuja dimensionalidade de entrada e saída são as mesmas a serem utilizadas. Muitas implementações usam Tipo U-Net arquiteturas.
6. A objetivo do treinamento é maximizar a probabilidade dos dados de treinamento. Isso se manifesta como o ajuste dos parâmetros do modelo para minimizar o limite superior variacional da probabilidade logarítmica negativa dos dados.
7. Quase todos os termos da função objetivo podem ser expressos como Divergências KL como resultado de nossa suposição de Markov. Esses valores tornar-se sustentável para calcular dado que estamos usando gaussianas, portanto omitindo a necessidade de realizar a aproximação de Monte Carlo.
8. Em última análise, usando um objetivo de treinamento simplificado treinar uma função que prevê o componente de ruído de uma determinada variável latente produz os melhores e mais estáveis ​​resultados.
9. A decodificador discreto é usado para obter probabilidades de log em valores de pixel como a última etapa do processo de difusão reversa.

Com esta visão geral de alto nível dos modelos de difusão em mente, vamos ver como usar modelos de difusão no PyTorch.

Modelos de difusão no PyTorch

Embora os Modelos de Difusão ainda não tenham sido democratizados no mesmo grau que outras arquiteturas/abordagens mais antigas em Aprendizado de Máquina, ainda existem implementações disponíveis para uso. A maneira mais fácil de usar um modelo de difusão no PyTorch é usar o denoising-diffusion-pytorch pacote, que implementa um modelo de difusão de imagem como o discutido neste artigo. Para instalar o pacote, basta digitar o seguinte comando no terminal:

pip install denoising_diffusion_pytorch

Exemplo mínimo

Para treinar um modelo e gerar imagens, primeiro importamos os pacotes necessários:

import torch
from denoising_diffusion_pytorch import Unet, GaussianDiffusion

Em seguida, definimos nossa arquitetura de rede, neste caso uma U-Net. o dim O parâmetro especifica o número de mapas de recursos antes da primeira redução de amostragem e o dim_mults O parâmetro fornece multiplicandos para este valor e downsamplings sucessivos:

model = Unet(
 dim = 64,
 dim_mults = (1, 2, 4, 8)
)

Agora que nossa arquitetura de rede está definida, precisamos definir o próprio Modelo de Difusão. Passamos o modelo U-Net que acabamos de definir junto com vários parâmetros – o tamanho das imagens a serem geradas, o número de passos de tempo no processo de difusão e uma escolha entre as normas L1 e L2.

diffusion = GaussianDiffusion(
 model,
 image_size = 128,
 timesteps = 1000, # number of steps
 loss_type = 'l1' # L1 or L2
)

Agora que o Modelo de Difusão está definido, é hora de treinar. Geramos dados aleatórios para treinar e, em seguida, treinamos o Modelo de Difusão da maneira usual:

training_images = torch.randn(8, 3, 128, 128)
loss = diffusion(training_images)
loss.backward()

Uma vez que o modelo é treinado, podemos finalmente gerar imagens usando o sample() método do diffusion objeto. Aqui geramos 4 imagens, que são apenas ruído, pois nossos dados de treinamento foram aleatórios:

sampled_images = diffusion.sample(batch_size = 4)

Treinamento sobre dados personalizados

A denoising-diffusion-pytorch O pacote também permite treinar um modelo de difusão em um conjunto de dados específico. Basta substituir o 'path/to/your/images' string com o caminho do diretório do conjunto de dados no Trainer() objeto abaixo e altere image_size ao valor adequado. Depois disso, simplesmente execute o código para treinar o modelo e, em seguida, experimente como antes. Observe que o PyTorch deve ser compilado com CUDA habilitado para usar o Trainer classe:

from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer
model = Unet(
 dim = 64,
 dim_mults = (1, 2, 4, 8)
).cuda()
diffusion = GaussianDiffusion(
 model,
 image_size = 128,
 timesteps = 1000, # number of steps
 loss_type = 'l1' # L1 or L2
).cuda()
trainer = Trainer(
 diffusion,
 'path/to/your/images',
 train_batch_size = 32,
 train_lr = 2e-5,
 train_num_steps = 700000, # total training steps
 gradient_accumulate_every = 2, # gradient accumulation steps
 ema_decay = 0.995, # exponential moving average decay
 amp = True # turn on mixed precision
)
trainer.train()

Abaixo você pode ver a redução progressiva de ruído de ruído Gaussiano multivariado para dígitos MNIST semelhante à difusão reversa:

Palavras finais

Modelos de difusão são uma abordagem conceitualmente simples e elegante para o problema de geração de dados. Seus resultados de última geração combinados com treinamento não adversário os impulsionou a grandes alturas, e outras melhorias podem ser esperadas nos próximos anos, devido ao seu status nascente. Em particular, os modelos de difusão foram considerados essenciais para o desempenho de modelos de ponta como DALL-E2.

Referências

[1] Aprendizado Profundo Não Supervisionado usando Termodinâmica Sem Equilíbrio

[2] Modelagem Gerativa por Estimativa de Gradientes da Distribuição de Dados

[3] Modelos probabilísticos de difusão sem ruído

[4] Técnicas aprimoradas para treinamento de modelos generativos baseados em pontuação

[5] Modelos probabilísticos de difusão de redução de ruído aprimorados

[6] Modelos de difusão vencem GANs na síntese de imagens

[7] GLIDE: Rumo à geração e edição de imagens fotorrealistas com modelos de difusão guiados por texto

[8] Geração de imagem condicional de texto hierárquica com CLIP Latents

Gostou deste artigo? Inscreva-se para mais atualizações de pesquisa de IA.

Avisaremos quando lançarmos mais artigos de resumo como este.

Relacionado