Resumo Técnico: Correspondência de Distribuição Contrastiva para Monte Carlo Sequencial Amortizado em Difusão Discreta

Declaração do Problema

Os modelos de difusão discreta tornaram-se frameworks poderosos para a geração de dados categóricos estruturados, como texto, sequências de DNA e proteínas. Uma aplicação central desses modelos é o alinhamento por recompensa, onde o objetivo é amostrar de uma distribuição-alvo que viés o prior pré-treinado em direção a valores mais altos de uma função de recompensa downstream (por exemplo, preferência humana ou viabilidade de design de proteínas).

Embora o Monte Carlo Sequencial (SMC) Torcido ofereça um método assintoticamente exato para amostrar de tais distribuições inclinadas por recompensa, sua aplicação em espaços de estados discretos enfrenta um gargalo fundamental: a estimação da função de torção ótima.

• Em domínios contínuos, a função de torção pode ser aproximada eficientemente usando a fórmula de Tweedie ou orientação baseada em gradientes.
• Em domínios discretos, a ausência de uma relação análoga força a dependência de aproximações de Monte Carlo (MC) custosas para estimar a função de torção em cada passo de inferência.
• Essa dependência cria uma sobrecarga computacional severa, particularmente quando a própria função de recompensa é cara de avaliar (por exemplo, simulações de dobramento de proteínas ou pontuação de preferência de LLM baseada em API). Aumentar o número de amostras MC ($M$) melhora a precisão, mas aumenta linearmente o tempo de inferência, tornando a abordagem impraticável para muitas aplicações do mundo real.

Metodologia

Os autores propõem a Correspondência de Distribuição Contrastiva (CDM), um framework projetado para amortizar o custo da inferência SMC aprendendo uma função de torção parametrizada via um objetivo de aprendizado contrastivo.

1. Objetivo de Correspondência de Distribuição Contrastiva

Em vez de aprender a função de torção via regressão direta (minimizando o Erro Quadrático Médio contra estimativas MC), a CDM minimiza a divergência Kullback-Leibler (KL) forward entre a distribuição-alvo $p^*_t$ e uma distribuição parametrizada $p^\phi_t$ modulada pela função de torção aprendida $\psi^\phi_t$.

O objetivo é formulado como:
$$\mathcal{L}_{\text{CDM}}(\phi) = \mathbb{E}_{t}\left[ \mathcal{D}_{\text{KL}}\left( p^*_t(\mathbf{x}_t) \big\| p^\phi_t(\mathbf{x}_t) \right) \right]$$

O gradiente dessa perda exibe uma estrutura contrastiva:
$$-\nabla_\phi \mathcal{L} = \mathbb{E}_{t}\left[ \underbrace{\mathbb{E}_{p^*_t}[\nabla_\phi \log \psi^\phi_t]}_{\text{Termo Positivo}} - \underbrace{\mathbb{E}_{p^\phi_t}[\nabla_\phi \log \psi^\phi_t]}_{\text{Termo Negativo}} \right]$$

• Termo Positivo: Aumenta o peso do log-torção em amostras extraídas da distribuição-alvo ótima (regiões de alta recompensa).
• Termo Negativo: Diminui o peso do log-torção em amostras extraídas da aproximação atual (regiões subótimas).

2. Esquema de Treinamento Eficiente

Para abordar a ineficiência de amostrar da distribuição-alvo $p^*_t$ para o termo positivo, os autores exploram uma propriedade única dos modelos de difusão: o kernel forward de forma fechada.

• Estimação de Gradiente Baseada no Forward: Em vez de amostrar independentemente de cada alvo intermediário $p^*_t$, o método primeiro amostra amostras positivas limpas ($\mathbf{x}_0$) da distribuição-alvo usando SMC. Em seguida, utiliza o kernel forward de forma fechada $p^{\text{base}}(\mathbf{x}_t | \mathbf{x}_0)$ para gerar múltiplas amostras positivas em qualquer instante de tempo intermediário $t$ a um custo negligenciável.
• Treinamento Baseado em Buffer: Um buffer de amostras positivas limpas é mantido e reutilizado em múltiplos instantes de tempo e atualizações de gradiente, reduzindo significativamente o número de avaliações de recompensa caras necessárias durante o treinamento.
• Amostragem Negativa: Amostras negativas são extraídas via Amostragem por Importância (IS) usando a torção parametrizada atual, o que é computacionalmente eficiente.

