Resumo Técnico: Few-step Cofolding com All-Atom Flow Maps (DECAF)

Problema

A modelagem computacional precisa de complexos biomoleculares 3D, particularmente para o cofolding proteína-ligante, tornou-se um paradigma dominante no design de fármacos baseado em estrutura (SBDD). Embora modelos generativos de todos os átomos modernos (ex: AlphaFold 3, Boltz-1) possam prever estruturas com resolução atômica, eles dependem de rollouts de difusão iterativos caros. Gerar estruturas de alta fidelidade tipicamente requer centenas de avaliações de função neural (NFEs), criando um gargalo computacional significativo tanto para implantação quanto para técnicas de busca em tempo de inferência.

Além disso, gerar estruturas fisicamente válidas ao nível atômico é crítico; modelagens locais imprecisas levam a falhas catastróficas, como choques estéricos, comprimentos/ângulos de ligação incorretos e cadeias laterais tensionadas. Restrições físicas estritas frequentemente deslocam o fardo computacional para métodos em tempo de inferência, onde a busca guiada por recompensa (ex: Monte Carlo Sequencial, MCTS) é usada para refinar as amostras. No entanto, esses métodos de busca complicam o custo de inferência, pois exigem múltiplas consultas de recompensa sobre múltiplas partículas durante a trajetória. A questão central de pesquisa é se é possível treinar um modelo de cofolding de todos os átomos que gere amostras otimizadas por recompensa de forma eficiente usando apenas um pequeno número de NFEs no tempo de inferência.

Metodologia: O Framework DECAF

Os autores introduzem o DENOISER COFOLDING ALL-ATOM FLOWMAP (DECAF), um framework projetado para destilar modelos de difusão de todos os átomos de estado da arte em flow maps de poucos passos. A metodologia aborda desafios técnicos específicos na adaptação de modelos de difusão padrão para flow maps de modelagem de proteínas de todos os átomos.

1. Reparametrização de Flow Map em Espaço-$\sigma$

Modelos de difusão padrão frequentemente utilizam a parametrização EDM (Elucidating the Design Space), que condiciona em um único nível de ruído $\sigma$. Esta formulação não se estende claramente para dual-time flow maps, e a reparametrização de tempo não linear $\sigma(t)$ usada no EDM leva à instabilidade numérica ao computar parciais de tempo ($\partial \sigma / \partial t$) necessárias para os objetivos de flow map.

Para contornar isso, o DECAF define o flow map diretamente no espaço de ruído $\sigma$. Em vez de mapear o tempo $t$, ele mapeia entre níveis de ruído $\rho \to \sigma$ (onde $\rho > \sigma$).

• $\sigma$-Velocidade: O framework define uma $\sigma$-velocidade instantânea $v_\sigma(x_\sigma) = dx_\sigma/d\sigma$.
• Velocidade Média: O flow map é parametrizado como um deslocamento usando a velocidade $\sigma$ média entre dois níveis de ruído, permitindo que o modelo dê "saltos" ao longo da trajetória da ODE de fluxo de probabilidade diretamente no espaço de ruído.
• Estabilidade: Esta reparametrização elimina o fator problemático $\partial \sigma / \partial t$, permitindo o treinamento estável e a destilação direta de modelos de cofolding pré-treinados no estilo EDM.

2. Parametrização do Denoiser com Alinhamento SE(3)

Uma inovação crítica é o uso de uma formulação baseada em denoiser em vez de um preditor de velocidade direta.

• Predição de Endpoint: O modelo prevê o endpoint alvo $\hat{x}_{tgt}$ (a estrutura limpa) em vez da velocidade instantânea.
• Alinhamento Rígido SE(3): A função de perda incorpora uma minimização sobre o grupo SE(3) (usando o algoritmo de Kabsch) para alinhar a estrutura prevista $\hat{x}_{tgt}$ com a verdade fundamental (ou a predição do professor/teacher). Isso impõe suavemente a simetria SE(3), o que é crítico para treinar modelos estáveis e reduzir a variância.
• Função de Perda: O objetivo de treinamento combina o endpoint previsto contra a predição do professor após o alinhamento rígido, suportando tanto o treinamento diagonal (correspondência de score) quanto o off-diagonal (jump).

3. DECAF-SEARCH: Framework de Inferência

Para aproveitar o flow map para otimização de recompensa, os autores introduzem o DECAF-SEARCH.

• Mecanismo de Lookahead: Diferente dos modelos de difusão padrão onde a estimativa de recompensa exige denoising completo, o flow map do DECAF fornece uma predição de "olhar adiante" (lookahead) do endpoint limpo $\hat{x}_0$ a partir de qualquer estado ruidoso intermediário $x_\sigma$. Isso permite uma estimativa de recompensa ($R(\hat{x}_0)$) de maior fidelidade sem completar a trajetória total.
• Algoritmo de Busca: O framework adapta métodos de condução (steering) em tempo de inferência (Feynman-Kac, MCTS) para o cenário de flow map. Ele utiliza um amostrador $\gamma$ estocástico para gerar partículas, avalia suas recompensas de lookahead e realiza ascensão de gradiente no espaço limpo ou reamostragem para conduzir a geração em direção a configurações de alta recompensa e fisicamente válidas.

