Modelos de deep learning inspirados na termodinâmica

Os modelos de difusão inspirados na física foram introduzidos pela primeira vez por Sohl-Dickstein et al em seu artigo de 2015, "Deep Unsupervised Learning using Nonequilibrium Thermodynamics". Seu algoritmo aplicou a dinâmica de Langevin (link externo a ibm.com), um método para modelar o movimento de sistemas moleculares, para sustentar a premissa básica dos modelos de difusão: transformar dados em ruído, para que você possa transformar o ruído em dados.

Uma nota sobre funções de densidade de probabilidade
Como a maioria dos modelos generativos, como os autocodificadores variacionais (VAEs), o algoritmo de Sohl-Dickstein modelou a densidade de probabilidade: a probabilidade relativa de uma variável amostrada aleatoriamente, x, cair dentro de um intervalo específico de valores. Essencialmente, modelar uma função de densidade de probabilidade para um conjunto de dados de treinamento permite que um algoritmo gere amostras com alta probabilidade de se ajustarem à distribuição de dados de treinamento. Ao gerar uma nova imagem, o modelo assume uma alta probabilidade de os valores de pixel serem distribuídos daquela maneira específica, com base na distribuição de probabilidade que ele aprendeu a partir de padrões em dados de treinamento.

Logicamente falando, as funções de densidade de probabilidade exigem que a probabilidade de todas as possibilidades somem 1. Em outras palavras, a chance percentual de todas as possibilidades deve somar exatamente 100%. Na prática, isso geralmente requer uma constante de normalização: um valor incorporado a uma função de probabilidade que reduz a probabilidade total a 1.

Calcular uma constante de normalização que funcione para todos os valores de variável possíveis é frequentemente intratável: tecnicamente solucionável, mas exigindo tempo infinito para calcular. Nesses casos, os modelos baseados em probabilidade devem ser restritos a arquiteturas de modelo específicas ou desenvolver soluções alternativas inteligentes que se aproximem da constante de normalização de uma forma tratável.

Modelos generativos baseados em pontuação
 

Independentemente do trabalho de Sohl-Dickstein, Yang Song e Stefano Ermon desenvolveram um tipo de modelo baseado em energia (link externo a ibm.com) chamado de rede de pontuação condicional de ruído em seu artigo de 2019, “Generative Modeling by Estimating Gradients of the Data Distribution”. O algoritmo deles modelou o gradiente (∇_x) do logaritmo ($Log$) da função de densidade de probabilidade $p\left(X\right)$. O gradiente da função de densidade de probabilidade logarítmica, escrito como ${\nabla }_X\log p\left(X\right)$é chamado de pontuação de Stein ou simplesmente de "função de pontuação"."

Ao contrário das funções convencionais de densidade de probabilidade, as funções de pontuação não exigem uma constante de normalização porque não modelam diretamente a densidade de probabilidade (e, portanto, não precisam normalizar a probabilidade total para 1). Em vez disso, elas são treinadas por meio da correspondência de pontuação: parâmetros do modelo de aprendizado, θ, que produz um modelo p_θ(x) cuja pontuação — em outras palavras, seu gradiente — corresponde à da distribuição de dados q(x) dos dados de treinamento.

Outro benefício desses modelos generativos baseados em pontuação (SGMs) é que, diferentemente dos modelos baseados em verossimilhança, eles não impõem muitas restrições à arquitetura do modelo de p_θ (x).

Explorando maneiras de melhorar o desempenho de seu modelo, Song e Ermon coincidentemente chegaram às mesmas técnicas empregadas por Sohl-Dickstein et al. Seu artigo observou que "perturbar dados com ruído gaussiano aleatório torna a distribuição de dados mais receptiva à modelagem generativa baseada em pontuação". O modelo deles, construído usando a arquitetura U-Net, originalmente desenvolvida para segmentação de imagens, da mesma forma aplicou a dinâmica de Langevin para gerar amostras.

Modelos probabilísticos de difusão de redução de ruído (DDPMs)

