Resumo Técnico: Um Framework de Campo Médio para o Controle Distribuicional de Modelos de Difusão em Tempo de Inferência

Declaração do Problema

Modelos de difusão são cada vez mais utilizados como amostradores controláveis, onde os processos de geração são direcionados durante a inferência para satisfazer objetivos específicos. Embora os métodos existentes lidem efetivamente com recompensas pontuais (definidas em amostras individuais, $r: \mathbb{R}^d \to \mathbb{R}$), muitas aplicações críticas exigem recompensas distribuicionais (definidas sobre a medida gerada, $R: \mathcal{P}(\mathbb{R}^d) \to \mathbb{R}$). Esses objetivos distribuicionais incluem calibrar populações para observações experimentais, incentivar a diversidade ou equilibrar modos.

Abordagens atuais para o direcionamento distribuicional frequentemente incorporam termos de gradiente da recompensa na dinâmica de inferência (por exemplo, potenciais repulsivos ou verossimilhanças de ensemble). No entanto, de forma semelhante aos primeiros métodos de direcionamento pontual, essas heurísticas carecem de garantias teóricas em relação à distribuição amostrada. Especificamente, simplesmente adicionar um termo de gradiente não garante que as partículas sigam a medida inclinada (tilted measure) $\mu^*$ que maximiza o objetivo regularizado por KL:
$$\arg \max_{\mu} \left\{ R(\mu) - \text{KL}(\mu \| p_1) \right\}$$
Ao contrário do caso pontual, onde o alvo inclinado é uma inclinação exponencial explícita ($\mu^* \propto p_1 e^r$), o caso distribuicional resulta em um alvo implícito onde o potencial de inclinação depende da própria medida. Atualmente, não existe um framework teoricamente fundamentado para corrigir a dinâmica e os pesos para visar essa distribuição implícita.

Metodologia

Os autores propõem um Framework de Campo Médio (Mean-Field Framework) para formalizar e resolver o controle distribuicional em tempo de inferência.

1. Formulação Teórica

O problema é formulado como o direcionamento de uma medida inclinada $\mu^*_t$ em cada passo de tempo $t \in [0,1]$, definida implicitamente por:
$$\mu^*_t(dx) = \frac{1}{Z_t} e^{\Psi_t(x, \mu^*_t)} p_t(x) dx$$
onde $\Psi_t(x, \mu) = \frac{\delta R_t}{\delta \mu}(x, \mu)$ é a primeira variação (derivada funcional) da recompensa dependente do tempo $R_t$. Isso cria uma relação de autoconsistência onde a medida alvo depende do potencial, que, por sua vez, depende da medida.

2. Dinâmica de McKean-Vlasov Ponderada

Para rastrear este caminho, os autores derivam um sistema de McKean-Vlasov ponderado. Diferente do direcionamento de gradiente padrão, o termo de direcionamento $\nabla_x \Psi_t(X_t; \mu_t)$ induz interação entre as partículas porque depende da medida empírica evolutiva $\mu_t$.

A dinâmica idealizada para uma partícula $X_t$ e seu log-peso $A_t$ é:
$$dX_t = \left[ b_t(X_t) + \epsilon_t s_t(X_t) + \epsilon_t \nabla_x \Psi_t(X_t; \mu_t) \right] dt + \sqrt{2\epsilon_t} dB_t$$
$$dA_t = \left[ b_t(X_t) \cdot \nabla_x \Psi_t(X_t; \mu_t) + \dot{\Psi}_t(X_t) \right] dt$$
Aqui, $\dot{\Psi}_t(x) = \frac{d}{dt} \Psi_t(x, \mu_t)$ é um termo de correção crucial. Ao contrário do caso pontual onde $\dot{\Psi}_t$ é simplesmente $\partial_t r$, aqui ele envolve a evolução temporal da medida $\mu_t$. Os autores mostram que $\dot{\Psi}_t$ satisfaz uma equação integral implícita envolvendo a segunda variação do funcional de recompensa.

3. Algoritmo de Partículas Finitas

Para implementar isso, os autores propõem um esquema de partículas interagentes ponderadas (Algoritmo 1):

• Atualizações de Posição: As partículas são atualizadas usando o drift base, o score e o gradiente da recompensa em relação à medida empírica atual $\hat{\mu}_t$.
• Atualizações de Log-Peso: Os pesos são atualizados para corrigir a discrepância entre a dinâmica direcionada e o caminho inclinado alvo. Isso requer o cálculo de $\dot{\Psi}_t$.
• Resolvendo para $\dot{\Psi}_t$: O artigo fornece dois métodos para calcular o termo de derivada temporal implícita:
  1. Sistema Linear Implícito: Resolvendo um sistema linear derivado da equação integral (Proposição 3.3).
  2. Iteração de Ponto Fixo: Um esquema de iteração de Picard leve que aproxima as atualizações de peso.
• Reamostragem (Resampling): A reamostragem padrão de Monte Carlo Sequencial (SMC) é empregada para evitar a degenerescência de pesos.

