Resumo Técnico: Yeti – Um Tokenizador Compacto de Estrutura Proteica para Reconstrução e Geração Multimodal

1. Declaração do Problema

A integração de sequências proteicas, estruturas 3D e anotações funcionais em uma representação multimodal unificada é crítica para o avanço da biologia computacional. No entanto, as abordagens atuais enfrentam três desafios principais:

1. Heterogeneidade de Dados: Dados biológicos abrangem modalidades diversas (restrições biofísicas, texto, imagens, dados volumétricos) com propriedades estatísticas distintas.
2. Incompletude e Ruído: Conjuntos de dados do mundo real frequentemente sofrem com modalidades ausentes, medições ruidosas e cobertura desigual no espaço estrutura-função.
3. Escalabilidade e Acoplamento: A maioria dos modelos opera em uma única escala ou modalidade, carecendo de mecanismos para transferir informações entre níveis (atômico a celular) ou para lidar com processos acoplados.

Especificamente, embora modelos generativos recentes tratem proteínas como objetos 3D usando espaços de coordenadas contínuos (por exemplo, difusão SE(3)), há uma alternativa convincente: codificar estruturas atômicas 3D como tokens discretos. Isso permite treinamento escalável e agnóstico à modalidade, semelhante a modelos de linguagem. No entanto, os tokenizadores de estrutura proteica existentes (por exemplo, ESM3, DPLM-2, Kanzi) frequentemente priorizam a precisão de reconstrução em detrimento das capacidades generativas. Há um desalinhamento onde alta fidelidade de reconstrução não garante diversidade de tokens de alta qualidade ou suporte à co-geração incondicional de sequência e estrutura. Além disso, muitos modelos de última geração dependem de contagens massivas de parâmetros ou paradigmas de treinamento em duas etapas onde os codificadores são congelados, limitando potencialmente a otimização ponta a ponta.

2. Metodologia

2.1 Arquitetura: Yeti

Yeti (Yielding Encoded Tokens for Intermodality) é um tokenizador compacto de estrutura proteica em etapa única, projetado para mapear coordenadas atômicas 3D contínuas para tokens discretos.

• Entrada: Estruturas proteicas representadas como coordenadas $C_\alpha$ centradas na média ($x \in \mathbb{R}^{L \times 3}$).
• Codificador: Um codificador baseado em Transformer que utiliza atenção multi-cabeça com codificação posicional rotativa (RoPE) mapeia entradas para embeddings latentes contínuos $Z \in \mathbb{R}^{L \times D}$.
• Quantização: Os embeddings são projetados em um espaço de quantização e processados por um Quantizador Livre de Consulta (LFQ). O espaço LFQ é um produto cartesiano de sub-livros de códigos binários ($C_i = \{-1, +1\}$). A quantização é realizada via uma operação de sinal, mapeando vetores para índices discretos sem a necessidade de uma tabela de consulta, resultando em um tamanho de vocabulário $K = 2^D$.
• Decodificador: Uma arquitetura StripedHyena (um modelo híbrido de processamento de sinais) atua como o decodificador. Diferentemente de decodificadores tradicionais que dependem de componentes invariantes a SE(3) como Atenção em Pontos Invariantes (IPA), o decodificador StripedHyena aprende um campo de velocidade condicional para reconstruir estruturas a partir de coordenadas ruidosas.

2.2 Objetivo de Treinamento: Flow Matching

O Yeti é treinado ponta a ponta usando um objetivo de Flow Matching, que interpola entre uma distribuição de ruído gaussiano ($p_0$) e a distribuição alvo de estrutura proteica ($p_{data}$).

• Função de Perda: O modelo minimiza a diferença entre o campo de velocidade condicional previsto $v_\theta(x_t, t, c)$ e o campo vetorial alvo ($x_1 - x_0$), onde $x_1$ é a estrutura limpa e $x_0$ é ruído gaussiano.
• Regularização de Entropia: Para evitar o colapso do livro de códigos e garantir alta utilização, uma penalidade de entropia é adicionada à função de perda: $L = L_{FM} + \lambda L_{entropy}$.
• Inferência: A Orientação Livre de Classificador (CFG) é empregada durante a inferência. Ao zerar aleatoriamente embeddings de tokens durante o treinamento, o modelo aprende campos de velocidade condicionais e incondicionais, permitindo decodificação guiada para refinar a precisão estrutural.

2.3 Co-geração Multimodal

Para validar a utilidade generativa dos tokens Yeti, os autores treinaram um Modelo de Difusão Mascarada (MDM) do zero. Este modelo gera conjuntamente sequências de aminoácidos e tokens de estrutura Yeti sem qualquer inicialização pré-treinada (diferentemente do DPLM-2, que inicia a partir de um modelo de sequência, ou do ESM3, que usa pipelines em cascata).

3. Contribuições Principais

