Understanding protein function with a multimodal retrieval-augmented foundation model

O artigo apresenta o PoET-2, um modelo fundamental de proteínas multimodal e aumentado por recuperação que combina aprendizado de contexto evolutivo e condicionamento estrutural para alcançar desempenho superior na previsão de efeitos de variantes e na aprendizagem de relações sequência-função, especialmente em conjuntos de dados pequenos.

O essencial

Imagine que as proteínas são como receitas de bolo extremamente complexas. Cada letra da receita é um ingrediente (um aminoácido), e a ordem em que você coloca esses ingredientes determina se o bolo vai ficar delicioso (funcional), seco (doente) ou se vai explodir na sua cara (tóxico).

Por décadas, os cientistas tentaram decifrar essas receitas olhando apenas para a lista de ingredientes. Mas a realidade é que o "bolo" (a proteína) dobra-se em uma forma 3D complexa, e pequenas mudanças na receita podem mudar tudo.

Aqui entra o PoET-2, o novo "chef de cozinha" criado pelos autores deste artigo. Vamos entender como ele funciona usando algumas analogias simples:

1. O Problema: O Chef que só sabe ler, não cozinhar

Antes do PoET-2, existiam outros "chefs" (modelos de IA) que liam milhões de receitas antigas. Eles eram bons em prever o que acontecia se você trocasse um ingrediente (uma mutação simples).

• O problema: Eles travavam se você tentasse adicionar um ingrediente novo no meio da receita (inserção) ou tirar um (deleção). Eles também não entendiam bem como ingredientes distantes na receita conversavam entre si (efeitos complexos). E, pior, eles precisavam de bibliotecas gigantescas (bilhões de parâmetros) para funcionar, o que os tornava lentos e caros.

2. A Solução: O PoET-2 é um "Chef com Memória de Família"

O PoET-2 é diferente. Ele não apenas lê a receita; ele tem três superpoderes:

A. O "Livro de Receitas da Família" (Aprendizado por Recuperação)

Imagine que você quer fazer um bolo de cenoura, mas não tem a receita exata. Em vez de tentar adivinhar do zero, você pega um livro de receitas de todas as famílias que fazem bolos de cenoura.

• O PoET-2 faz isso. Quando você pergunta sobre uma proteína, ele vai buscar instantaneamente outras proteínas "primas" (da mesma família evolutiva) que já existem na natureza. Ele olha para elas para entender as regras não escritas daquela família.
• A mágica: Isso permite que ele aprenda com poucos exemplos. Ele não precisa ser um gênio com uma biblioteca de 100 bilhões de livros; ele é inteligente porque sabe onde procurar a informação certa na hora.

B. O "Olho de Raio-X" (Multimodalidade: Sequência + Estrutura)

A maioria dos chefs antigos só lia a lista de ingredientes (sequência). O PoET-2, além de ler, consegue "ver" a forma 3D do bolo.

• Ele entende que se você colocar dois ingredientes que se atraem magneticamente, eles vão ficar um em cima do outro, independentemente de estarem longe na lista.
• Ele pode receber um esboço da forma do bolo e dizer: "Se você mudar este ingrediente aqui, o bolo vai desmoronar". Isso é crucial para prever se uma mutação vai causar uma doença.

C. O "Duplo Modo" (Gerador e Analista)

O PoET-2 tem dois chapéus:

1. O Criador (Modo Causal): Ele pode escrever uma receita do zero, ingrediente por ingrediente, criando novas proteínas que nunca existiram.
2. O Analista (Modo Mascarado): Ele pode olhar para uma receita incompleta (com buracos) e dizer o que falta, entendendo o contexto global. Isso é ótimo para criar mapas de como a proteína funciona.

3. O Que Ele Conseguiu Fazer? (Os Resultados)

