Resumo Técnico: SAEs de Média por Turno para Descoberta de Características e Atribuição de Longo Contexto

1. Declaração do Problema

Os Autoencoders Esparsos (SAEs) emergiram como uma ferramenta primária para extrair características interpretáveis de modelos de linguagem, facilitando a análise de circuitos, o direcionamento comportamental e o monitoramento de segurança. No entanto, as arquiteturas de SAE padrão operam em ativações de tokens individuais. Essa granularidade por token cria dois gargalos significativos para a interpretabilidade em escala:

1. Limitações de Escala: O número de características ativas escala linearmente com o comprimento do contexto. Em cenários de contexto longo ou conversas de múltiplos turnos, o volume de dados resultante torna a análise subsequente (como gráficos de atribuição) complexa e difícil de inspecionar.
2. Perda de Contexto de Alto Nível: Muitos objetivos de interpretabilidade envolvem entender por que um modelo produziu uma resposta específica, exigindo a identificação de propriedades de alto nível (ex: tópico, persona, traços de segurança) que estão distribuídas por muitos tokens. Os SAEs por token padrão frequentemente aprendem detectores de baixo nível para palavras específicas ou padrões sintáticos. Como o objetivo de treinamento é reconstruir cada posição de forma independente, não há pressão para que o modelo separe essas características de alto nível e de amplitude de intervalo (span-wide) do ruído ao nível do token.

2. Metodologia

2.1 SAEs de Média por Turno

A proposta central é treinar SAEs na média do estado oculto através de um intervalo de tokens (especificamente, um único turno Humano ou Assistente) em vez de em ativações de tokens individuais.

• Transformação de Entrada: Dado um turno $T$ com tokens $x_t$, a entrada é a média do estado oculto $\bar{x} = \frac{1}{|T|} \sum_{t \in T} x_t$.
• Arquitetura: A arquitetura do SAE permanece idêntica às implementações padrão (usando BatchTopK para esparsidade). Como a média de vetores em $\mathbb{R}^{d_{model}}$ retém a mesma dimensionalidade que os vetores individuais, nenhuma mudança estrutural é necessária.
• Intuição: A média atua como um filtro passa-baixa. Ela suaviza detalhes de alta frequência específicos do token (como escolhas de palavras específicas ou marcadores sintáticos) enquanto preserva componentes de baixa frequência que correspondem a conceitos amplos como tópico, estilo e traços comportamentais.

2.2 Arquiteturas Híbridas

Para abordar a limitação de que a média descarta o detalhe por token, os autores propõem duas abordagens híbridas para capturar tanto a estrutura de alto nível quanto os resíduos finos:

1. SAE Desacoplado: Treina dois SAEs independentes. Um reconstrói a média do turno ($\bar{x}$) e o outro reconstrói o resíduo ($x_t - \bar{x}$). Embora isso separe os sinais, os modelos são causalmente desconectados, complicando a análise de atribuição.
2. SAE Aninhado (estilo Matryoshka): Treina um único SAE com um codificador e decodificador compartilhados, mas uma função de perda especializada. O espaço de características é dividido em uma partição interna $[0:h]$ e uma partição externa $[h:d_{sae}]$.
   • Função de Perda: $L = \sum_t \| \text{dec}(\text{enc}(x_t)) - x_t \|^2 + \alpha \cdot \| \text{dec}_h(\text{enc}(\bar{x})[0:h]) - \bar{x} \|^2$.
   • A partição interna é forçada a reconstruir a média do turno, enquanto o conjunto completo de características reconstrói as ativações por token. Isso permite que um único modelo aprenda simultaneamente características de alto nível do turno e detalhes ao nível do token.

2.3 Adaptação de Gráfico de Atribuição

Os autores adaptam os gráficos de atribuição para o nível do turno. Em vez de nós representando características em posições de tokens específicos, os nós representam características ativas em um turno específico e camada.

• Pesos de Aresta: Calculados pela média da atribuição de gradiente através do turno de origem e normalizados pelo comprimento do turno ($1/|T_A|$ e $1/|T_B|$) para garantir que os pesos sejam agnósticos ao comprimento do turno.
• Validação de Intervenção: Os pesos das arestas são validados perturbando as características de origem ao longo de suas direções de decodificação em todos os tokens de um turno e medindo a mudança nas pré-ativações da característica alvo.

3. Principais Contribuições

