Resumo Técnico: Modelo de Difusão Sem Simulação e de Tempo Finito

Declaração do Problema

Modelos de difusão generativa alcançaram um desempenho de estado da arte em vários domínios, mas enfrentam um compromisso fundamental entre treinamento sem simulação e geração de tempo finito.

• Treinamento sem simulação permite que o objetivo de treinamento seja avaliado sem simular a equação diferencial estocástica (SDE) de referência, reduzindo significativamente os custos computacionais. Os Modelos Baseados em Score (SBMs) convencionais alcançam isso utilizando um processo de Ornstein–Uhlenbeck (OU) como referência. No entanto, o processo OU apenas se aproxima da distribuição a priori assintoticamente ($T \to \infty$), tornando a geração de tempo finito ineficiente ou imprecisa para dados de alta dimensão.
• Geração de tempo finito requer que o processo de referência conecte a distribuição de dados empíricos e a distribuição a priori dentro de um horizonte de tempo finito prescrito. Abordagens como a ponte de Schrödinger alcançam isso, mas geralmente exigem a simulação da SDE de referência durante o treinamento (por exemplo, via Correspondência de Score Implícita) porque a equação de Fokker–Planck correspondente é analiticamente intratável.

O desafio central abordado por este trabalho é construir um processo de referência que satisfaça simultaneamente ambas as propriedades: permitir o treinamento sem simular a SDE de referência e garantir que o processo conecte as distribuições a priori e empírica dentro de um tempo finito fixo.

Metodologia

Os autores propõem um framework que inverte o procedimento de design convencional. Em vez de especificar primeiro uma SDE de referência e analisar suas distribuições induzidas, os autores primeiro prescrevem uma família tratável de distribuições condicionais dependentes do tempo $\{\rho_t\}_{t \in [0,1]}$ e então constroem uma SDE de referência que realiza essas distribuições como suas marginais.

1. Construção do Processo de Referência

Seja $\mu$ a distribuição empírica e $\pi$ a a priori. Os autores definem uma família de distribuições condicionais $\rho_t(z|x)$ satisfazendo:

1. Regularidade: $\rho_t$ é $C^{1,2}$ em relação ao tempo e ao espaço.
2. Tratabilidade: Amostras de $\rho_t(\cdot|x)$ podem ser extraídas diretamente.
3. Condições de Contorno: As marginais induzidas $p_t(z) = \int \rho_t(z|x)\mu(x)dx$ satisfazem $p_0 = \pi$ e $p_1 = \mu$.

Dadas essas condições, os autores derivam coeficientes $\alpha(t, z, x)$ (drift condicional) e $b(t, z)$ (coeficiente de difusão) de modo que $\rho_t$ satisfaça a equação de Fokker–Planck:
$$\partial_t \rho_t = -\nabla_z \cdot [\alpha(t, z, x)\rho_t] + \frac{1}{2} \sum_{i,j} \partial_{z_i}\partial_{z_j} [\Gamma_{ij}(t, z)\rho_t]$$
onde $\Gamma = bb^\top$. O processo de referência $P$ é então definido pela SDE:
$$dZ_t = a(t, Z_t)dt + b(t, Z_t)dW_t, \quad Z_0 \sim p_0$$
onde o drift $a(t, z)$ é a esperança de $\alpha$ em relação a $\rho_t$.

2. Objetivo Sem Simulação

Ao aplicar o teorema de Girsanov, a divergência Kullback–Leibler (KL) de espaço de trajetória entre o processo de referência $P$ e o processo de geração $Q_\theta$ é reformulada. Crucialmente, como $\rho_t$ é prescrita e tratável, a função objetivo pode ser avaliada sem simular a SDE de referência:
$$L_{SF}(\theta) = \frac{1}{2n} \sum_{i=1}^n \int_0^1 dt \, \mathbb{E}_{Z \sim \rho_t(\cdot|x_i)} \left[ \| \alpha(t, Z, x_i) - s_\theta(t, Z) \|_{\Gamma(t, Z)^{-1}}^2 \right]$$
Este objetivo depende apenas da amostragem direta de $\rho_t$ e do drift condicional $\alpha$, evitando a necessidade do score da marginal $p_t$ (que tipicamente requer reversão temporal e é intratável para pontes de tempo finito).