Em 2020, Ho et al propuseram usar a abordagem de Sohl-Dickstein para gerar imagens de alta qualidade usando inferência variacional em seu artigo seminal, "Modelos probabilísticos de difusão de redução de ruído" (DDPMs). O artigo mostrou que maximizar o limite inferior de evidência (ELBO)— uma maneira de reescrever problemas de otimização baseados em probabilidade para serem tratáveis — para treinar modelos de difusão é essencialmente equivalente à combinação de objetivos de correspondência de pontuação usados para treinar SGMs.

Implementando a abordagem de Sohl-Dickstein usando correspondência de pontuações, Ho et al demonstraram que modelos probabilísticos de difusão podem alcançar qualidade de imagem competitiva com GANs, que na época eram de última geração. Essas conexões foram exploradas ainda mais por Song, Ermon, Sohl-Dickstein e outros (incluindo Diederik P. Kingma, criador do VAE) em seu artigo de 2021, "Score-Based Generative Modeling through Stochastic Diferencial Equations".

Posteriormente naquele ano, Dhariwal e Nichol, aproveitando insights do artigo mencionado anteriormente, publicaram “Diffusion Models Beat GANs on Image Synthesis”, estabelecendo firmemente os modelos de difusão como os mais avançados existentes.

Os DDPMs, em vez dos SGMs, geralmente continuam sendo o modo dominante de modelos de difusão, embora com melhorias introduzidas por meio de pesquisas subsequentes. Por exemplo, o influente artigo de 2022 “High-Resolution Image Synthesis with Latent Diffusion Models” marcou avanços importantes em eficiência e relação custo/benefício.

O objetivo do processo de difusão para a frente é transformar dados limpos do conjunto de dados de treinamento, como uma imagem ou amostra de áudio, em ruído puro. O método mais comum envolve a injeção iterativa de ruído gaussiano até que toda a distribuição de dados seja gaussiana.

Em termos matemáticos, esse processo passo a passo é formulado como uma cadeia de Markov: um tipo de processo estocástico, um processo aleatório que segue certas regras probabilísticas, para modelar dados sequenciais de séries temporais. Em uma cadeia de Markov, o resultado em cada etapa de tempo é influenciado apenas pela etapa de tempo imediatamente anterior a ela. Simplificando: x_t, o estado da cadeia de Markov x na etapa de tempo t, é diretamente influenciado apenas por x_t-1. A função matemática que define a transição de qualquer x_t para x_t+1 é chamada de kernel de transição.

Em cada etapa de tempo t, uma pequena quantidade de ruído gaussiano é adicionada a x_t-1 e a imagem é, então, redimensionada, para manter um tamanho de imagem constante, apesar da injeção contínua de pixels aleatórios. Nessa formulação, x₀ é o ponto de dados limpos original; x₁ é o ponto de dados após a primeira etapa de tempo, com uma pequena quantidade de ruído gaussiano adicionado a ele; x_T é o estado final do processo de difusão para a frente. Se T for grande o suficiente (ou seja, após etapas suficientes), x_T convergirá para ruído gaussiano puro.

Definimos cada etapa para a frente como $q\left(X_t\right|X_{t-1}):$prever o estado da distribuição de dados q (x_t), dado q(x_t-1). Em um DDPM padrão, esse processo para a frente não envolve nenhum aprendizado de máquina: o resultado final da cadeia de Markov será sempre uma distribuição gaussiana e, portanto, não requer otimização.

Adição de ruído

Em um DDPM, o ruído gaussiano adicionado em cada etapa na cadeia de Markov não é constante nem arbitrário. Em vez disso, o ruído é derivado da estrutura da imagem original, e a taxa na qual ele é adicionado aumenta constantemente a cada etapa consecutiva.

Variar a quantidade de ruído melhora a estabilidade do treinamento do modelo e melhora o desempenho geral, ao equilibrar duas prioridades concorrentes. Conforme observado por Yang Song em seu post de blog sobre modelos generativos baseados em pontuação (link externo a ibm.com):

