Resumo Técnico: Tokenização de Proteínas Adaptativa

Declaração do Problema

As abordagens atuais para a tokenização da estrutura de proteínas baseiam-se na localidade, onde os tokens são criados através do agrupamento (pooling) de informações de resíduos espacialmente vizinhos ao longo de uma sequência proteica. Embora esses métodos alcancem uma reconstrução de alta fidelidade, os autores identificam duas limitações críticas:

1. Acúmulo de Erro na Geração: Modelos generativos baseados em tokenização local sofrem de propagação de erro. Um único token amostrado incorretamente pode levar a desvios significativos na estrutura prevista, fazendo com que abordagens discretas tenham desempenho inferior aos modelos de difusão contínua em métricas generativas.
2. Baixa Escalabilidade: O número de tokens escala linearmente com o tamanho da proteína. Isso torna o modelamento de complexos proteicos grandes computacionalmente caro e limita a capacidade de comprimir proteínas de forma eficaz.
3. Limitações de Representação: Tokenizadores locais frequentemente exigem operações de mean-pooling para gerar representações de tamanho fixo para tarefas subsequentes, o que pode descartar informações globais críticas e escalar mal com o tamanho da proteína.

Metodologia

Os autores propõem a Tokenização de Proteínas Adaptativa (APT), um método de tokenização global onde tokens sucessivos contribuem com níveis crescentes de detalhe para uma representação global, em vez de corresponderem a vizinhanças locais específicas.

Arquitetura

A APT é implementada como um autoencoder de difusão com um gargalo (bottleneck) discreto:

• Entrada: Uma sequência de coordenadas $C_\alpha$ brutas (normalizadas para massa zero) de comprimento $L$.
• Encoder: Um transformador de atenção bidirecional mapeia as posições brutas para uma sequência latente $c \in \mathbb{R}^{L \times d}$.
• Discretização: A sequência latente é discretizada usando Quantização Escalar Finita (FSQ) com níveis $(8, 5, 5, 5)$, resultando em um tamanho de código (codebook) efetivo de 1.000.
• Mecanismo de Adaptatividade: Para impor a adaptatividade, o modelo emprega um dropout aninhado (nested dropout). Durante o treinamento, um limite superior $U$ é amostrado uniformemente de $[1, \min(L, k_{max})]$. Os tokens além deste limite são descartados. Isso incentiva o modelo a colocar informações globais críticas (componentes de baixa frequência) nos primeiros tokens e detalhes de alta frequência nos tokens subsequentes.
• Decoder: Um decoder de difusão, treinado usando um objetivo de correspondência de fluxo (flow-matching objective), reconstrói as coordenadas atômicas condicionadas aos tokens discretos. O decoder utiliza equivariância estocástica, aprendendo simetrias em vez de impô-las explicitamente.
• Predição de Tamanho: O tamanho da proteína é regredido a partir do primeiro token usando uma perda de entropia cruzada, permitindo que o modelo desacople o tamanho da proteína do comprimento da sequência de condicionamento.

Treinamento e Inferência

• Treinamento: O modelo é treinado em $\approx 473.000$ previsões sintéticas do AlphaFold2 a partir do banco de dados AFDB agrupado por Foldseek. O treinamento envolve duas etapas: (1) treinamento do autoencoder com rotações aleatórias para aprender a equivariância estocástica, e (2) treinamento de um modelo autorregressivo (estilo GPT) sobre os tokens APT.
• Estratégias de Inferência:
  • Tail Dropout (Dropout de Cauda): Durante a geração, o modelo pode interromper a geração de tokens precocemente. Qualquer prefixo da sequência de tokens é um sinal de condicionamento válido. A cauda (detalhes de alta frequência) é delegada ao decoder de difusão.
  • Parada Baseada em Entropia: Os autores introduzem critérios de parada baseados na entropia do token (limite finito, entropia mínima baseada em spline) para determinar o ponto ideal para interromper a geração, equilibrando a fidelidade da representação e a capacidade generativa.
  • Recozimento de Classificador (Classifier Annealing): Para mitigar artefatos causados pelo afastamento do manifesto real de dados, os autores utilizam o recozcimento de classificador, interpolando entre campos de orientação condicional e incondicional durante a decodificação.

Principais Contribuições

1. Tokenizador Global Adaptativo: Introdução de um tokenizador onde os tokens representam descritores globais em uma hierarquia de grosseiro para fino, em vez de vizinhanças locais.
2. Validação de Desempenho: Demonstração de que o modelo iguala ou supera os tokenizadores locais de estado da arte (SOTA) (ex: DPLM2, ESM3, Kanzi) em tarefas generativas, de aprendizado de representação e de reconstrução.
3. Técnicas de Inferência: Desenvolvimento de amostragem baseada em entropia e recozcimento de classificador para gerenciar a complexidade de amostragem e mitigar a exposição ao erro, levando a uma melhor projetabilidade (designability).
4. Aplicações Zero-Shot: Demonstração do desacoplamento do tamanho da proteína do condicionamento, permitindo aplicações zero-shot como encolhimento de proteína (reduzindo a contagem de resíduos enquanto mantém a estrutura) e maturação de afinidade via escalonamento em tempo de inferência.

Resultados

• Reconstrução: No CATH, CAMEO e conjuntos de teste AFDB, a APT alcança RMSD e TM-scores comparáveis aos modelos de difusão contínos SOTA. Notavelmente, usando apenas 32–64 tokens, obtêm-se RMSDs $< 2$ Å, afirmando a compressibilidade das estruturas proteicas.
• Geração: Um modelo autorregressivo treinado nos tokens APT alcança uma projetabilidade de 0,871 (fração de amostras com scRMSD $< 2$ Å), superando significativamente modelos de difusão discreta como DPLM2 (0,486) e ESM3. Isso é alcançado sem técnicas de amostragem "best-of-N".
• Aprendizado de Representação: Na tarefa de classificação CATH, a sondagem não linear (MLP) nos tokens APT supera a sondagem nas representações de ESM3 e DPLM2. Crucialmente, a APT fornece representações globais de tamanho fixo sem mean-pooling, e mesmo representações altamente comprimidas (16 tokens) superam modelos maiores.
• Aplicações:
  • Encolhimento de Proteína: Gerou com sucesso versões menores de proteínas (ex: hemoglobina, beta-barril) preservando a estrutura global e local ao condicionar o processo de difusão em menos resíduos.
  • Maturação de Afinidade: Demonstrou a capacidade de evoluir ligantes fracos em ligantes fortes usando busca de feixe (beam search) guiada por funções de recompensa (ex: iPAE) no espaço latente.

Significância

O artigo afirma que a Tokenização de Proteínas Adaptativa fornece uma abordagem consciente da tarefa para escalar tarefas biológicas. Ao mover-se da tokenização local para a global, o método resolve o compromisso entre fidelidade de reconstrução e capacidade generativa. A capacidade de comprimir proteínas em vetores de tamanho fixo sem perda de informação via mean-pooling oferece um caminho prático para o modelamento de grandes complexos proteicos e para realizar o escalonamento em tempo de inferência. Os autores posicionam este trabalho como um passo em direção a modelos generativos capazes de raciocinar através de múltiplas escalas de comprimento, abordando um desafio fundamental na bioengenharia de fronteira.

Nota: Os autores reconhecem limitações, especificamente que os tokens APT são eficazes para tarefas globais, mas menos adequados para tarefas locais como o andaime de motivos (motif scaffolding), e que o tokenizador atual não considera cadeias laterais ou composição de aminoácidos, que são críticos para certas funções.