3. Construções Práticas

O artigo fornece construções específicas para dois casos:

• Prior Gaussiana: $\rho_t$ é escolhida como uma distribuição Gaussiana com média e variância dependentes do tempo. Isso recupera uma forma semelhante à correspondência de score de denoising, mas com um horizonte de tempo finito.
• Priors Não-Gaussianas: Uma construção de "push-forward" é utilizada. Um processo de difusão base com distribuição estacionária $\pi$ é transformado via uma bijeção suave dependente do tempo $\phi^x_t$. Isso permite que o framework lide com priors de cauda pesada ou não-gaussianas (ex: distribuição de Johnson's SU) mantendo a propriedade sem simulação.

Principais Contribuições

1. Framework Unificado: O artigo introduz um framework geral que alcança simultaneamente o treinamento sem simulação e a geração de tempo finito, resolvendo o compromisso inerente aos modelos de difusão convencionais.
2. Reinterpretação da Correspondência de Score: Os autores demonstram que a correspondência de score não é fundamental para o treinamento de modelos de difusão. Em vez disso, ela emerge naturalmente como uma consequência de reverter o tempo de um processo de referência definido na direção dado-para-prior. Ao construir o processo de referência diretamente na direção prior-para-dado, o objetivo de treinamento torna-se independente da função de score derivada da reversão temporal.
3. Conexão com Flow Matching: O artigo mostra que o Conditional Flow Matching (CFM) surge como o limite de baixo ruído ($\epsilon \to 0$) do proposto framework estocástico. Isso fornece um elo teórico entre modelos de difusão estocásticos e modelos baseados em fluxo determinísticos, esclarecendo que o CFM pode ser visto como um limite determinístico de um objetivo de KL de espaço de trajetória.
4. Priors Não-Gaussianas: A metodologia se estende a priors não-gaussianas via construção push-forward, demonstrando flexibilidade além da suposição padrão Gaussiana.

Resultados

Os autores conduziram experimentos numéricos em conjuntos de dados de brinquedo bidimensionais (mistura gaussiana, espiral, tabuleiro de xadrez e dois meses) usando tanto priors Gaussianos padrão quanto de Johnson's SU (cauda pesada).

• Desempenho: O método proposto aprendeu com sucesso a estrutura qualitativa de todos os conjuntos de dados com ambos os tipos de prior.
• Eficiência: Diferente dos SBMs, que exigem o ajuste do horizonte de tempo $T$ para equilibrar a convergência para a prior contra erros de discretização, o método proposto opera em um intervalo fixo $[0, 1]$. Os experimentos mostraram que o método proposto alcança geração de alta qualidade sem a necessidade de ajuste do horizonte de tempo, enquanto escolhas inapropriadas de $T$ em SBMs degradaram a qualidade da amostra.
• Comparação: Os resultados foram comparáveis a um SBM de preservação de variância (VP-SBM) padrão com um grande horizonte de tempo ($T=10$), mas sem o overhead computacional associado a longos tempos de simulação ou a sensibilidade à seleção de $T$.

Significância e Alegações

O artigo afirma fornecer uma reinterpretação estrutural dos modelos de difusão existentes. Ao mudar o paradigma de design de "especificar SDE $\to$ analisar distribuições" para "especificar distribuições $\to$ derivar SDE", os autores revelam que:

• A dependência de correspondência de score em SBMs é um artefato do procedimento de reversão temporal necessário para alinhar o processo de referência com a direção de geração.
• A distinção entre modelos de difusão estocásticos e flow matching determinístico é uma questão de nível de ruído, sendo o CFM o limite de baixo ruído do objetivo de KL estocástico proposto.

O trabalho sugere que o design das distribuições condicionais $\{\rho_t\}$ é um aspecto crítico, porém subexplorado, da modelagem de difusão que impacta diretamente a qualidade da geração. Os autores concluem que seu princípio de construção oferece um caminho robusto para o design de modelos de difusão de tempo finito eficientes, sem os gargalos computacionais das abordagens convencionais.