• Prever o Impossível: Enquanto outros modelos ficavam confusos com "buracos" na receita (inserções e deleções), o PoET-2 acertou em cheio. Ele consegue prever o efeito de mutações complexas que mudam o tamanho da proteína.
• Medicina de Precisão: Ele é muito bom em dizer se uma mutação genética em humanos é "inofensiva" ou "perigosa" (patogênica), ajudando a diagnosticar doenças.
• Eficiência: Ele é pequeno e rápido (apenas 182 milhões de parâmetros). Enquanto outros modelos precisam de supercomputadores gigantescos, o PoET-2 roda em computadores mais comuns, tornando a engenharia de proteínas acessível.
• Aprendizado Rápido: Em testes onde só havia poucos dados de experimentos (como em laboratórios pequenos), o PoET-2 aprendeu muito mais rápido e com mais precisão do que os gigantes anteriores.

Resumo da Ópera

O PoET-2 é como um chef de cozinha que não apenas memorizou milhões de receitas, mas também aprendeu a ler a linguagem dos ingredientes e a ver a forma do bolo antes mesmo de assá-lo.

Ele usa a sabedoria da "família" (outras proteínas similares) para não cometer erros, consegue lidar com receitas que mudam de tamanho e é tão eficiente que cabe na mochila de um pesquisador, não exigindo um data center inteiro para funcionar. Isso abre portas para criar novos remédios, enzimas industriais e entender doenças com uma velocidade e precisão que nunca tivemos antes.

Resumo técnico

Resumo Técnico: PoET-2

1. O Problema

Os Modelos de Linguagem de Proteínas (PLMs) têm avançado significativamente na previsão de estruturas e na compreensão de sequências naturais. No entanto, a literatura atual aponta três desafios críticos que limitam sua aplicação em engenharia de proteínas e medicina:

• Limitação em Mutações Complexas: A maioria dos PLMs baseia-se em modelos de linguagem mascarada (MLM) que são limitados a prever efeitos de substituições únicas. Eles falham em prever o impacto de inserções e deleções (indels) e de efeitos epistáticos (interações entre múltiplas mutações).
• Ineficiência em Cenários Supervisionados: Embora os PLMs funcionem bem em cenários "zero-shot" (sem dados experimentais), eles frequentemente exigem grandes quantidades de dados mutagênicos para aprender relações sequência-função em cenários supervisionados, especialmente para posições não observadas durante o treinamento.
• Custo e Generalização: O aumento da capacidade do modelo (escalando para bilhões de parâmetros) tem melhorado a previsão de estrutura, mas muitas vezes não melhora (ou até piora) a modelagem de aptidão (fitness) e a previsão de função, além de aumentar drasticamente os custos computacionais e o risco de memorização.

2. Metodologia: A Arquitetura PoET-2

Os autores propõem o PoET-2, um modelo fundamental de proteínas multimodal, de aumento por recuperação (retrieval-augmented) e com objetivos de treinamento duais.

• Arquitetura Híbrida (Encoder-Decoder):

  • Encoder: Utiliza um transformador hierárquico com duas etapas de atenção. É equivariante à ordem das proteínas no contexto (ou seja, a ordem em que as proteínas de referência são apresentadas não altera a saída). Isso permite o aprendizado "in-context" de restrições evolutivas específicas de uma família proteica sem a necessidade de modelos massivos.
  • Decodificadores Duais: O modelo possui dois decodificadores distintos:
    1. Decodificador Autoregressivo (CLM): Treinado com objetivo de modelagem de linguagem causal. Gera sequências e calcula log-verossimilhanças exatas, permitindo pontuação de variantes com indels e múltiplas mutações.
    2. Decodificador Bidirecional (MLM): Treinado com objetivo de modelagem de linguagem mascarada. Gera embeddings ricos e contextualizados para tarefas de representação e previsão supervisionada.

• Multimodalidade e Aumento por Recuperação:

  • O modelo processa simultaneamente sequências e estruturas (coordenadas de backbone atômico N, Cα, C e confiança pLDDT).
  • Utiliza um mecanismo de recuperação (retrieval) onde o "prompt" de entrada consiste em um conjunto de proteínas homólogas (contexto) e, opcionalmente, uma proteína de consulta (query) com restrições parciais de sequência ou estrutura.
  • O modelo aprende a gerar novas proteínas condicionadas às restrições evolutivas e estruturais inferidas a partir desse contexto.

• Viés de Atenção Baseado em Estrutura:

  • Incorpora um viés de atenção aprendido baseado nas distâncias Cα-Cα discretizadas entre resíduos, permitindo que o modelo entenda a proximidade 3D independentemente da distância linear na sequência.