Principais Contribuições

1. Primeiro Flow Map de Proteína de Todos os Átomos: O design do DECAF-Boltz, o primeiro flow map de todos os átomos que destila com sucesso um professor Boltz-1.
2. Inovações Metodológicas:
   • Uma nova reparametrização em espaço-$\sigma$ que permite o treinamento estável de flow maps para modelagem de todos os átomos.
   • Uma parametrização baseada em denoiser que permite o alinhamento rígido ponderado por SE(3), crucial para estabilidade e precisão.
3. Busca em Tempo de Inferência: A introdução do DECAF-SEARCH, um framework que utiliza o lookahead do flow map para melhorar a estimativa de recompensa e o alinhamento, permitindo uma busca eficiente com menos NFEs.
4. Validação Empírica: Demonstração de que o DECAF supera os baselines baseados em difusão em regimes de baixo NFE e iguala o desempenho de orçamento total com significativamente menos computação.

Resultados Experimentais

Os autores avaliaram o DECAF nos benchmarks Runs N' Poses (RnP) e PoseBusters, comparando-o com o Boltz-1x (uma variante de potencial físico do Boltz-1) e outros modelos de fronteira como AlphaFold 3 e Pearl.

Desempenho em Orçamentos de Baixo NFE

• Benchmark RnP: Em orçamentos estritos de NFE (ex: 20–50 NFEs), o DECAF-SEARCH supera significativamente o Boltz-1x (ajustado) em todas as métricas (PB-Valid, Taxa de Sucesso, lDDT-PLI). Por exemplo, com 20 NFEs, o DECAF atinge um escore PB-Valid de 83,5% comparado a 75,9% do Boltz-1x.
• Generalização: O ganho de desempenho mantém-se em todos os quartis de similaridade de bolsa (pocket-similarity), incluindo alvos fora de distribuição, confirmando a robustez do modelo.
• Eficiência: O DECAF iguala o desempenho do Boltz-1x em 600 NFEs usando apenas 30 NFEs (uma redução de 20x em computação) no PoseBusters.

Destilação de Modelos de Estado da Arte

• DECAF-Pearl: O framework foi aplicado para destilar o modelo de cofolding de estado da arte Pearl. O DECAF-Pearl iguala a taxa de sucesso do professor enquanto usa 5x menos NFEs (40 NFEs vs. simulação completa).
• Comparação: O DECAF-Pearl supera outros baselines potentes (incluindo a simulação completa do Pearl) em métricas de validade física e taxa de sucesso em orçamentos de baixo NFE.

Fronteira de Pareto

A análise da fronteira de computação ótima mostra que o DECAF-SEARCH domina o Boltz-1x em todos os orçamentos de computação de inferência. A estratégia de busca ideal varia conforme o orçamento:

• Baixo NFE ($\le$ 30): SMC baseado em partículas (FK-steering) é o ideal.
• NFE Moderado (100–250): Estimativa de gradiente de Monte Carlo (MC-GRAD) é mais eficaz.
• Alto NFE ($\ge$ 142): A busca baseada em MCTS produz os melhores resultados.

Significância e Alegações

O artigo afirma que o DECAF representa um passo significativo para tornar a geração biomolecular de todos os átomos computacionalmente viável para aplicações de alto rendimento (high-throughput).

• Aceleração: Ao comprimir trajetórias de difusão em flow maps de poucos passos, o DECAF reduz a latência de inferência em ordens de magnitude (reduções de 5x a 20x em NFEs) sem sacrificar a precisão estrutural ou a valididade física.
• Otimização de Recompensa: O lookahead do flow map permite uma busca guiada por recompensa mais eficiente e precisa, abordando o "imposto de inferência" associado a múltiplas consultas de recompensa em modelos de difusão padrão.
• Impacto Prático: A aceleração permite novos fluxos de trabalho, como o cofolding de bibliotecas inteiras de ligantes contra um alvo para triagem virtual e a geração de grandes conjuntos de dados sintéticos para treinar modelos downstream, o que era anteriormente impraticável devido ao custo das simulações de difusão completas.
• Modéstia: Os autores reconhecem que os modos de falha (ex: ligações sp2 não planas) são herdados dos sinais de recompensa dos modelos professores (Boltz-1x) e não do procedimento de busca em si. Eles também observam que a relação entre o orçamento de NFE e a taxa de sucesso não é monotônica, sugerindo trabalhos futuros em busca adaptativa e treinamento conjunto com recompensas específicas de tarefa.

O código é disponibilizado em https://github.com/genesistherapeutics/decaf.