Resumo Técnico: Emyx – Geração de Proteínas All-Atom Rápida e Eficiente

Declaração do Problema

O design computacional de enzimas requer a geração de scaffolds proteicos que posicionem precisamente resíduos catalíticos e ligantes. Esta tarefa exige tanto precisão geométrica quanto diversidade estrutural. Os modelos generativos all-atom atuais, como o RFdiffusion3 e o Proteína-Complexa, herdam arquiteturas complexas da predição de estrutura de proteínas (ex: representações de pares densos, pilhas de embeddings pesadas). Estas arquiteturas resultam em altos custos de treinamento, diversidade de amostras limitada e, frequentemente, exigem etapas caras de folding inverso para gerar sequências de alta qualidade. Além disso, as métricas de avaliação existentes (ex: RMSD de autossuperfície de átomos pesados) podem superestimar o sucesso ao falhar em verificar a recuperação da dobra global, levando a designs que satisfazem restrições de motivos locais, mas falham em formar proteínas estáveis e dobráveis.

Metodologia

Arquitetura do Modelo

O Emyx é um modelo de correspondência de fluxo condicional (conditional flow matching) de 140M de parâmetros, projetado para a geração de proteínas all-atom. Ao contrário dos predecessores que dependem de atualizações de pares densos e camadas de embedding específicas de domínio, o Emyx concentra sua capacidade dentro de blocos transformer padrão.

• Representação: O modelo utiliza uma representação unificada de dois níveis. No nível atômico, resíduos e ligantes são mapeados para uma matriz "Rep14" de tamanho fixo (14 slots de átomos), onde posições não utilizadas são preenchidas com "átomos fantasmas" que herdam coordenadas de backbone. No nível de token, resíduos e HETATMs constituem tokens únicos.
• Conectividade: Em vez de representações de pares densos, o Emyx utiliza um grafo de conectividade esparso com um orçamento de arestas fixo. As arestas são construídas com base em vizinhos de sequência, ligações químicas, k-NN de ligantes e conectividade de motivo.
• Design do Transformer: A arquitetura emprega um pipeline "patchify–process–unpatchify":
  1. Codificador de Átomos: Blocos transformer leves codificam o detalhe atômico local dentro da representação Rep14.
  2. Downcast: Cross-attention com portão (gated cross-attention) agrupa as representações de átomos em embeddings de nível de token.
  3. Tronco de Token: Um tronco transformer profundo (18 camadas) processa o contexto estrutural de longo alcance usando auto-atenção esparsa com vieses de pares aprendidos.
  4. Upcast: Um segundo cross-attention transmite o contexto do token de volta para os átomos.
  5. Decodificador de Átomos: Blocos transformer finais refinam o detalhe atômico local antes do cabeçalho de velocidade de saída.
• Modificações: Os blocos DiT (Diffusion Transformer) padrão são modificados com projeções de gargalo para condicionamento, portão sigmoide para limitar atualizações residuais e profundidade estocástica (drop-path) para regularização.

Treinamento e Objetivo

O Emyx é treinado usando correspondência de fluxo condicional (conditional flow matching) sobre coordenadas cartesianas.

• Objetivo: O modelo aprende um campo de velocidade dependente do tempo $v_\theta(x_t, t, c, G)$ para transportar amostras de uma distribuição de ruído para a distribuição de dados. A perda de treinamento combina uma regressão $\ell_2$ nas velocidades para átomos não pertencentes ao motivo com uma perda auxiliar de teste de diferença de distância local diferenciável (lDDT) para impor qualidade estrutural.
• Condicionamento: O modelo é condicionado em restrições geométricas esparsas (geometria do motivo, posições de ligantes) e características globais (composição de estrutura secundária, raio de giração). Estas características são estocasticamente mascaradas durante o treinamento para permitir a geração tanto condicional quanto incondicional.
• Eficiência: O treinamento utiliza um interpolante linear entre ruído e dados, o que produz gradientes de menor variância e convergência mais rápida em comparação com cronogramas de difusão de preservação de variância. O modelo foi treinado em dados experimentais do PDB em aproximadamente 3,55 dias em 8 GPUs NVIDIA H200 (681,6 horas de GPU).

Amostragem e Reparametrização

Uma contribuição metodológica chave é a derivação de uma reparametrização exata do interpolante de correspondência de fluxo para a estrutura de nível de ruído EDM (Elucidating the Design Space of Diffusion-based Generative Models).