• Eficiência:

  • O modelo possui apenas 182 milhões de parâmetros, tornando-o significativamente mais leve que concorrentes de grande escala (como ESM-2 ou ESM-3), mas com desempenho superior em várias tarefas.

3. Principais Contribuições

1. Previsão Zero-Shot de Indels e Mutações de Alta Ordem: O PoET-2 é o primeiro modelo a demonstrar desempenho de ponta na previsão de efeitos de inserções/deleções e mutações com 3 ou mais substituições, superando métodos anteriores em até 20%.
2. Eficiência de Dados em Aprendizado Supervisionado: Em cenários de poucos exemplos (few-shot), os embeddings do PoET-2 superam métodos anteriores (como Kermut) com drasticamente menos dados de treinamento, alcançando o mesmo desempenho com apenas 100-250 pontos de dados que outros modelos exigiriam milhares.
3. Arquitetura Equivariante e Multimodal: A combinação de aumento por recuperação com condicionamento estrutural e sequencial em uma arquitetura leve permite a generalização para famílias proteicas não vistas no treinamento.
4. Desempenho em Benchmarks Clínicos: Estabelece novos recordes (SOTA) na previsão de patogenicidade de variantes clínicas humanas, tanto para substituições quanto para indels.

4. Resultados Experimentais

Os resultados foram avaliados no benchmark ProteinGym, cobrindo dados de varredura de mutação profunda (DMS) e dados clínicos.

• Zero-Shot (Sem dados experimentais):

  • Indels: O PoET-2 alcançou um coeficiente de correlação de Spearman ($\rho$) de 0.566, superando o anterior melhor (PoET-1) em ~0.05 e modelos não-PoET em ~0.10.
  • Mutações Múltiplas: Para variantes com 3+ mutações, o PoET-2 superou o estado da arte (VenusREM) com ganhos de $\Delta\rho \approx 0.09$.
  • Variantes Clínicas: Melhorou a AUROC em 0.018 para indels e 0.008 para substituições em relação ao PoET-1.
  • Ensemble: A combinação simples de PoET-2 com o VenusREM resultou no melhor desempenho geral em todos os benchmarks.

• Supervisionado (Com dados experimentais):

  • Em tarefas de regressão de aptidão usando Gaussian Processes (GP), o PoET-2 superou consistentemente o Kermut e o ESM-2 em todas as divisões de validação (Randômica, Modulo e Contígua).
  • Eficiência de Dados: O PoET-2 treinado com apenas 100 pontos de dados igualou o desempenho do ESM C treinado com ~2600 pontos, demonstrando uma capacidade superior de aprendizado a partir de dados limitados.

• Papel da Estrutura:

  • A condicionamento estrutural melhorou significativamente a previsão zero-shot (especialmente para estabilidade), mas teve impacto limitado ou nulo na previsão supervisionada, sugerindo que os embeddings do modelo já codificam implicitamente informações estruturais críticas.

5. Significado e Impacto

O PoET-2 representa um avanço fundamental na biologia computacional ao demonstrar que modelos menores e mais eficientes podem superar modelos massivos quando combinados com estratégias de recuperação de contexto e aprendizado multimodal.

• Acessibilidade: Com apenas 182M de parâmetros, o modelo pode ser executado em hardware acessível (GPUs de consumo), democratizando o acesso a ferramentas de ponta de engenharia de proteínas.
• Aplicações Práticas: A capacidade de prever indels e interações epistáticas é crucial para o design de enzimas, terapias gênicas e compreensão de doenças genéticas complexas onde mutações raras e múltiplas são comuns.
• Paradigma de Treinamento: O trabalho sugere que a escalabilidade pura (mais parâmetros) não é a única solução; a integração inteligente de dados evolutivos (recuperação) e estruturais em arquiteturas especializadas é um caminho mais eficiente para o avanço da IA em biologia.

Em suma, o PoET-2 estabelece um novo estado da arte na previsão de efeitos de variantes de proteínas, oferecendo uma ferramenta robusta, eficiente e versátil para pesquisadores e engenheiros de proteínas.