• Um ruído maior melhora a capacidade do modelo de aprender com precisão em regiões de "baixa densidade" de dados de treinamento (categorias e conceitos visuais que têm menos representação nos dados de treinamento) ao preencher essas regiões com dados ruidosos. Mas também pode corromper dados em excesso, reduzindo a precisão geral.
• Um ruído menor causa menos corrupção da distribuição de dados original, mas produz um desempenho ruim em regiões de baixa densidade.
• Portanto, para alcançar o melhor dos dois mundos, os modelos de difusão usam várias escalas de ruído no treinamento.

Lembre-se de que qualquer distribuição gaussiana (normal) tem uma média, $\mu$ , e uma variância, $\Sigma$ . Outro parâmetro, β, serve como um fator de escala para a média e variância do ruído gaussiano no kernel de transição que define cada etapa na cadeia de Markov para a frente. Alterar o valor de β em uma determinada etapa resulta na alteração do ruído gaussiano adicionado nessa etapa. β₁ é a variância na etapa de tempo 1; β_t é a variância na etapa de tempo t e assim por diante, até β_T.

O valor da taxa de β em cada etapa é, por sua vez, determinado pela tabela de variâncias.

Tabela de variâncias

Em cada etapa t, a imagem é ligeiramente deslocada de sua iteração na etapa anterior (pela média), e é adicionado ruído a essa versão deslocada da imagem (de acordo com a variância). A magnitude de cada deslocamento e a adição de ruído são determinados pelo valor de  ${\beta }_t$: como ${\beta }_t$ aumenta de acordo com a tabela de variâncias, a taxa de difusão aumenta constantemente como resultado. β é sempre um valor entre 0 e 1: portanto, 0 < β₁ < β₂ < … < β_T < 1.

Escolher uma tabela de variâncias específica para β é uma consideração importante. Geralmente é definida manualmente como um hiperparâmetro, seja fixado a um valor constante ou procedendo de acordo com alguma fórmula com um valor inicial e final predeterminados para β. No artigo do DDPM, Ho et al usaram uma tabela linear com 1.000 etapas em que β₁ = 10^-4 e β_T = 0,02. Pesquisas posteriores descobriram melhorias no desempenho e na eficiência com outros tipos de tabelas, como uma tabela de cossenos,^[1] ou tornando a tabela em si outro parâmetro aprendido.^[2]

O valor de β_t determina a média e a variância do ruído gaussiano adicionado na etapa t.

• A média $\mu$ do ruído gaussiano adicionado na etapa de tempo t, ${\mu }_t$, é calculada como ${\mu }_t=\left(1-{\beta }_t\right)X_{t-1}$. Em linguagem simples, a média do ruído adicionado em cada etapa t é simplesmente uma versão em escala da imagem da etapa anterior, xt-1. O tamanho de _βt determina o quanto essa média se desvia da etapa anterior: quando _βt é muito pequeno, esse deslocamento é muito pequeno, pois $\left(1-{\beta }_t\right)\approx (1-0)\approx 1$— e o ruído adicionado, assim, se assemelharão muito à imagem original. À medida que o valor de β_t aumenta, esse deslocamento se torna mais significativo.
  
  
• A variância do ruído gaussiano adicionado na etapa de tempo t é calculada como ${\Sigma }_t={\beta }_{t}I$, onde $I$ é a matriz identidade (link externo a ibm.com). Um β_t maior resulta em mais ruído. Um β_t muito pequeno resulta em ruído insignificante.

Resumindo, em cada etapa t, a imagem é ligeiramente deslocada de sua iteração na etapa anterior (de acordo com a média), e o ruído é adicionado a essa versão deslocada da imagem (de acordo com a variância). A magnitude de cada deslocamento e a adição de ruído são determinadas pelo valor de β_t: à medida que β_t aumenta de acordo com a tabela de variâncias, a taxa de difusão aumenta constantemente como resultado.

