Resumo Técnico: Prototype Transformer (ProtoT)

Declaração do Problema

Os modelos de linguagem (LMs) autorregressivos de última geração, como o GPT-4 e o LLaMA, alcançam um desempenho superior em diversos domínios, mas sofrem com a falta de transparência em seus processos de raciocínio. Seus caminhos computacionais internos são amplamente opacos, o que gera desafios para compreender como os outputs são gerados, riscos de alucinação e dificuldades em edições comportamentais direcionadas. Os métodos atuais de interpretabilidade são primariamente post-hoc (ex: visualização de atenção, sondagem, autoencoders esparsos) e operam em arquiteturas que não foram projetadas com a interpretabilidade como uma restrição primária. Além disso, os mecanismos padrão de autoatenção incorrem em custos computacionais quadráticos em relação ao comprimento da sequência, limitando a escalabilidade.

Metodologia: O Prototype Transformer (ProtoT)

Os autores introduzem o ProtoT, uma arquitetura de LM autorregressiva que substitui o mecanismo de autoatenção de custo quadrático por um módulo de comunicação baseado em protótipos de custo linear. O ProtoT mantém a estrutura geral do LLaMA-3 (incluindo redes feed-forward SwiGLU e RMS pre-layernorm), mas altera fundamentalmente a camada "mixer".

Arquitetura Central

1. Mixer Baseado em Protótipos: Em vez de autoatenção, o ProtoT utiliza $R$ vetores de parâmetros aprendíveis chamados protótipos. Esses protótipos atuam como filtros que definem $R$ canais de comunicação independentes.
2. Portões de Escrita e Leitura (Write and Read Gates):
   • Portão de Escrita (Write Gate): Agrega informações de tokens passados ($x_{j}, j < i$) nos $R$ canais. Ele computa pontuações de similaridade entre os tokens de entrada e os protótipos, usando um softmax sobre os protótipos para rotear a informação para canais específicos. Isso cria um "gargalo de roteamento semântico" onde múltiplos significados são forçados a uma representação de compromisso por canal, incentivando a especialização semântica.
   • Portão de Leitura (Read Gate): Recupera informações dos $R$ canais com base na similaridade entre o token atual e os protótipos.
3. Média de Prefixo (PM) & Desconto Temporal: A informação agregada em cada canal é mantida como uma média de prefixo das ativações passadas. Um parâmetro de desconto temporal aprendível ($\beta_k$) aplica uma Média Móvel Exponencial (EMA) à média do prefixo, permitindo que diferentes protótipos operem em diferentes escalas de tempo (dependências de curto prazo vs. longo prazo).
4. Complexidade Linear: A arquitetura alcança complexidade computacional linear ($O(R)$) em relação ao comprimento da sequência, pois os protótipos não interagem entre si (ao contrário da interação $O(R^2)$ em espaços latentes do tipo Perceiver) e a recorrência permite a geração de tokens em tempo constante ($O(1)$) via cache.
5. Portão Alpha (Alpha Gate): Um portão escalar do tipo ReZero é aplicado na saída de cada módulo mixer para depurar a estabilidade do treinamento e garantir que o mixer contribua para as previsões.

Principais Contribuições

1. Inovação Arquitetural: A introdução do ProtoT, uma nova arquitet verdade que substitui a autoatenção por comunicação multicanal baseada em protótipos, alcançando complexidade computacional linear enquanto mantém um desempenho competitivo.
2. Interpretabilidade Emergente: Demonstração de que os protótipos aprendem automaticamente conceitos desentranhados e nomeáveis (ex: "mulher", "Nova Zelândia", "COVID-19") durante o treinamento através de vários níveis de abstração. Isso ocorre sem supervisão explícita, impulsionado pela pressão arquitetural para especializar os canais.
3. Edição Comportamental Direcionada: Evidência de que o raciocínio do modelo pode ser editado cirurgicamente ao intervir em protótipos específicos (ex: reinicializando ou mascarando um protótipo "feminino"). Essas intervenções alteram significativamente a probabilidade de tokens relacionados (ex: diminuindo "mulheres" em ~16%) enquanto causam degradação negligenciável na perplexidade geral do modelo (<1% na maioria dos casos).
4. Análise de Robustez: Demonstração de que o ProtoT exibe robustez a perturbações lexicais que preservam o significado (sinônimos, erros de digitação) e que essa robustez é mediada pelos mecanismos de roteamento de protótipos.

Resultados Experimentais

Os autores avaliaram o ProtoT contra três baselines: LLaMA (autoatenção padrão), Mamba (modelo de espaço de estados) e DeltaNet (atenção linear), utilizando um backbone fixo (6 camadas, 256 tamanho oculto) e treinando em 250M de tokens do FineWeb-Edu.

• Escalabilidade: O ProtoT escala bem com o tamanho do modelo e o volume de dados. Em configurações de larga escala (512 tamanho oculto, 12 camadas), o ProtoT atinge uma perplexidade de 29.5, superando o DeltaNet (31.5) e o Mamba (26.5), embora fique atrás do LLaMA (25.8).
• Geração de Texto: Em avaliações de preferência humana pareadas (LLM-as-a-judge), o ProtoT supera o LLaMA e o DeltaNet, posicionando-se próximo ao Mamba.
• Tarefas de Downstream: No benchmark GLUE, o ProtoT permanece competitivo com baselines densos, frequentemente igualando ou superando o Mamba e o DeltaNet, e demonstrando forte robustez cross-domain.
• Throughput: O ProtoT escala melhor que o LLaMA com o aumento do comprimento do contexto devido ao seu compute linear, superando o LLaMA aos 32k tokens, embora o DeltaNet mantenha o maior throughput em contextos muito longos.
• Métricas de Interpretabilidade: Os protótipos do ProtoT pontuam mais alto em desentranhamento e cobertura de temas principais comparados às cabeças de atenção do LLaMA e às features de Autoencoders Esparsos (SAE).
• Eficácia da Intervenção: A interrupção de protótipos específicos (ex: L9 P7 para "feminino") causou quedas significativas nas probabilidades de tokens alvo (ex: -16.60% para "mulheres") com danos colaterais mínimos à coerência da sequência.

Significância e Alegações

O artigo afirma que o ProtoT representa um passo promissor em direção a modelos de linguagem autorregressivos que são interpretáveis por design. Ao integrar a interpretabilidade diretamente na arquitetura via canais de protótipos, o modelo:

• Captura conceitos coerentes e nomeáveis automaticamente.
• Permite edições cirúrgicas e direcionadas ao comportamento do modelo sem desestabilizar todo o sistema.
• Oferece uma alternativa de compute linear à autoatenção que é robusta a perturbações.

Os autores reconhecem limitações, observando que, embora o desempenho seja competitivo com baselines lineares, ele atualmente subperforma modelos de autoatenção estabelecidos como o LLaMA em termos de perplexidade bruta. Eles também alertam que os rótulos dos protótipos devem ser tratados como aproximações e não como verdade absoluta, já que alguns protótipos permanecem polissêmicos. No entanto, o trabalho sugere que equilibrar capacidade com transparência é viável, oferecendo um caminho para modelos de linguagem mais seguros e controláveis.