VFUSE: Virulent Feature Understanding with Sparse autoEncoders

Resumo Técnico: VFUSE – Compreensão de Características Virulentas com Autoencoders Esparsos

Definição do Problema

Modelos generativos para o design de proteínas, particularmente modelos de difusão de todos os átomos como o RoseTTAFold3 (RF3) e o RFDiffusion3 (RFD3), possuem a capacidade de duplo uso para projetar proteínas perigosas (ex: neurotoxinas, proteínas de evasão imune) ao estruturar novos scaffolds em torno de templates perigosos. Os métodos atuais de triagem de biossegurança dependem de homologia de sequência ou classificadores de nível de sequência. Essas abordagens sofrem de duas limitações críticas: operam post-hoc (sinalizando sequências após a geração) e carecem de interpretabilidade mecanística, o que significa que não conseguem explicar quais características estruturais específicas dispararam um alerta de perigo ou intervir durante o processo de amostragem. Além disso, embora os Autoencoders Esparsos (SAEs) tenham sido aplicados com sucesso a modelos de linguagem de proteínas (PLMs) e modelos de visão, eles ainda não foram utilizados para sondar as representações internas de modelos de difusão de todos os átomos.

Metodologia

Os autores introduzem o VFUSE, um framework de interpretabilidade mecanística que treina Autoencoders Esparsos nas ativações de modelos de difusão-transformer para auditar características associadas a perigos.

1. Modelo & Coleta de Ativações:

   • O estudo foca no RFDiffusion3 (um modelo de difusão de todos os átomos de 18 blocos) e no RoseTTAFold3 (seu contraparte de dobramento).
   • As ativações ($d=768$) são capturadas via hooks de PyTorch em blocos específicos do fluxo residual: blocos 6, 8 e 12 para RFD3, e blocos 12 e 16 para RF3.
   • Para o RFD3, os inputs são gerados via difusão parcial (em $\partial t = 5$, aproximadamente 5Å de ruído) em estruturas PDB de templates perigosos para simular o processo de design do modelo. Para o RF3, os inputs são dobrados diretamente.

2. Construção do Dataset:

   • Sequências Virulentas: Obtidas do SafeProtein (275 pares) e ToxinPred3 (1.200 sequências apenas para RF3).
   • Sequências Benignas: Sequências de comprimento correspondente (100–300 aa) do UniProt, filtradas para excluir palavras-chave de toxina/virulência.
   • Controle de Homologia: Para evitar que as sondas memorizem a identidade da família de dobras em vez de aprender sinais de virulência, as sequências são agrupadas em 30% de identidade usando mmseqs2. As divisões de treino/teste garantem que nenhum homólogo próximo apareça em ambos os conjuntos.

3. Arquitetura & Treinamento do SAE:

   • Arquitetura: Um Matryoshka BatchTopK SAE é empregado. Este combina:
     • Matryoshka: Tamanhos de prefixo aninhados que permitem a reconstrução em múltiplas granularidades.
     • BatchTopK: Impõe esparsidade mantendo apenas os top $K$ pré-ativações através de um lote, substituindo as penalidades L1.
   • Hiperparâmetros: Tamanho do dicionário $m=12.288$ (expansão de 16x), $K=80$ latentes ativos, treinado por 20.000 passos.
   • Performance: Os SAEs treinados explicam 96,9% da variância de ativação retida com uma esparsidade $L_0$ média de 79,9.

4. Protocolo de Avaliação:

   • Sondagem (Probing): Vetores latentes do SAE ($z$) com pooling médio e ativações brutas ($x$) são usados para treinar sondas de regressão logística com regularização L2.
   • Métricas: O AUROC é avaliado sob dois regimes: divisões estratificadas aleatórias e divisões agrupadas por homologia.
   • Pontuação de Características (Feature Scoring): O AUROC univariado é computado para cada característica contra os rótulos de classe. A significância é determinada via testes Mann–Whitney U com correção de FDR de Benjamini–Hochberg ($q < 0,05$).

Principais Resultados

1. Latentes do SAE Superam Ativações Brutas:

   • No RFD3 Bloco 12, as características codificadas pelo SAE superam significativamente as ativações brutas na configuração de homologia agrupada, atingindo um AUROC de 0,817 comparado a 0,763 para características brutas (uma melhoria de +0,054).
   • Esta melhoria é específica para o Bloco 12; blocos anteriores (6, 8) mostram ganhos neutros ou negativos com os SAEs.
   • Para o RF3 (dataset ToxinPred3), as características do SAE do Bloco 12 também mostram uma melhoria de +0,041 sobre as características brutas em divisões agrupadas.

2. Redução na Memorização de Homologia:

   • A lacuna entre o AUROC de divisão aleatória e o de homologia agrupada (um proxy para memorização de família) é menor para as características do SAE do RFD3 Bloco 12 (lacuna de 0,13) comparado ao Bloco 6 (lacuna de 0,22). Isso sugere que o SAE do Bloco 12 captura sinais estruturais discriminativos de classe e generalizáveis à família, em vez de memorizar famílias de dobras.

3. Identificação de Características de Perigo Monossemânticas:

   • O estudo identifica características individuais do SAE que disparam seletivamente em designs perigosos.
   • Top Feature: A Característica 639 no RFD3 Bloco 12 atinge um AUROC de 0,819 ($q \approx 10^{-13}$) para detectar neurotoxinas de víboras.
   • Localização: Estas características exibem localização por resíduo. Para a Característica 639, os top-3 resíduos ativadores (99, 104, 109) agrupam-se dentro de uma única alfa-hélice da estrutura da Ammodytoxin A. Outras características localizam-se de forma semelhante em regiões específicas de proteínas de evasão imune de Monkeypox, proteínas de inativação de ribossomos e toxinas salivares de mosquitos.

4. Performance Competitiva:

   • A melhor sonda VFUSE (RFD3 Bloco 12 SAE, divisão aleatória, AUROC 0,877) é competitiva com os classificadores de sequência de estado da arte (ex: DTVF com 0,92), apesar de usar uma sonda logística simples e um dataset menor. Crucialmente, o VFUSE fornece racionalidades por resíduo, algo que modelos baseados apenas em sequência não possuem.

Significância e Alegações

O artigo afirma ser o primeiro SAE treinado em um modelo de difusão de todos os átomos e a primeira auditoria de virulência ao nível de característica de um modelo de design de proteínas.

• Insight Mecanístico: O trabalho demonstra que o aprendizado de dicionário em profundidades específicas (Bloco 12) desentrelaça conceitos estruturais polissêmicos em direções monossemânticas que são mais propensas à detecção de perigo de baixa capacidade do que as representações brutas do modelo.
• Segurança e Interpretabilidade: O VFUSE estabelece um caminho para monitoramento em tempo de execução e auditoria interpretável de modelos generativos de proteínas. Ao identificar resíduos específicos e motivos estruturais associados à virulência, o método vai além da triagem de "caixa preta" para verificações de segurança explicáveis.
• Direção Futura: Embora os autores notem limitações (dataset pequeno, falta de sucesso em testes de direcionamento/steering preliminares), o trabalho pavimenta o caminho para o design de proteínas seguro e interpretável, permitindo a detecção de características perigosas ao nível da característica, em vez de apenas ao nível da sequência.