Como a adição de ruído gaussiano começa gradualmente, e o ruído em si é sempre derivado da estrutura essencial da imagem original na etapa anterior, as qualidades essenciais da imagem original são retidas por muitas etapas. Isso permite que o modelo aprenda de forma significativa os padrões e a estrutura da distribuição de dados original durante o processo de difusão reversa.

O truque da reparametrização

Uma falha desse processo passo a passo é que ele é complicado e computacionalmente caro: para cada imagem em um conjunto de dados de treinamento que pode conter milhares ou milhões de imagens, o processo para a frente exigiria dezenas ou centenas de etapas individuais.

Em vez de adicionar ruído repetidamente, a fórmula para o processo para a frente pode ser reescrita de forma inteligente, reparametrizando a expressão de  $1–{\beta }_t$ como um novo parâmetro,  ${\alpha }_t$. Por meio de uma "propriedade agradável" das cadeias de Markov, esse novo parâmetro pode ser estendido ainda mais para um parâmetro adicional, ${\overline{\alpha }}_t$, derivado da multiplicação iterativa de ${\alpha }_t$ em cada etapa progressiva da cadeia até aquele ponto. Este parâmetro adicional reflete essencialmente a relação sinal/ruído (SNR) de x_t: em outras palavras, quanto da imagem original permanece na etapa de tempo t.

Por exemplo, em x₁, uma pequena quantidade de ruído foi adicionada uma vez. O valor de  ${\overline{\alpha }}_t$ é próximo de 1, o que significa que a imagem ainda retém a maior parte do seu "sinal" original. Em uma etapa posterior, como x₅₀, o ruído foi adicionado muitas vezes. Porque ${\alpha }_t=1–{\beta }_t$, o valor de  ${\alpha }_t$ é sempre menor que 1. Desde ${\overline{\alpha }}_{50}={\alpha }_1\bullet {\alpha }_2\bullet ...\bullet {\alpha }_{49}\bullet {\alpha }_{50}$, o valor de  ${\overline{\alpha }}_t$ na etapa 50 estará muito mais próximo de 0, o que significa que uma maior parte da imagem original foi substituída por ruído. Na etapa de tempo T, x_T é totalmente ruído, e o valor de  ${\overline{\alpha }}_t$ se aproxima de 0.

Embora a derivação complexa da equação esteja além do escopo deste artigo, há duas conclusões importantes para entender a importância desse truque de reparametrização:

• O estado de x em qualquer etapa de tempo t agora pode ser definido como $X_t=\sqrt{{\overline{\alpha }}_t}\bullet X_0+\sqrt{1-{\overline{\alpha }}_t}\bullet {\epsilon }_0$, onde ${\epsilon }_0$ é o ruído adicionado na primeira etapa. Como o valor de ${\overline{\alpha }}_t$ diminui a cada etapa, a influência de x₀ diminui, e a influência de ε₀ aumenta.
  
  
• Porque ${\overline{\alpha }}_t$ é derivado de $1–{\beta }_t$ e o valor de ${\beta }_t$ é determinado pela tabela de variâncias; reescrever a fórmula dessa forma permite o cálculo direto de x_t em qualquer etapa de tempo t sem ter que passar por todo o processo passo a passo para a frente.

Nos modelos de difusão, o processo de difusão reversa é onde ocorre o aprendizado de máquina real. Ao aprender a executar a reversão das etapas de ruído do processo para a frente, o modelo está essencialmente aprendendo a reduzir o ruído gaussiano puro em uma imagem limpa. Uma vez que a rede neural tenha sido treinada, essa capacidade pode ser usada para gerar novas imagens a partir do ruído gaussiano por meio de difusão reversa passo a passo.

Em teoria, a tarefa do modelo pode ser considerada o simples reverso da difusão para a frente. O processo para a frente, começando com o ponto de dados x₀ amostrado da distribuição real de dados q(x) do conjunto de dados de treinamento, é definido como $q\left(X_t\right|X_{t-1})$: isto é, dado $q\left(X_{t-1}\right)$, calcular  $q\left(X_t\right)$. Seu oposto, difusão reversa, seria definido como $q\left(X_{t-1}\right|X_t)$. Mas, na prática, a computação $q\left(X_{t-1}\right|X_t)$ é intratável.