• Este mapeamento permite que o modelo de correspondência de fluxo utilize o cronograma de ruído Karras e mecanismos de churn estocástico sem necessidade de retreinamento ou mudanças arquiteturais.
• A reparametrização une a eficiência de treinamento da correspondência de fluxo com os métodos de amostragem de estado da arte projetados para modelos de difusão, melhorando significativamente as taxas de sucesso sobre a integração Euler–Maruyama de SDE padrão.

Principais Contribuições

1. Gerador All-Atom Eficiente: O Emyx alcança desempenho de estado da arte com uma contagem de parâmetros significativamente menor (140M) e custo de treinamento (4× menor que o RFdiffusion3) ao dispensar arquiteturas herdadas de predição dispendiosas.
2. Protocolo de Avaliação Estrito: Os autores propõem um protocolo "sc-RMSD estrito" que exige tanto a recuperação da dobra global (backbone RMSD < 2.0 Å) quanto a precisão da geometria catalítica local (motivo tip-atom RMSD < 1.5 Å). Eles demonstram que métricas anteriores (heavy-atom sc-RMSD) superestimam as taxas de sucesso em cerca de 2× porque falham em penalizar designs que recuperam o sítio ativo, mas não a dobra global.
3. Reparametrização EDM: O artigo introduz um método agnóstico ao modelo para converter modelos de velocidade de correspondência de fluxo na estrutura EDM, permitindo estratégias de amostragem superiores (cronograma Karras, churn estocástico) sem retreinamento.
4. Análise Espectral: Os autores fornecem uma análise espectral revelando um fundamental "gap gerador–preditor". Preditores inicializados com sinais ricos de MSA/template exibem maior rank efetivo e menor overtraining, enquanto geradores inicializados a partir de coordenadas com ruído operam em um regime de rank inferior. O Emyx demonstra a menor fração de overtraining entre os geradores, atribuída aos seus gargalos arquiteturais e profundidade estocástica.

Resultados

Avaliado contra o benchmark AME (Atom-level Motif Enzyme - 41 sítios ativos catalíticos):

• Taxa de Sucesso: Sob o protocolo sc-RMSD estrito, o Emyx alcança uma taxa de sucesso de 13,4% (39/41 alvos resolvidos), superando o Proteína-Complexa (8,8%) e o RFdiffusion3 (6,7%).
• Novidade Estrutural: O Emyx produz os scaffolds mais novos estruturalmente, com um TM-score mediano de 0,475 para o hit mais próximo no PDB (menor é mais novo), comparado a 0,499 para o Proteína-Complexa e 0,512 para o RFdiffusion3.
• Diversidade: O Emyx gera o maior número de clusters estruturais únicos (441) entre os designs bem-sucedidos.
• Validade Geométrica: O Emyx passa em verificações de validade geométrica (comprimentos/ângulos de ligação) em 82,2% dos designs, superando significativamente o RFdiffusion3 (0,8%) e o Proteína-Complexa (49,9%). Os autores observam que os baixos scores de geometria do RFdiffusion3 são provavelmente devidos a artefatos de pós-processamento onde as coordenadas do motivo de entrada sobrescrevem as geradas, criando ligações não físicas.
• Eficiência: O Emyx treina em ~682 horas de GPU (4× mais rápido que o RFdiffusion3) e é >2× mais rápido na inferência em uma única GPU A10G em todos os comprimentos de cadeia.

Significância e Alegações

O artigo afirma que a simplicidade arquitetural é benéfica para geradores de proteínas. Diferente dos preditores, que requerem sinais coevolutivos ricos (MSA, templates) para inicializar representações de pares complexas, os geradores condicionam-se em restrições geométricas esparsas. Portanto, as arquiteturas pesadas herdadas da predição de estrutura são desnecessárias e prejudiciais à eficiência e diversidade.

O Emyx demonstra que concentrar a capacidade em blocos transformer padrão, combinado com treinamento de correspondência de fluxo eficiente e amostragem do tipo EDM, pode alcançar um desempenho superior em tarefas de design de enzimas com significativamente menos recursos computacionais. O trabalho também destaca a importância de protocolos de avaliação rigorosos (sc-RMSD estrito) para distinguir entre designs localmente precisos, mas globalmente instáveis, e scaffolds de enzimas verdadeiramente viáveis. Finalmente, a reparametrização EDM proposta oferece um caminho para que outros modelos de correspondência de fluxo adotem técnicas de amostragem avançadas sem necessidade de reformulação arquitetural.