🤖 machine learning

WriteSAE: Sparse Autoencoders for Recurrent State

Este artigo apresenta o WriteSAE, o primeiro autoencoder esparsos projetado para decompor e editar as gravações de cache de matriz de modelos de linguagem recorrentes híbridos e de espaço de estados, fatorando os átomos do decodificador em sua forma nativa de atualização de rank-1, permitindo intervenções comportamentais precisas e previsões de efeito altamente acuradas.

Autores originais: Jack Young

Publicado 2026-05-14

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Jack Young

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine um modelo de linguagem grande (como os que escrevem histórias ou respondem perguntas) como uma fábrica massiva e de alta velocidade. Dentro desta fábrica, há uma "placa de memória" especial onde a máquina anota seus pensamentos atuais antes de avançar para a próxima etapa.

Há muito tempo, cientistas têm tentado entender o que está escrito nesta placa usando ferramentas chamadas Autoencoders Esparsos (SAEs). Pense nessas ferramentas como um par de óculos que permite aos pesquisadores ver as "ideias" ou "características" individuais sobre as quais a máquina está pensando. No entanto, havia um problema: esses óculos só funcionavam na saída da fábrica, e não na maneira específica como a máquina escreve em sua placa de memória em certos modelos avançados.

Este artigo apresenta uma nova ferramenta chamada WriteSAE. Eis como ela funciona, explicada de forma simples:

1. O Problema: A Forma Errada dos Óculos

Em modelos mais antigos, a máquina escrevia seus pensamentos como uma simples lista de números. Os óculos antigos (SAEs padrão) foram projetados para ler listas.
Mas em modelos mais novos e rápidos (como Gated DeltaNet, Mamba-2 e RWKV-7), a máquina não escreve uma lista. Em vez disso, ela escreve uma matriz (uma grade de números) usando um truque matemático específico chamado "atualização de rank-1".

A Analogia: Imagine que a máquina está pintando um quadro. As ferramentas antigas tentavam descrever a pintura observando a tela finalizada. Mas as novas máquinas pintam adicionando uma única pincelada específica (uma única linha de cor) de cada vez a uma grade complexa. As ferramentas antigas não conseguiam ver aquela única pincelada porque estavam procurando por uma lista completa de cores, e não por uma única pincelada.

2. A Solução: WriteSAE

Os autores construíram o WriteSAE, um novo conjunto de óculos projetado especificamente para ver aquela única pincelada.

A Correspondência de Forma: Em vez de tentar ler uma lista, os "átomos" do WriteSAE (as peças que ele procura) têm exatamente a mesma forma da pincelada que a máquina usa. São padrões minúsculos de pincelada única.
O Resultado: Como a forma corresponde perfeitamente, o WriteSAE consegue isolar exatamente o que a máquina está escrevendo em sua memória em qualquer momento dado.

3. Os Experimentos: Trocando as Pinceladas

Para provar que isso funciona, os pesquisadores realizaram uma "cirurgia" na memória da máquina.

A Troca: Eles encontraram um momento em que a máquina escreveu uma pincelada específica. Eles apagaram aquela pincelada e a substituíram por uma pincelada "aprendida" de seu novo dicionário (WriteSAE).
O Teste: Eles compararam isso com dois outros cenários:
1. Apagá-la completamente (deixando um espaço em branco).
2. Colocar um rabisco aleatório.
O Desfecho: Quando eles trocaram pela pincelada do WriteSAE, a máquina continuou funcionando quase exatamente como antes (92,4% das vezes). Quando a apagaram ou usaram um rabisco aleatório, a máquina ficou confusa e cometeu erros.
A Metáfora: É como substituir uma engrenagem específica de um relógio por uma engrenagem sob medida que se encaixa perfeitamente. O relógio continua a marcar o tempo. Se você remover a engrenagem ou colocar uma pedra aleatória, o relógio para.

4. O Que Eles Encontraram

Dois Tipos de Pinceladas: Eles descobriram que a máquina usa dois tipos principais de pinceladas:
- Registradores: São pinceladas precisas e focadas que realizam tarefas específicas (como marcar o início de uma frase ou um nome próprio).
- Pacotes: São pinceladas mais dispersas que parecem realizar uma mistura de coisas.
Prever o Futuro: Eles encontraram uma fórmula matemática que pode prever exatamente como mudar uma pincelada específica alterará a próxima palavra da máquina. É como saber que, se você ajustar essa única engrenagem específica, o relógio tocará um segundo mais cedo.
Editando a Máquina: Eles usaram com sucesso essa ferramenta para "instalar" novos comportamentos. Por exemplo, eles puderam forçar a máquina a continuar falando sobre um tópico específico (como uma palavra de "meio-ranking") que ela normalmente não escolheria, simplesmente inserindo a pincelada correta na placa de memória.

