Resumo Técnico: Amostras e Pontes de Difusão Discretas

Declaração do Problema

O artigo aborda o desafio de amostrar a partir de uma distribuição $p(x) \propto e^{-E(x)}$ sobre um espaço de estados discreto $\mathcal{S} = \{1, \dots, C\}^d$, onde a função de energia $E$ é conhecida, mas a constante de normalização $Z$ é intratável. Embora amostradores de difusão tenham sido adaptados com sucesso para espaços contínuos usando controle estocástico ótimo e correspondência no espaço de trajetórias, sua aplicação a espaços discretos permanece subdesenvolvida. Amostradores de difusão discretos existentes frequentemente falham em aproveitar técnicas de aprendizado por reforço (RL) off-policy que são padrão em configurações contínuas, levando a uma cobertura de modos subótima e qualidade de amostragem, particularmente em distribuições altamente multimodais. Além disso, o problema de conectar duas distribuições arbitrárias (a ponte de Schrödinger) em espaços discretos tem sido amplamente restrito a configurações onde amostras de ambas as distribuições estão disponíveis (dados-a-dados), deixando o cenário "dados-a-energia" (onde uma distribuição é definida apenas por uma densidade não normalizada) inexplorado.

Metodologia

Os autores propõem um framework que integra estratégias de treinamento off-policy em amostradores de difusão discretos e generaliza esses métodos para pontes de Schrödinger discretas.

1. Treinamento Off-Policy para Amostradores de Difusão Discretos

A metodologia central envolve treinar um kernel forward $\overrightarrow{p}_\theta$ para corresponder a uma distribuição alvo, minimizando uma divergência entre a distribuição da trajetória do processo forward e a distribuição da trajetória do processo backward (ruído). As funções objetivo estudadas são divergências de segundo momento, especificamente:

• Equilíbrio de Trajetória (TB): Trata a constante de normalização como um parâmetro aprendível.
• Log-Variância (LV): Minimiza a variância da razão logarítmica das probabilidades forward e backward.

Para melhorar a exploração e a cobertura de modos, os autores introduzem três estratégias off-policy adaptadas de RL e Redes de Fluxo Generativo (GFlowNets):

• Buffer de Replay: Armazena estados amostrados $X_N$ e os reutiliza amostrando trajetórias backward $\overleftarrow{p}^{\otimes N}$ para treinamento.
• Buffer Ponderado por Importância: Atribui probabilidades de amostragem aos elementos do buffer com base nos pesos de importância das trajetórias que os geraram, priorizando regiões de alta probabilidade que o amostrador atual pode perder.
• Exploração MCMC: Refina amostras do buffer de replay usando passos de Cadeia de Markov Monte Carlo (MCMC) (por exemplo, Swendsen-Wang ou Metropolis-Hastings) antes de serem usadas para treinamento. Isso aproveita a função de energia diretamente para guiar o amostrador em direção a regiões de alta densidade sem exigir execuções do modelo.

2. Pontes de Schrödinger Dados-a-Energia Discretas

O artigo generaliza o problema de amostragem para encontrar uma ponte de Schrödinger (SB) entre duas distribuições arbitrárias $p_0$ e $p_N$. Diferentemente de trabalhos anteriores de SB discreto que assumem amostras para ambas as extremidades, este trabalho aborda o cenário dados-a-energia onde $p_N$ é definida apenas por uma função de energia.

• O método emprega um procedimento de Ajuste Proporcional Iterativo (IPF).
• Alterna entre ajustar o kernel backward (via máxima verossimilhança em amostras de $p_0$) e o kernel forward (via perda de log-variância em amostras da $p_N$ definida por energia).
• Técnicas off-policy (buffers e MCMC) são integradas ao loop IPF para facilitar a exploração da distribuição definida por energia.

3. Aplicação a Espaços Latentes