1. Tokenizador Compacto e Ponta a Ponta: Introdução do Yeti, um tokenizador em etapa única treinado ponta a ponta com flow matching, evitando o paradigma de duas etapas "congelar codificador" comum em trabalhos anteriores.
2. Alta Utilização do Livro de Códigos: O Yeti alcança superior utilização de livro de códigos e diversidade de tokens em comparação com modelos maiores (ESM3, DPLM-2) e outros métodos de quantização (FSQ, VQ-VAE), apesar de ter significativamente menos parâmetros.
3. Eficiência de Parâmetros: O Yeti alcança precisão de reconstrução competitiva com apenas 62,5M de parâmetros, aproximadamente 10x menos que o ESM3 (648M) e 2x menos que o DPLM-2 (118M).
4. Prova de Conceito de Co-geração: Demonstração de que um modelo multimodal compacto treinado do zero em tokens Yeti pode realizar co-geração incondicional de sequências e estruturas proteicas, alcançando resultados comparáveis a modelos 10x maiores.
5. Análise da Dinâmica de Dobramento: Identificação de um surgimento hierárquico de estrutura secundária durante a trajetória de decodificação por flow matching, caracterizado por compactação global precoce seguida por um fenômeno de "compromisso tardio" onde a topologia global se consolida nos passos finais.

4. Resultados

4.1 Análise do Livro de Códigos

• Utilização: Em um livro de códigos de 8.192 entradas, o Yeti alcança alta perplexidade e entropia, superando significativamente o ESM3 e o Kanzi. No conjunto de dados CATH, o Yeti utiliza ~66% de sua capacidade (Perplexidade = 5414), com uma diversidade intra-estrutura ($\bar{U}$) de 98,8%, quase saturando a métrica e indicando atribuições de tokens altamente distintas para resíduos.
• Comparação: O Yeti supera o DPLM-2 na utilização do livro de códigos, apesar de ambos usarem LFQ, atribuído ao paradigma de treinamento ponta a ponta versus a abordagem de codificador congelado do DPLM-2.

4.2 Precisão de Reconstrução

• Métricas: O Yeti alcança um TM-score de 0,96 no CAMEO, 0,97 no CASP14 e 0,95 no CASP15 e CATH.
• Eficiência: Embora o ESM3 mostre RMSD ligeiramente menor (devido à sua escala massiva), os TM-scores do Yeti são comparáveis ou superiores, indicando que ele recupera efetivamente a topologia global correta. O decodificador baseado em flow do Yeti produz amostras com variância em posições atômicas locais (penalizando o RMSD), mas topologia global consistente, uma característica benéfica para geração multi-conformacional.

4.3 Co-geração Incondicional

• Desempenho: Um modelo de 224M de parâmetros treinado do zero em 2,6M de proteínas (usando tokens Yeti) alcançou um TM-score de auto-consistência (scTM) de 0,70 e pLDDT de 76,12.
• Comparação: Este desempenho é comparável ao DPLM-2 (scTM 0,87) e significativamente melhor que o ESM3-Open (scTM 0,46), apesar do DPLM-2 depender de um modelo de sequência pré-treinado de 650M de parâmetros e o ESM3 usar um pipeline em cascata.
• Designabilidade: O ProteinMPNN recuperou 30,6% das sequências originais geradas a partir das estruturas geradas, confirmando a consistência geométrica e a designabilidade dos pares co-gerados.

4.4 Dinâmica da Trajetória Flow

A análise da trajetória de decodificação revelou um surgimento não monótono de características:

• Fase Precoce ($t < 0,6$): Compactação global rápida (diminuição do Raio de Giração).
• Fase Tardia ($t > 0,9$): Coalescência rápida de elementos de estrutura secundária ($\alpha$-hélices e $\beta$-folhas) e consolidação da topologia global.
  Isso sugere que o tokenizador desacopla efetivamente a densidade de granulação grossa de detalhes estruturais de alta frequência.

5. Significado e Alegações

O artigo afirma que o Yeti serve como uma base compacta e expressiva para o treinamento de modelos proteicos multimodais. Seu significado principal reside em demonstrar que:

1. Tokenização Discreta é Viável para Geração: A co-geração incondicional de alta qualidade de sequência e estrutura é possível sem pré-treinamento massivo ou pipelines em cascata, desde que o tokenizador seja expressivo e o objetivo de treinamento (flow matching) esteja alinhado com objetivos generativos.
2. Eficiência sobre Escala: Um modelo menor e cuidadosamente projetado (62,5M de parâmetros) pode superar ou igualar modelos muito maiores (600M+) em diversidade de tokens e reconstrução, desafiando a noção de que a escala é o único caminho para o desempenho em modelagem proteica.
3. Representação Unificada: A capacidade de treinar um modelo do zero em dados conjuntos de sequência-estrutura valida o potencial para uma representação molecular universal que captura o acoplamento intrínseco entre sequências de aminoácidos e estruturas 3D.

Os autores posicionam o Yeti não como uma solução final para todas as tarefas de design proteico, mas como uma prova de conceito que estabelece a viabilidade de sua estratégia de tokenização. Eles reconhecem limitações, como a dependência de estruturas de conformação única e a necessidade de maior escalabilidade para igualar o desempenho de modelos especializados e altamente ajustados como o La-Proteina. No entanto, eles afirmam que o Yeti fornece uma espinha dorsal escalável, eficiente e eficaz para futuros modelos proteicos multimodais capazes de co-gerar sequências e estruturas altamente plausíveis.