Em vez disso, a tarefa de treinamento é formulada por meio de duas soluções alternativas:

• Conforme descrito anteriormente, q(x) é aproximado com uma rede neural p_θ(x) que constitui o próprio modelo de difusão real. O objetivo do treinamento é aprender os parâmetros do modelo θ que fazem com que a produção de p_θ( $X_{t-1}|X_t$ ), corresponda à produção de q( $X_{t-1}|X_t$ ).
  
  
• Esse modelo, p_θ(x), não prevê diretamente $X_{t-1}$, nem mesmo o ruído específico adicionado entre x_t-1 e x_t. Em vez disso, prevê todo o ruído presente in x_t, em seguida, remove uma fração desse ruído (com base no estado da tabela de variâncias nessa etapa) para chegar a x_t-1. O artigo original do DDPM simplificou ainda mais esse processo, estimando apenas a média do ruído gaussiano, embora modelos posteriores também prevejam a variância.

Lembre-se novamente de que a média do ruído gaussiano adicionado na difusão para a frente não é arbitrária: embora seja de fato aleatória, a estrutura do ruído é inicialmente derivada da estrutura da imagem original x₀. Portanto, ao aprender a prever com precisão o ruído por meio da difusão reversa, o modelo não apenas aprende a reduzir o ruído da imagem, mas também aprende implicitamente a estrutura de x₀.

Função de perda para treinamento de modelos de difusão

O objetivo específico de treinamento usado para modelos de difusão está intimamente relacionado ao termo de perda de reconstrução usado para otimizar autocodificadores variacionais (VAEs). Como os VAEs, os modelos de difusão são otimizados maximizando o limite inferior variacional (VLB), também chamado de limite inferior de evidência (ELBO), de uma combinação de termos de perda múltipla.

A maximização do VLB é usada na inferência variacional para aproximar a função de pontuação intratável ${\nabla }_X\log \left(p\right(X\left)\right)$: em vez de minimizar diretamente o erro, reformula a equação como maximizando a estimativa mínima (ou limite inferior) da precisão das previsões do modelo.

Os termos de perda usados refletem a divergência de Kullback-Leibler (ou "divergência KL", geralmente denotada como D_KL) entre os resultados das etapas de difusão para a frente de q e as etapas reversas previstas por p_θ. A divergência KL é usada para medir a diferença entre duas distribuições de probabilidade — por exemplo, entre a distribuição dos valores de pixels em uma imagem e a distribuição dos valores de pixels em outra.

Especificamente, a função de perda para modelos de difusão combina três termos de perda: L_T, L_t e L₀.

•  L_T reflete a divergência KL entre q$\left(X_T\right|X_0)$ e p_θ(x_T). Em outras palavras, a diferença entre o resultado final totalmente ruidoso do processo para a frente q e o ponto de partida do processo reverso. Esse termo geralmente pode ser ignorado, pois x_T é gaussiano e q não tem parâmetros que possam ser aprendidos.
  
  
• L_t reflete a divergência KL entre $q\left(X_{t-1}\right|X_t,X_0)$ e $p_{\theta }\left(X_{t-1}\right|X_t)$ em cada etapa. Em outras palavras, a precisão de cada uma das previsões de redução de ruído de p_{θ durante a difusão reversa, em comparação com cada etapa de ruídocorrespondente durante o processo de difusão para a frente da imagem original, x 0.}
  
  
• L₀ medidas $-\log p_{\theta }\left(X_0\right|X_1)$. Em outras palavras, L₀ reflete a probabilidade logarítmica negativa da previsão do modelo da imagem totalmente sem ruído x₀. O gradiente de L₀ é o termo de correspondência de pontuação descrito anteriormente no artigo. O termo de perda é negativo, de modo que minimizar a função de perda se torna o equivalente a maximizar a probabilidade das previsões do modelo.
  

