Constrained Flow Optimization via Sequential Fine Tuning for Molecular Design

Resumo Técnico: Otimização de Fluxo com Restrições via Ajuste Fino Sequencial para Design Molecular

1. Declaração do Problema

O artigo aborda o desafio de adaptar modelos de fundação generativos, especificamente modelos de fluxo e difusão, para aplicações de descoberta científica do mundo real, como o design molecular e a engenharia de proteínas. Embora esses modelos se destaquem em capturar distribuições de dados complexas, a amostragem padrão é insuficiente para tarefas de descoberta, onde o objetivo é gerar candidatos que maximizem recompensas específicas (ex: afinidade de ligação, momento de dipolo) enquanto satisfazem rigorosamente restrições específicas do domínio (ex: acessibilidade sintética, limites de toxicidade ou estabilidade de energia).

Métodos de ajuste fino existentes, frequentemente baseados em aprendizado por reforço (RL) ou teoria de controle, tipicamente otimizam funções de recompensa sob regularização KL para evitar o desvio distributivo. No entanto, eles têm dificuldade em lidar com restrições rígidas de forma confiável. Abordagens ingênuas, como incorporar restrições como termos ponderados na função objetivo, são consideradas pouco confiáveis porque a ponderação ideal entre recompensas e restrições varia entre as tarefas e as fases de treinamento, exigindo um ajuste ineficiente de tentativa e erro. Além disso, esses métodos não podem garantir que as amostras geradas satisfaçam restrições explícitas, frequentemente produzindo amostras de alta recompensa que violam propriedades físicas ou químicas críticas.

A questão central levantada é: Como um modelo de fluxo ou difusão pré-treinado pode ser ajustado para negociar de forma confiável e previsível o equilíbrio entre a otimização da recompensa e a satisfação de restrições?

2. Metodologia

Os autores propõem um framework formal e uma solução algorítmica para este problema.

Framework de Otimização Generativa com Restrições:
O problema é formulado como encontrar uma política $\pi^*$ que maximiza a recompensa esperada $r(x)$ enquanto satisfaz uma restrição esperada $c(x) \leq B$, sujeito à regularização KL em relação a um modelo pré-treinado $\pi_{pre}$. O objetivo é:
$$\max_{\pi} \mathbb{E}_{x \sim p^{\pi}_1}[r(x)] - \alpha D_{KL}(p^{\pi}_1 || p^{pre}_1) \quad \text{sujeito a} \quad \mathbb{E}_{x \sim p^{\pi}_1}[c(x)] \leq B$$
Esta formulação generaliza tanto a otimização guiada por recompensa quanto a geração com restrições (onde $r$ é constante).

Otimização de Fluxo com Restrições (CFO):
Para resolver isso, os autores introduzem a Otimização de Fluxo com Restrições (CFO), um algoritmo baseado no esquema de Lagrangiano Aumentado (AL). O CFO transforma o problema com restrições em uma sequência de subproblemas de ajuste fino não restringidos e com regularização de entropia. O algoritmo opera iterativamente com as seguintes etapas:

1. Construção do Objetivo Aumentado: Em cada iteração $k$, uma função objetivo aumentada $f_k(x)$ é construída. Esta função combina a recompensa $r(x)$ com um termo de penalidade derivado da violação da restrição, ajustado por um multiplicador de Lagrange $\lambda_k$ e um parâmetro de penalidade $\rho_k$:
   $$f_k(x) = r(x) - \frac{\rho_k}{2} \left[ \max\left(0, c(x) - B - \frac{\lambda_k}{\rho_k}\right) \right]^2$$
2. Etapa de Ajuste Fino: Um solver de ajuste fino padrão (ex: Adjoint Matching, DiffusionNFT) é usado para otimizar a política $\pi_k$ contra o objetivo aumentado $f_k$, regularizado pela divergência KL em relação ao modelo pré-treinado.
3. Atualizações de Variáveis Duais:
   • Multiplicador de Lagrange ($\lambda$): Atualizado com base no gap de restrição empírico $G_k = \mathbb{E}[c(x)] - B$. Se a restrição for violada ($G_k > 0$), $\lambda$ é diminuído para fortalecer a penalidade; se for satisfeita, ela é relaxada.
   • Parâmetro de Penalidade ($\rho$): Atualizado com base em uma estatística de contração $V_k$. Se o gap da restrição não contrair suficientemente, $\rho$ é aumentado para impor uma maior viabilidade.

Esta abordagem sequencial permite que o algoritmo equilibre automaticamente a maximização da recompensa e a satisfação da restrição sem ajuste manual de pesos.

3. Principais Contribuições

O artigo apresenta quatro contribuições primárias:

• Framework Formal: Formula a Otimização Generativa com Restrições via ajuste fino de fluxo, capturando a tarefa prática de adaptação guiada por recompensa sob quaisquer restrições.
• Algoritmo CFO: Introduz o CFO, um método baseado em Lagrangiano Aumentado que resolve o problema com restrições via ajuste fino sequencial de forma comprovável.
• Garantias Teóricas: Sob suposições realistas de um solver aproximado e condições de oráculo específicas, os autores fornecem garantias de convergência. Eles provam que o CFO retorna uma política que minimiza a inviabilidade e, sob condições mais fortes (erro de aproximação nulo), alcança a otimalidade global.
• Validação Empírica: O método é avaliado tanto em configurações ilustrativas de baixa dimensão quanto em tarefas de design molecular de alta dimensão, demonstrando sua capacidade de equilibrar recompensas e restrições de forma eficaz.

4. Resultados Experimentais

Os autores avaliam o CFO em duas frentes principais:

Configurações Ilustrativas de Baixa Dimensão:

• Otimização com Restrições: Usando uma mistura de Gaussianas, o CFO conseguiu direcionar o modelo para regiões de alta recompensa (próximo a um ponto alvo) enquanto permanecia dentro de regiões de restrição válidas (triângulos). Comparado ao Adjoint Matching padrão (que ignora restrições), o CFO reduziu significativamente as violações de restrição (de 2,47 para 0,12) mantendo uma alta recompensa.
• Geração com Restrições: Ao definir a recompensa como uma constante, o CFO recuperou capacidades de geração com restrições. Variar o limite da restrição $B$ permitiu o relaxamento controlado da restrição, demonstrando a flexibilidade do método.
• Agnosticismo de Solver: O CFO apresentou desempenho consistente quando pareado tanto com solvers de primeira ordem (Adjoint Matching) quanto com solvers livres de gradiente (DiffusionNFT), confirmando sua modularidade.

Tarefas de Design Molecular:

• Maximização de Dipolo com Restrição de Energia: Usando o modelo FlowMol no conjunto de dados GEOM Drugs, o CFO otimizou o momento de dipolo (recompensa) enquanto impunha um limite superior na energia total (restrição, proxy para estabilidade).
  • O CFO aumentou o momento de dipolo de 6,55 D (pré-treinado) para 8,39 D enquanto satisfazia a restrição de energia (energia esperada -82,28 Ha vs. limite -80 Ha).
  • Em contraste, o ajuste fino sem restrições (Adjoint Matching) alcançou recompensas semelhantes, mas violou a restrição de energia (energia esperada -78,31 Ha).
• Comparação com Baselines: O CFO superou as baselines de penalidade de peso fixo ajustadas manualmente. O ajuste manual exigiu uma busca exaustiva sobre 18 ordens de magnitude de pesos para encontrar uma única configuração que igualasse o desempenho do CFO, destacando a ineficiência do ajuste manual.
• Validade e Eficiência: Embora o ajuste fino tenha reduzido a porcentagem de moléculas quimicamente válidas (um desafio conhecido no design generativo), o CFO manteve uma validade maior do que as baselines sem restrições em níveis de recompensa comparáveis. O CFO também demonstrou eficiência computacional, exigindo apenas uma passagem de ajuste fino em comparação com os altos custos de amostragem de métodos de restrição em tempo de inferência, como o DiffOpt.

5. Significância e Alegações

O artigo afirma que o CFO fornece um mecanismo confiável e previsível para a otimização generativa com restrições. Sua significância reside em:

• Troca Automatizada: Elimina a necessidade de ajuste manual e ineficiente de tentativa e erro dos pesos das restrições, que é um grande gargalo nos fluxos de trabalho atuais de design generativo.
• Rigor Teórico: Oferece garantias de convergência tanto para a satisfação de restrições quanto para a otimalidade da recompensa, fundamentando o método na teoria de otimização estabelecida (Lagrangiano Aumentado).
• Utilidade Prática: Ao aplicar com sucesso a tarefas de design molecular de alta dimensão, os autores demonstram que o método escala para problemas científicos do mundo real onde as restrições (como estabilidade química) não são negociáveis.
• Modularidade: O framework é agnóstico ao solver de ajuste fino subjacente, permitindo que seja integrado com várias estratégias de otimização (baseadas em gradiente ou livres de gradiente) e aplicável a diferentes modalidades de dados (contínuos, discretos ou mistos).

Os autores concluem que o CFO permite a adaptação de modelos de fluxo pré-treinados para maximizar recompensas esperadas enquanto impõe restrições arbitrárias, oferecendo uma ferramenta robusta para aplicações de descoberta científica.