PairSAE: Mechanistic Interpretability from Pair Representations in Protein Co-Folding

Resumo Técnico: PairSAE – Interpretabilidade Mecanística a partir de Representações de Pares em Co-dobramento de Proteínas

Declaração do Problema
Modelos de fundação para biologia estrutural, como o AlphaFold e os modelos Boltz, alcançaram uma precisão notável na previsão de estruturas e complexos biomoleculares. No entanto, compreender as características internas que impulsionam esses resultados permanece um desafio significativo. Embora os Autoencoders Esparsos (SAEs) tenham conseguido o desmembramento de características "sobrepostas" (superposed) em modelos de transformadores de linguagem (LLMs) e modelos de linguagem de proteínas (pLMs), mapeando-as para conceitos monossemânticos, eles não se transferem de forma limpa para modelos de previsão de estrutura.

Diferente dos LLMs padrão, os modelos de previsão de estrutura utilizam arquiteturas pairformer que mantêm tanto as incorporações (embeddings) de sequência quanto tensores de interação de pares. A aplicação ingênua de SAEs a esses modelos apresenta dois problemas críticos:

1. Explosão Quadrática: Aprender dicionários separados para representações de pares escala quadraticamente com o comprimento da sequência, tornando a análise computacionalmente proibitiva.
2. Conceitos Obscurecidos: As características podem estar sobrepostas conjuntamente nos espaços de sequência e de par. Tratá-las separadamente falha em capturar conceitos distribuídos entre ambas as representações, enquanto tratá-las independentemente obscurece a natureza conjunta das características aprendidas.

Metodologia: PairSAE
Os autores introduzem o PairSAE, um framework projetado para aprender um conjunto compartilhado de características ao nível de token que reconstrui conjuntamente as representações de sequência e de par. O método consiste em três estágios primários:

1. Resumo de Interações de Pares via SVD de N-modos:
   Para evitar a escala quadrática dos tensores de pares ($Z \in \mathbb{R}^{N_{tok} \times N_{tok} \times n_z}$), os autores realizam uma Decomposição de Valor Singular (SVD) de N-modos. Eles extraem as matrizes de modo-1 e modo-2 dos vetores singulares à esquerda ($U^{(1)}$ e $U^{(2)}$), que codificam os papéis de interação de linha e coluna de cada token. Estes são truncados para os $r$ vetores singulares superiores e concatenados para formar uma incorporação de resumo por token $m_i$. Este passo comprime a informação de pares em uma representação de nível de token que preserva os papéis de interação sem a complexidade quadrática.

2. Arquitetura de Autoencoder Compartilhado:
   O codificador PairSAE recebe a concatenação da incorporação de sequência original ($s_i$) e do resumo de interação derivado via SVD ($m_i$) como entrada. Ele produz um vetor de característica esparsa $h_i$ usando uma não-linearidade indutora de esparsidade (uma composição de BatchTopK e ReLU).
   Crucialmente, o decodificador usa este conjunto de características compartilhadas $h$ para reconstruir tanto:

• A incorporação de sequência: $\hat{s}_i = D_s h_i + b_{dec}^s$
• A incorporação de par: $\hat{z}_{i,j} = D_z^{row} h_i + D_z^{col} h_j + b_{dec}^z$
  Isso força o modelo a aprender uma base de características unificada que explica tanto o contexto local da sequência quanto as interações de par.

3. Função de Perda e Treinamento:
   O modelo é treinado usando uma perda Matryoshka SAE, que fatia o decodificador e os vetores de características através de larguras aninhadas ($c_1 < c_2 < c_3 = D$) para garantir robustez entre tamanhos de dicionário. A perda total inclui termos de reconstrução para as representações de sequência e de par, ponderados igualmente, além de uma perda auxiliar para evitar características "mortas". O treinamento é acelerado computando uma aproximação de Monte Carlo para o termo de perda de par, amostrando um único $j$ para cada $i$.

Principais Contribuições

• Arquitetura PairSAE: Um método inovador para aplicar o aprendizado de dicionário esparso a modelos baseados em pairformer ao resumir tensores de pares via SVD de N-modos e decodificar tanto para espaços de sequência quanto de par a partir de um conjunto de características compartilhado.
• Descoberta de Características Interpretáveis: O método recupera com sucesso características que se alinham com conceitos biofísicos, evitando as armadilhas dos SAEs convencionais quando aplicados a modelos de biologia estrutural.
• Aplicação à Afinidade de Ligação: O framework é aplicado às ativações do Boltz-2 em complexos proteína-ligante (dataset PLINDER) para ligar as representações latentes às previsões de afinidade de ligação.

Resultos
Os autores avaliaram as representações do Boltz-2 em complexos proteína-ligante:

• Probing Linear: As características do PairSAE demonstraram alto poder preditivo para anotações de resíduos UniProt. Características das camadas 33 e 64 (terceiro passo de recycling) alcançaram pontuações F1 de 29,1% e 24,0% ao nível de token, e 53,2% e 49,6% ao nível de complexo, superando significativamente classificadores baseados em neurônios do ESM2-650M (que alcançaram 0,8% e 19,7%, respectivamente).
• Alinhamento de Conceitos: Visualizações mostraram características específicas ativando em motivos estruturais distintos, como regiões transmembrana, pontes de dissulfeto e sítios ativos de protease, com pontuações F1 elevadas contra a verdade fundamental (ground truth) do UniProt.
• Previsão de Afinidade: Usando regressão LASSO nas características de max-pooling do PairSAE, os autores previram valores de afinidade do Boltz-2 com um $R^2$ de 0,528 em um conjunto de teste (PoseBusters). Características específicas (ex: característica 2299) foram identificadas como correlacionadas fortemente com complexos de alta afinidade, ativando-se especificamente em ligantes e mostrando diferenças significativas de grupo em comparação com complexos inativos.

Significância e Alegações
O artigo afirma que o PairSAE consegue ligar o espaço latente de modelos de fundação para biologia estrutural a conceitos estruturais interpretáveis. Ao abordar as restrições arquitetônicas específicas dos modelos pairformer, o PairSAE esclarece "o que o modelo sabe" em relação ao dobramento e ligação sem sucumbir à complexidade quadrática ou aos problemas de sobreposição de características que limitam os SAEs convencionais.

Os autores posicionam este trabalho como um passo em direção à avaliação e descoberta fundamentada em biologia estrutural, permitindo que pesquisadores avaliem a confiabilidade da previsão e a plausibilidade biofísica. Eles observam que, embora o método mapeie efetivamente características de sequência e par para anotações conhecidas, permanece uma limitação em mapear características que ativam ligantes para conceitos específicos devido à falta de anotações detalhadas de ligantes. Um trabalho futuro é proposto para escalar o estudo para mais camadas, integrar estas características em pipelines automatizados e desenvolver métodos de controle (steering) para o design de proteínas interpretável.