Embora sua derivação matemática complexa esteja além do escopo deste artigo, o VLB pode ser simplificado para o erro médio quadrático (MSE) entre o ruído previsto pelo modelo, ${\epsilon }_{\theta }$ e o verdadeiro ruído adicionado no processo para a frente, $\epsilon$, em cada etapa de tempo. Isso explica por que a produção do modelo é uma previsão de ruído em cada etapa, em vez da própria imagem sem ruído.

Calculando o gradiente da função de perda durante a retropropagação e, em seguida, ajustando os pesos do modelo para minimizar a função de perda por meio do gradiente descendente, as previsões do modelo em todo o conjunto de dados de treinamento se tornarão mais precisas.

Depois que o modelo de difusão tiver aprendido a estimar com precisão o ruído a ser subtraído em cada etapa, ele poderá ser usado para gerar novas imagens por amostragem a partir de uma imagem ruidosa aleatória x_T da distribuição de dados que aprendeu e eliminando-a por T etapas. Semelhante aos VAEs, introduzir um pequeno elemento de aleatoriedade no processo de amostragem permite que os modelos de difusão produzam novas imagens que se assemelham aos dados de treinamento, em vez de reproduzir diretamente as imagens de treinamento.

Ao contrário do processo de treinamento de difusão reversa, a quantidade de etapas no processo de geração não precisa corresponder à quantidade de etapas usadas no processo para a frente. Isso é possível porque o modelo é treinado para prever todo o ruído em cada etapa da imagem, em vez da quantidade específica de ruído a ser removida nessa etapa.

Menos etapas implicam maior velocidade e menores demandas computacionais, com uma possível compensação em detalhes finos; mais etapas normalmente melhoram a precisão, embora ao custo de diminuição da velocidade e aumento dos custos computacionais.

Modelos de difusão guiada

Enquanto um modelo de difusão padrão pode produzir variações de alta qualidade de imagens de treinamento de forma aleatória, a maioria dos usos práticos de um modelo de geração de imagens exige algum controle sobre a produção do modelo. Os modelos de difusão guiada permitem que um usuário condicione as imagens geradas com uma orientação específica .

A forma mais comum de modelo de difusão guiada é um modelo de difusão de texto para imagem que permite aos usuários condicionar a produção com um prompt de texto, como "uma girafa usando uma cartola". Isso envolve o emparelhamento de um modelo de difusão com um grande modelo de linguagem (LLM) separado para interpretar o prompt de texto, introduzido pela primeira vez pelo Google no artigo "Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding".

Os modelos de difusão padrão visam prever a função de pontuação condicional ${\nabla }_X\log p\left(X\right)$: em outras palavras, o gradiente do logaritmo da probabilidade de que a imagem x gerada pelo modelo p se ajuste aos dados de treinamento x. Os modelos de difusão guiada introduzem uma categoria visual específica, y (por exemplo, "girafa") e preveem a função de pontuação condicional ${\nabla }_X\log p\left(X\right|y)$: em outras palavras, a probabilidade da imagem x, dado que ela deve se encaixar na categoria y.

Os métodos de difusão guiada podem ser separados em duas categories:

• A difusão guiada por classificador requer um modelo de classificador separado para aprender as incorporações vetoriais para cada categoria y para a qual o modelo de difusão será treinado para produzir recursos visuais. Então, essa incorporação vetorial é então usada para condicionar a produção em cada etapa t. O modelo de difusão não requer treinamento extra, mas só poderá condicionar os resultados das categories específicas aprendidas pelo classificador.
  
  
• A orientação livre de classificadores não requer um modelo separado, mas exige um modelo de difusão em dois estágios para ser treinado especificamente para orientação condicional. Isso normalmente envolve um modelo de dois estágios: no primeiro estágio, um algoritmo de incorporação como o CLIP (link externo a ibm.com) gera um y de incorporação para o prompt. No segundo estágio, um modelo de difusão usa essa incorporação para condicionar sua produção. Apesar da sobrecarga de treinamento extra, isso tem o benefício de permitir a orientação zero-shot para categorias de imagens não vistas.