Resumo Técnico: Destilação de Mapa de Transição Estocástica para Inferência Probabilística Rápida

Declaração do Problema

Os modelos de difusão estabeleceram desempenho de última geração em qualidade de geração, diversidade e cobertura da distribuição em comparação com predecessores como GANs e VAEs. No entanto, esse desempenho é frequentemente acompanhado por custos de inferência elevados devido à necessidade de integrar equações diferenciais ordinárias (EDOs) de tempo contínuo ou equações diferenciais estocásticas (EDEs) ao longo de muitas etapas.

Embora existam técnicas de aceleração existentes, como destilação (por exemplo, Modelos de Consistência) e abordagens baseadas em Transporte Ótimo, elas frequentemente enfrentam um compromisso entre qualidade/diversidade de geração e complexidade de inferência. Uma limitação crítica dos métodos de destilação atuais é sua dependência de inferência determinística (por exemplo, via EDO de fluxo de probabilidade). Dada uma amostra de ruído fixa, a saída é única. Isso restringe sua utilidade em tarefas downstream que inerentemente requerem inferência estocástica, como amostragem posterior de difusão para problemas inversos, ajuste fino baseado em energia ou melhoria do alinhamento com consultas de usuários. Além disso, muitos frameworks de destilação existentes exigem um modelo professor pré-treinado, otimização de dois níveis ou a simulação e cache de trajetórias, o que complica o treinamento e a escalabilidade.

Metodologia: Destilação de Mapa de Transição Estocástica (STMD)

Os autores propõem a Destilação de Mapa de Transição Estocástica (STMD), um framework sem professor projetado para acelerar a inferência de modelos de difusão enquanto preserva a estrutura de transição estocástica do processo subjacente.

Conceito Central