Os autores aplicam esses amostradores à amostragem de posterior "terceirizada". Eles treinam amostradores de difusão no espaço latente discreto de um Autoencoder Variacional Quantizado Vetorial (VQ-VAE) pré-treinado. O objetivo é amostrar a partir do posterior $p(z|y) \propto p_{latent}(z)p(y|f(z))$ sem observar amostras condicionais não enviesadas, efetivamente condicionando o modelo generativo aprendendo a amostrar em seu espaço latente.

Principais Contribuições

1. Algoritmos Off-Policy para Amostradores Discretos: A introdução de técnicas de RL off-policy (buffers de replay, ponderação por importância e exploração MCMC) no treinamento de amostradores de difusão discretos. O artigo demonstra que essas técnicas melhoram significativamente o desempenho em relação a baselines on-policy.
2. Pontes de Schrödinger Dados-a-Energia: A primeira generalização de amostradores de difusão discretos para o problema de ponte de Schrödinger dados-a-energia, permitindo a construção de mapas de transporte entre uma distribuição baseada em amostras e uma distribuição definida por energia.
3. Amostragem de Posterior em Espaço Latente: A demonstração de amostradores de difusão discretos para inferência posterior em espaços latentes discretos de modelos generativos de imagem, mostrando a capacidade de condicionar em rótulos de classe (por exemplo, dígitos ímpares/pares) sem retreinar o modelo generativo.

Resultados Experimentais

Os autores avaliam seus métodos em benchmarks sintéticos e aplicações do mundo real:

• Modelos de Ising e Potts: Em reticulados $16 \times 16$, métodos off-policy (especificamente TB com Buffer + MCMC) superam consistentemente métodos on-policy e a linha de base Amostrador Neural de Difusão Mascarada (MDNS). Eles alcançam distâncias de Sinkhorn e erros de magnetização/correlação mais baixos, particularmente em baixas temperaturas (alto $\beta$) onde o colapso de modos é prevalente.
• Densidades Sintéticas Discretizadas: Em distribuições multimodais desafiadoras (40GMM e ManyWell) discretizadas via códigos Gray, métodos off-policy descobrem com sucesso todos os modos, enquanto métodos on-policy e MDNS frequentemente falham em cobrir a distribuição completa.
• Pontes de Schrödinger: Em tarefas de ponte entre distribuições simples e multimodais, métodos off-policy escalam efetivamente para alvos complexos onde amostradores on-policy colapsam.
• Amostragem Terceirizada: Visualizações de dígitos MNIST gerados a partir do espaço latente mostram que tanto amostradores on-policy quanto off-policy podem gerar com sucesso dígitos das classes desejadas (por exemplo, ímpares vs. pares), validando a utilidade da abordagem para geração condicional.

Significado e Afirmações

O artigo afirma que integrar técnicas de treinamento off-policy é um passo crítico para frente para amostradores de difusão discretos, abordando a questão do colapso de modos e melhorando a eficiência de amostragem em espaços discretos de alta dimensão e multimodais. Ao fechar a lacuna entre difusão discreta e RL off-policy, os autores mostram que esses amostradores podem lidar com problemas complexos de física estatística (Ising/Potts) e tarefas de otimização combinatória mais efetivamente do que métodos anteriores.

Além disso, o trabalho estende o escopo teórico das pontes de Schrödinger para o cenário dados-a-energia em domínios discretos, oferecendo uma nova ferramenta para transporte entre distribuições onde uma é definida por uma função de energia. Finalmente, a aplicação a espaços latentes demonstra um caminho prático para condicionar modelos generativos pré-treinados sem exigir acesso à distribuição completa de dados ou o retratamento do próprio modelo generativo.

Os autores reconhecem limitações, incluindo a necessidade de fixar um kernel de ruído e o custo computacional do armazenamento de trajetórias, mas sugerem que o treinamento em escalas de tempo mais grosseiras (como suportado por extensões da teoria de tempo contínuo) pode mitigar alguns custos. Eles identificam direções futuras, como aprender funções de energia a partir de dados e escalar esses métodos para modelos generativos de imagem maiores.