5. Os Limites

O artigo é muito cuidadoso ao afirmar que isso funciona melhor em modelos que escrevem dessa maneira específica de "pincelada única" (rank-1).

Se um modelo escrever de uma maneira mais complexa (como usando duas pinceladas ao mesmo tempo ou um padrão diagonal), a ferramenta não funciona tão perfeitamente, embora ainda encontre alguns padrões.
A ferramenta funciona muito bem no modelo Qwen3.5 (um tipo específico de IA), e eles mostraram que funciona em outros modelos semelhantes também, mas a "fórmula mágica" para prever o futuro muda dependendo da arquitetura do modelo.

Resumo

O WriteSAE é uma nova ferramenta que nos permite ver e editar a maneira específica como modelos avançados de IA escrevem em sua memória interna. Ao ajustar a forma da ferramenta à forma da escrita da máquina, os pesquisadores podem trocar pensamentos específicos, prever como a máquina reagirá e até mesmo direcionar a máquina a dizer coisas que ela normalmente não diria, tudo sem quebrar a máquina. É a primeira vez que cientistas realizaram com sucesso esse tipo de "cirurgia" diretamente no local de escrita da memória desses tipos específicos de modelos.

Resumo Técnico: WriteSAE para Decomposição de Estado Recorrente

Declaração do Problema

Modelos de linguagem híbridos e de espaço de estado recorrentes (por exemplo, Gated DeltaNet, Mamba-2, RWKV-7, Qwen3.5) utilizam um cache matricial para propagação de estado, onde cada token escreve um produto externo de posto 1 ( $k_t v_t^\top$ ) em uma matriz $d_k \times d_v$ . Os Autoencoders Esparsos (SAEs) existentes operam sobre fluxos residuais, lendo o estado após a emissão. Consequentemente, os SAEs padrão não conseguem acessar ou decompor diretamente o mecanismo de escrita nativo, que ocorre a montante na recorrência. Além disso, como o estado é atualizado via produtos externos de posto 1, um átomo decodificador baseado em vetor de um SAE padrão não pode substituir uma única escrita no cache sem romper a correspondência estrutural necessária para intervenção causal (correção de cache).

Metodologia

O artigo introduz o WriteSAE, o primeiro autoencoder esparso projetado para decompor e editar a escrita no cache matricial de modelos recorrentes.

Arquitetura e Treinamento

Correspondência de Forma Nativa: Os átomos decodificadores do WriteSAE são restritos a ser produtos externos de posto 1 ( $v_i w_i^\top$ ), correspondendo ao primitivo de escrita nativo ( $k_t v_t^\top$ ) da arquitetura hospedeira. Isso permite que um único átomo substitua uma atualização específica de slot no cache.
Codificador Bilinear: Para codificar eficientemente o estado matricial $S_t$ , o modelo utiliza um codificador bilinear ( $a_i = v_i^\top S_t w_i$ ) em vez de um codificador denso achatado. Isso reduz a contagem de parâmetros por átomo de $d_k d_v$ (16.384 para $128 \times 128$ ) para $d_k + d_v$ (256), uma redução de 64x.
Objetivo de Treinamento: O modelo é treinado em estados centrados na média ( $x = \text{vec}(S_t - M)$ ) usando uma restrição de esparsidade TopK e uma perda auxiliar para reviver características mortas. O corpus de treinamento consiste em sequências OpenWebText processadas através do Qwen3.5-0.8B (e outros modelos) em camadas e cabeças específicas.

Classificação de Características

O dicionário treinado é partitionado em duas classes com base na similaridade do cosseno mediana entre a direção de escrita do átomo e a direção de escrita nativa:

Registradores: Átomos com alta alinhamento de cosseno às escritas nativas (recuperáveis do cache).
Pacotes: Átomos com direções de escrita dispersas através do cache.

Testes de Intervenção Causal

O artigo propõe três mecanismos primários de validação:

Substituição de Slot de Cache: Substituir a escrita nativa ( $k_t v_t^\top$ ) por um átomo SAE com norma de Frobenius correspondente ( $\hat{S}_t = S_t - \Delta_{nat} + \Delta_{atom}$ ) e medir a divergência KL a jusante.
Deslocamento de Logit em Forma Fechada: Derivar uma expressão de três fatores para prever deslocamentos de logit por token: $\Delta \ell \approx G_{t_0 \to t} \langle w_i, q_t \rangle \langle v_i, W_U[\text{tok}] \rangle$ , onde $G$ é o produto do portão, $q_t$ é a consulta de leitura e $W_U$ é a linha de desincorporação.
Intervenções Diretas no Cache: Apagar ou instalar átomos em posições específicas para observar mudanças nas probabilidades de tokens e no comportamento de geração.