Diferentemente dos modelos de difusão baseados em pontuação que tipicamente modelam a média da distribuição posterior, a STMD destila o mapa de transição completo associado à EDE de amostragem. O método aprende um gerador condicional que prevê a probabilidade de transição $p(x_{t'} | x_t)$ para $t' < t$ em uma ou poucas etapas, evitando a integração numérica custosa da EDE reversa.

Formulação Técnica

1. Fluxo de Média Condicional (MF): O framework parametriza as transições da EDE usando um modelo de Fluxo de Média Condicional. Em vez de aprender um mapa de fluxo determinístico, ele aprende a velocidade média ao longo da trajetória do processo de difusão.
2. Parametrização em $x_0$: Para modelar a transição $p(x_{t'} | x_t)$, os autores utilizam uma parametrização em $x_0$. Eles expressam a transição como uma integral sobre a posterior $p(x_0 | x_t)$ e a medida de ponte conhecida $p(x_{t'} | x_0, x_t)$. Como a medida de ponte para EDEs que preservam variância (VP) está disponível em forma fechada, o problema reduz-se a aprender um modelo condicional $p_\theta(x_0 | x_t)$.
3. Objetivo de Treinamento: O modelo $u_\theta$ é treinado para aproximar a velocidade média condicional. Os autores definem uma função de perda de Fluxo de Média Condicional (CMF):
   $$L_{dCMF} = \mathbb{E} \left[ \left\| u_\theta(z_s, r, s, x_t, t) - \text{sg}\left( (z_1 - x_0) - (s-r) \frac{d}{ds}u_\theta(z_s, r, s, x_t, t) \right) \right\|^2 \right]$$
   onde $z_s$ é um interpolante de linha reta entre os dados $x_0$ e uma amostra de prioridade $z_1$, condicionado ao estado de difusão $x_t$ e ao tempo $t$. O operador $\text{sg}$ denota o stop-gradient.
4. Treinamento Sem Professor: O processo de treinamento não requer um modelo professor pré-treinado. Ele depende de amostras da distribuição conjunta $p(x_0, x_t)$, que podem ser geradas amostrando $x_0$ da distribuição de dados e aplicando o kernel da EDE direta (que é conhecido em forma fechada).

Inferência

A inferência é realizada integrando o Fluxo de Média condicional aprendido para trás no tempo. O processo envolve:

1. Amostragem de um ruído inicial $x_1$ (ou $z_1$).
2. Aplicação iterativa do campo de velocidade aprendido $u_\theta$ para avançar do tempo $s$ para $s-\Delta s$.
3. Em cada etapa de tempo de difusão $t$, amostragem da distribuição de ponte $p(x_{t-\Delta t} | x_0, x_t)$ para atualizar o estado, simulando efetivamente a transição estocástica em grandes etapas.

Contribuições Principais

1. Framework Inovador: A proposta da STMD, um framework de destilação baseado em Fluxos de Média Condicionais que permite aprender mapas de transição de EDEs de difusão de alta dimensão usando um gerador de uma ou poucas etapas via um objetivo de regressão, sem qualquer professor pré-treinado.
2. Convergência Teórica: A derivação de limites de convergência para o algoritmo de Fluxo de Média na distância de Wasserstein-2. Até onde se sabe dos autores, esta é a primeira análise de convergência do algoritmo de Fluxo de Média nesta métrica. Eles estendem isso para provar que seu mapa de transição condicional aprendido converge para o kernel de transição da EDE reversa subjacente.
3. Validação Empírica: Demonstração do método em tarefas de geração de imagens e inpainting nos conjuntos de dados MNIST, CIFAR-10 e CelebA, mostrando desempenho competitivo com baselines enquanto mantém a estocasticidade.

Resultados Experimentais

Os autores avaliaram a STMD em três conjuntos de dados:

• MNIST: Usando uma arquitetura base U-Net, o modelo foi comparado contra o algoritmo original de Fluxo de Média e um Modelo Probabilístico de Difusão com Remoção de Ruído (DDPM). Os autores utilizaram a Distância Fréchet (FD) no espaço latente de um classificador (já que o FID é pouco confiável para MNIST). Os resultados mostraram que a STMD desempenha de forma comparável às baselines enquanto suporta amostragem estocástica.
• CIFAR-10: Experimentos semelhantes foram conduzidos usando uma U-Net. O modelo alcançou pontuações competitivas de Distância Fréchet Inception (FID) em comparação com DDPM e baselines de Fluxo de Média para um número fixo de iterações de treinamento, mantendo a estocasticidade completa do amostrador de difusão.
• CelebA: Os experimentos utilizaram uma arquitetura DiT latente. O modelo gerou amostras incondicionais com um FID de 8,28 (usando 4 etapas de inferência e 2 etapas de Fluxo de Média) e executou com sucesso tarefas de inpainting de imagens.

Em todos os casos, os modelos foram treinados do zero com as mesmas arquiteturas de rede base e contagens de iterações que as baselines para garantir comparação justa.

Significado e Alegações

O artigo afirma que a STMD aborda a lacuna entre a destilação determinística rápida e a necessidade de inferência estocástica em tarefas downstream. Ao destilar o mapa de transição em vez de apenas o mapa de fluxo, a STMD preserva a estrutura estocástica do processo de difusão.

Os autores enfatizam que sua abordagem é sem professor, eliminando a necessidade de otimização de dois níveis ou cache de trajetórias, o que torna o processo de treinamento mais eficiente e escalável. A contribuição teórica de estabelecer limites de convergência de Wasserstein-2 fornece uma base rigorosa para a abordagem de Fluxo de Média, preenchendo uma lacuna entre garantias teóricas e aplicação prática.

O trabalho se posiciona como uma solução para tarefas que requerem inferência probabilística, como problemas inversos e amostragem posterior, onde os métodos de destilação determinística são insuficientes. Os autores permanecem modestos quanto ao desempenho atual, observando que o número limitado de iterações de treinamento impede resultados de última geração, mas o estudo comparativo valida a eficácia do método em relação às baselines existentes.