O framework recupera o direcionamento de recompensa pontual padrão como um caso especial onde $R(\mu) = \int r(x) d\mu(x)$.

Principais Contribuições

1. Formalização do Direcionamento Distribuicional: Os autores formalizam o direcionamento ao nível de distribuição como o alvo de uma medida inclinada onde a inclinação depende da própria medida, estabelecendo um análogo distribuicional para alvos de recompensa inclinados pontuais.
2. Derivação da Dinâmica de McKean-Vlasov Ponderada: Eles derivam um sistema dinâmico fundamentado com reponderação Feynman-Kac corretiva que garante que as partículas interagentes ponderadas rastreiem o caminho inclinado desejado $\mu^*_t$.
3. Recuperação Teórica e Fundamentação: O framework recupera o direcionamento pontual padrão como um caso especial e fornece uma base teórica para métodos heurísticos de direcionamento em nível de lote existentes (ex: amostragem de diversidade, orientação experimental).
4. Garantias de Convergência: O artigo prova que o sistema de partículas finitas converge para a dinâmica idealizada conforme o número de partículas $N \to \infty$.

Resultados Experimentais

Os autores validam o framework tanto em ambientes de baixa dimensão quanto de alta dimensão:

• Baixa Dimensão (Mistura 1D):

  • Tarefa: Inclinar uma mistura Gaussiana bimodal para uma distribuição alvo com diferentes pesos de modo usando uma recompensa de MMD (Maximum Mean Discrepancy) ao quadrado.
  • Descobertas: O direcionamento de gradiente padrão empurra as partículas em direção ao alvo, mas falha em corresponder precisamente à forma da distribuição inclinada $\mu^*$. Em contraste, o método de campo médio visa $\mu^*$ com precisão, minimizando o gap de otimalidade $J(\mu) - J(\mu^*)$ quando a força de direcionamento $\lambda$ corresponde ao parâmetro do objetivo. Isso se mantém através de vários cronogramas de ruído (noise schedules).

• Alta Dimensão (Conformação de Proteína):

  • Tarefa 1 (HIV-1 Protease): Direcionar ensembles gerados pelo Boltz-2 para corresponder às distribuições de distância de resíduos observadas experimentalmente (dados DEER). O modelo base estava enviesado para o estado "fechado"; o objetivo era recuperar as proporções do estado "semi-aberto".
    • Resultado: O direcionamento de campo médio correspondeu precisamente à distribuição de distância de resíduos empírica, enquanto o direcionamento apenas por gradiente produziu uma distribuição com uma forma ligeiramente diferente e picos mais fortes.
  • Tarefa 2 (Adenilato Quinase): Direcionar para uma distribuição 2D de ângulos de ligação derivada de simulações de Dinâmica Molecular.
    • Resultado: O direcionamento de campo médio produziu um deslocamento mais forte em direção à distribuição alvo comparado aos métodos de apenas gradiente.
  • Tarefa 3 (Cristalografia de Raios-X): Direcionar gerações do Protenix para ajustar um mapa de densidade eletrônica (estrutura 4OLE) que exibe conformações heterogêneas.
    • Resultado: Para a mesma força de direcionamento, o direcionamento de campo médio alcançou um maior alinhamento de cosseno com a densidade alvo e colocou mais amostras na conformação alternativa (Conformação B) comparado ao direcionamento apenas por gradiente, sugerindo que métodos de apenas gradiente subdirecionam em relação ao verdadeiro alvo do problema inverso.

Significância e Alegações

O artigo afirma preencher a lacuna entre o controle pontual e o distribuicional ao fornecer uma fundamentação teórica sólida para o segundo.

• Corretude: A principal alegação é que o procedimento proposto alvos exatamente a distribuição inclinada pela recompensa prescrita $\mu^*$ no limite de infinitas partículas, corrigindo as discrepâncias inerentes ao direcionamento de gradiente heurístico.
• Interpretabilidade: O framework esclarece que os métodos de direcionamento de nível de lote existentes estão efetivamente realizando passos de gradiente, mas carecem da reponderação corretiva necessária para amostrar da verdadeira medida posterior/inclinada.
• Praticidade: Embora o método introduza um custo computacional adicional (resolver para $\dot{\Psi}_t$), os autores observam que isso é frequentemente negligenciável em comparação com as avaliações de rede em configurações de alta dimensão (ex: proteínas) e que o solver de ponto fixo é eficiente.
• Limitações: Os autores reconhecem que a dependência de reamostragem pode ser uma limitação para modalidades como imagens (potencialmente levando a amostras duplicadas) e que o método requer funções de recompensa diferenciáveis. Eles sugerem que a combinação com técnicas de redução de variância (ex: solvers da equação de Poisson) poderia mitigar os problemas de reamostragem.

O trabalho posiciona-se não como um substituto para todos os métodos existentes, mas como um embasamento rigoroso que explica por que certas heurísticas funcionam e fornece um mecanismo para garantir que elas visem a distribuição correta.