3. Parametrização

Para garantir que a função de torção aprendida não introduza sobrecarga de inferência, a função de torção $\psi^\phi_t$ é parametrizada como uma cabeça escalar leve anexada à camada de características final do backbone de difusão pré-treinado. Isso permite que o modelo produza tanto logits quanto a estimativa de torção em uma única passagem forward, adicionando menos de 5% de sobrecarga computacional em comparação com o modelo base.

Contribuições Principais

1. Framework CDM: Introdução de um novo framework de amortização SMC para difusão discreta que reduz o custo de aplicar a função de torção no tempo de inferência para uma operação de tempo constante.
2. Treinamento Nativo de Difusão: Projeto de um esquema de treinamento que explora o processo forward de forma fechada dos modelos de difusão para permitir treinamento eficiente com funções de recompensa caras, evitando a ineficiência de amostragem de abordagens baseadas em regressão padrão.
3. Versatilidade e Sinergia: Demonstração de que a CDM é agnóstica à distribuição de proposta, permitindo que seja combinada com propostas ajustadas (por exemplo, d1, DRAKES) para ganhos de desempenho sinérgicos.
4. Objetivo Superior: Validação de que o objetivo de aprendizado contrastivo produz desempenho superior e convergência mais rápida em comparação com objetivos de torção baseados em regressão padrão (Soft Value).

Resultados Experimentais

Os autores avaliaram a CDM em quatro aplicações diversas: geração de texto tóxico, design de sequências de DNA regulatórias, viabilidade de design de proteínas e alinhamento de modelo de linguagem grande de difusão (dLLM).

• Desempenho de Escala: Em todas as tarefas, a CDM superou consistentemente as linhas de base (Best-of-N, SMC padrão com $M$ variável e Soft Value) quando comparadas pelo tempo de parede.
• Recompensas Caras: Na tarefa de viabilidade de design de proteínas, onde a avaliação da recompensa é o gargalo dominante, a CDM alcançou uma escala superior em comparação tanto com Best-of-N quanto com Soft Value, enquanto o SMC padrão falhou em escalar efetivamente devido à sobrecarga de inferência.
• Compatibilidade com Ajuste Fino: Quando integrada com propostas ajustadas (d1 e DRAKES), a CDM melhorou ainda mais o desempenho e mitigou problemas de colapso de modo frequentemente observados em modelos ajustados, mantendo a diversidade enquanto alcançava altas recompensas.
• Eficiência de Treinamento: A CDM demonstrou convergência mais rápida em tempo de parede em comparação com Soft Value, atribuída à eficácia do objetivo de aprendizado contrastivo com amostras negativas.

Significado e Alegações

O artigo afirma que a CDM aborda o principal gargalo na inferência SMC torcida para difusão discreta: a dependência de estimação de Monte Carlo custosa. Ao aprender a função de torção via um objetivo contrastivo que explora o processo forward de difusão, o método permite alinhamento por recompensa eficiente e escalável.

Os autores enfatizam que sua abordagem é ortogonal ao ajuste fino de propostas; pode ser combinada com métodos existentes de ajuste fino para alcançar ganhos de desempenho adicionais. Além disso, o método permanece aplicável mesmo quando a função de recompensa é não diferenciável ou computacionalmente cara, um cenário onde a orientação baseada em gradientes e o SMC padrão lutam. Os resultados sugerem que a CDM estabelece uma nova fronteira de Pareto para o trade-off entre qualidade da recompensa e tempo de inferência na modelagem generativa discreta.