1. SAEs de Média por Turno: Demonstrou que o treinamento em estados ocultos médios produz características que capturam características de alto nível do turno (tópico, estilo, traços de segurança) de forma mais completa do que as características por token.
2. Arquitetura de SAE Aninhado: Introduziu uma arquitetura de modelo único que aprende conjuntamente características de média por turno e por token usando uma perda tipo Matryoshka, alcançando alta qualidade de reconstrução para ambos os componentes do sinal.
3. Atribuição Escalável: Adaptou os gráficos de atribuição para o nível do turno, permitindo uma análise compacta de conversas de múltiplos turnos. Os autores validaram esses gráficos através de experimentos de intervenção e métricas de suficiência (completude e substituição).

4. Resultados

4.1 Qualidade de Característica e Reconstrução

• Reconstrução: O SAE Aninhado atinge um Erro Quadrático Médio Normalizado (nMSE) por token de 0,167, aproximando-se do SAE puramente por token (0,162), enquanto também atinge um nMSE de média por turno de 0,116. Em contraste, um SAE padrão por token tem desempenho ruim em entradas de média (nMSE 0,355), indicando que a média está fora da distribuição para modelos treinados por token.
• Descoberta Qualitativa: Em exemplos envolvendo erros matemáticos, estilos técnicos entusiastas e conformidade regulatória, as características de média por turno identificaram com sucesso propriedades comportamentais de alto nível (ex: "respostas incorretas", "estilo técnico entusiasta") que as características por token perderam ou capturaram apenas parcialmente.
• Generalização de Longo Contexto: SAEs de média por turno treinados em turnos de chat (comprimento médio de ~150 tokens) generalizaram com sucesso para documentos 150x mais longos (24k–32k tokens) do LongBench v2, identificando conceitos de alto nível como "macroeconomia" e "computação gráfica" que as características por token falharam em detectar.

4.2 Avaliação Quantitativa

Usando um framework de LLM-como-juiz para comparar descrições de características com resumos estruturados de turnos:

• Discriminação (correspondência de 10 vias): As características por token (max-pooling) superaram as características de média por turno (95,0% vs. 73,9%), pois dependem de detalhes específicos de nível de token para distinguir textos.
• Cobertura (ranking de pareamento e 5 vias): As características de média por turno superaram significativamente as características por token (vencendo 87,9% das comparações pareadas). Elas forneceram descrições mais completas das características de alto nível de um turno.
• Métrica de Embedding: As características de média por turno alcançaram a maior similaridade de cosseno (0,663) com campos de resumo estruturados, confirmando sua superior cobertura de conteúdo semântico.

4.3 Análise de Atribuição

• Compactação de Gráfico: Os gráficos de média por turno reduziram a contagem de nós de ~128K (por token) para ~5K para uma conversa de 10 turnos, tornando a inspeção viável.
• Estudos de Caso:
  • Atribuição de System Prompt: Gráficos de média por turno identificaram corretamente quais parágrafos do system prompt influenciaram consultas específicas do usuário, mostrando mudanças de atribuição comparáveis aos gráficos por token.
  • Rastreamento de Degeneração: Em uma conversa de 14 turnos onde o modelo degradou para texto repetitivo, o gráfico de média por turno traçou uma cadeia causal coerente: "Avisos de segurança de IA" $\to$ "Armas nucleares" $\to$ "Detecção de prompt adversário" $\to$ "Saída degenerada". Os gráficos por token falharam em encontrar a característica de detecção adversária, rastreando-a em vez disso para conceitos não relacionados como "compostos químicos".
• Intervenção e Suficiência: Os pesos das arestas mostraram correlação significativa com os efeitos de intervenção. O SAE Aninhado alcançou as pontuações de suficiência de "completude" e "substituição" mais altas, superando tanto o SAE puramente por token quanto as agregações post-hoc de características por token.

5. Significância e Alegações

O artigo afirma que os SAEs de média por turno tornam as técnicas de interpretabilidade práticas em comprimentos de contexto longos. Ao mudar a unidade de análise de tokens para intervalos (spans), o método:

• Produz características que descrevem características de alto nível (tópico, persona, segurança) de forma mais completa do que as características por token.
• Permite a construção de gráficos de atribuição compactos e interpretáveis para conversas de múltiplos turnos que, de outra forma, seriam grandes demais para analisar.
• Demonstra que modificar o sinal de treinamento (via média) pode forçar os SAEs a aprender características na granularidade desejada (conceitos de alto nível) independentemente da contagem de tokens.
• Mostra que essas características generalizam para intervalos muito mais longos do que os vistos durante o treinamento, sugerindo robustez para a análise de modelos de fronteira.

Os autores observam limitações, incluindo a escala atual dos experimentos (modelo 7B, corpus LMSYS-Chat-1M) e a dependência de LLMs para métricas de avaliação, identificando a aplicação a modelos de fronteira e a criação de benchmarks de verdade fundamental (ground-truth) como trabalhos futuros necessários.