Contribuições Principais

Dicionário Correspondente à Arquitetura: Um dicionário onde os átomos são produtos externos de posto 1, permitindo a substituição direta de escritas nativas no cache, uma capacidade anteriormente indisponível para SAEs de fluxo residual.
Validação de Substituição: Demonstração de que átomos aprendidos podem substituir escritas nativas com alta fidelidade. No Qwen3.5-0.8B (Camada 9, Cabeça 4), os átomos superaram a ablação com norma correspondente em 92,4% de 4.851 ativações.
Previsão em Forma Fechada: Uma fórmula sem parâmetros e de três fatores que prevê deslocamentos de logit por ativação com um $R^2$ mediano de 0,98 em 200 células de átomo-por-épsilon.
Transferência Cross-Arquitetura: Validação de que o sucesso da substituição e a partição de características (Registrador vs. Pacote) mantêm-se através de diferentes arquiteturas de recorrência matricial (Gated DeltaNet, RWKV-7, Mamba-2), embora as taxas de sucesso variem conforme o "posto de escrita" do substrato.

Resultados Experimentais

Desempenho de Substituição

Qwen3.5-0.8B (Gated DeltaNet): Na Camada 9 Cabeça 4, os átomos superaram a ablação em 92,4% das ativações. O teste populacional sobre 87 átomos mantém-se em 89,8%. A ordenação estrita $KL_{atom} < KL_{ablate} < KL_{random}$ mantém-se em 89,5% das ativações.
Mamba-2-370M: A substituição tem sucesso em 88,1% sobre 2.500 ativações, apesar do Mamba-2 utilizar uma atualização de estado diagonal em vez de um produto externo completo de posto 1.
RWKV-7: Mostra desempenho intermediário, consistente com sua regra de escrita de posto 2.

Previsão de Deslocamento de Logit

A expressão em forma fechada rastreia com precisão os efeitos medidos. Na Camada 9 Cabeça 4, o $R^2$ mediano entre os deslocamentos de logit previstos e medidos é 0,98. No entanto, esta forma fechada falha (R² negativo) no Mamba-2 e no Qwen3.5-4B, indicando que o fator de portão $G$ é específico do substrato.

Intervenções Comportamentais

Apagamento: Apagar um átomo de registrador específico (F412) em suas posições naturais de ativação reduziu o log-probabilidade do token alvo ("espaço") em um mediano de -0,116 nats ( $p=1,07 \times 10^{-6}$ ).
Direcionamento de Geração: A instalação sustentada de uma direção em forma fechada em três posições consecutivas aumentou a taxa "alvo-em-continuação" para tokens de rank médio de 33,3% para 100% sob decodificação gulosa.
Amplificação ao Nível de Trecho: Amplificar características correlacionadas com limites em um modelo de 4B reduziu a geração de novas linhas em 33% (de 16,8 para 11,2 por 400 tokens) em uma dose de 5x, demonstrando controle sobre a estrutura do documento.

Escopo Arquitetural

O artigo estabelece uma hierarquia de eficácia do "posto de escrita":

Gated DeltaNet (Posto-1): Maior separação de cosseno de registradores (0,262) e fidelidade de substituição.
RWKV-7 (Posto-2): Separação intermediária (0,180).
Mamba-2 (Diagonal/Escalar): Menor separação (0,0575), ainda assim alcança ~88% de sucesso na substituição.

Significado e Alegações

O artigo alega fornecer o primeiro dicionário causal para o local de escrita de recorrência matricial, indo além da análise de fluxo residual. Ele demonstra que:

O Posto de Escrita é Determinante: A capacidade de decompor o estado em "registradores" interpretáveis correlaciona-se com o posto da regra de escrita nativa (Posto-1 > Posto-2 > Diagonal).
Substituição Causal é Possível: Átomos aprendidos podem substituir escritas nativas com divergência KL a jusante mínima, validando o conceito de "registrador" em modelos recorrentes.
Interpretabilidade em Forma Fechada: Para o Gated DeltaNet, o efeito de uma característica sobre os logits pode ser previsto analiticamente sem ajuste de parâmetros, ligando o dicionário diretamente aos mecanismos de portão e leitura do modelo.

Os autores notam modestamente limitações: o coeficiente em forma fechada não se transfere entre substratos (por exemplo, falha no Mamba-2) e a identidade por átomo é específica da semente (embora propriedades de nível de classe sejam estáveis). O trabalho estabelece um novo marco para interpretabilidade mecânica em modelos de espaço de estado, mostrando que o primitivo de "escrita" é um alvo viável para aprendizado de dicionário e edição causal.