Resumo Técnico: VASAE – Nomeação de Direções do Dicionário de SAE com Ancoragem Alinhada ao Vocabulário

1. Declaração do Problema

Autoencoders Esparsos (SAEs) são cada vez mais utilizados para decompor os fluxos residuais de alta dimensão de modelos Transformer em códigos esparsos sobre dicionários aprendidos, auxiliando na interpretabilidade mecanística. No entanto, uma limitação significativa dos SAEs padrão é que suas características aprendidas (direções do dicionário) não estão intrinsecamente conectadas ao vocabulário de tokens do modelo.

Na prática padrão, a nomeação de características é um processo post hoc. Pesquisadores devem inspecionar contextos onde características específicas ativam ou executar procedimentos de interpretação automatizados separados para atribuir rótulos semânticos. Esse descompasso surge porque a geometria do dicionário aprendido e a geometria da geometria dos embeddings de tokens são tipicamente não relacionadas; uma característica pode reconstruir efetivamente ativações enquanto reside longe de qualquer vetor de embedding de token, tornando-a sem um nome "intrínseco" ao nível do vocabulário.

A questão levantada pelo artigo é: Podemos treinar características de SAE para que permaneçam flexíveis o suficiente para uma reconstrução de alta qualidade, mas sejam geometricamente ancoradas suficientemente próximas de embeddings de tokens para adquirir nomes de tokens intrínsecos durante o treinamento, sem sacrificar o desempenho de reconstrução?

2. Metodologia: Autoencoder Esparso com Vocabulário Alinhado (VASAE)

Os autores propõem o VASAE, uma modificação ao objetivo de treinamento do SAE padrão que introduz um mecanismo de ancoragem alinhada ao vocabulário.

Arquitetura Central

• Entrada: O VASAE opera nos vetores pós-fluxo residual ($h^{(\ell)}$) de uma camada $\ell$ de um Transformer.
• Encoder: Utiliza um encoder esparso top-k padrão (mapa linear + ReLU + seleção Top-K) para produzir um código esparso $z$.
• Decoder: Utiliza uma matriz de dicionário aprendível $W_D$ (onde as colunas são vetores de características $d_i$) para reconstruir o fluxo residual.
• Âncoras: O método utiliza os embeddings de tokens de entrada fixos do modelo ($W_E$) como direções de referência. Esses embeddings compartilham a mesma dimensionalidade do fluxo residual (ou são projetados para tal).

O Objetivo de Ancoragem

A inovação principal é a adição de um termo de perda de âncora ao erro de reconstrução padrão.

1. Alinhamento com o Token Mais Próximo: Para cada característica aprendida $d_i$, o método calcula a similaridade de cosseno com todos os embeddings de tokens fixos $w_v$.
   $$s_i = \max_{v \in V} \cos(d_i, w_v)$$
   O token $v^*_i$ que alcança este máximo define o nome de token intrínseco para a característica $i$.
2. Perda de Âncora: O objetivo de treinamento incentiva as características a permanecerem próximas de seu embedding de token mais próximo:
   $$L_{anchor} = -\frac{1}{d_{sparse}} \sum_{i=1}^{d_{sparse}} s_i$$
   Minimizar esta perda maximiza a similaridade de cosseno entre as características e seus tokens de vocabulário mais próximos.
3. Perda Total: O objetivo combinado é:
   $$L_{vasae} = L_{recon} + \lambda_{anchor} L_{anchor}$$
   onde $\lambda_{anchor}$ é um hiperparâmetro que controla a força da restrição de alinhamento.

Distinção de Vinculação Rígida (Hard-Tying)

Os autores distinguem explicitamente o VASAE de uma linha de base de decoder de "vinculação rígida" (hard-tied), onde os pesos do decoder são fixos para serem exatamente iguais à matriz de embeddings de tokens ($W_D = W_E$). A vinculação rígida força as características a serem exatamente os embeddings de tokens, o que o artigo hipotetiza ser insuficiente para uma reconstrução de alta qualidade. O VASAE mantém o dicionário aprendível, mas o puxa em direção à geometria do vocabulário via perda de âncora.

3. Configuração Experimental

O método foi avaliado em dois modelos Transformer decoder-only:

• GPT-2-small: 12 camadas, $d_{model}=768$, tamanho do vocabulário $|V|=50.257$.
• Llama-3.1-8B: 32 camadas, $d_{model}=4096$, tamanho do vocabulário $|V|=128.256$.

Linhas de Base (Baselines):

1. SAE Padrão: Treinado apenas com perda de reconstrução.
2. Decoder de Vinculação Rígida (Hard-Tied Decoder): Pesos do decoder fixos aos embeddings de tokens; apenas o encoder é treinado.

Métricas:

• Reconstrução: Erro Quadrático Médio (MSE), Variância Explicada (VE) e recuperação de Entropia Cruzada (CE) após substituição do fluxo residual.
• Alinhamento: Taxa de alinhamento geométrico de características (porcentagem de características com score de alinhamento $s_i \ge 0,8$) e cobertura do vocabulário.
• Estudos de Caso: Inspeção qualitativa dos nomes de tokens intrínsecos atribuídos a características ativas em contextos específicos.

4. Principais Resultados

Qualidade de Reconstrução

O VASAE alcança métricas de reconstrução comparáveis ao SAE Padrão, superando significativamente a linha de base de Decoder de Vinculação Rígida.

• GPT-2-small: Tanto o VASAE quanto o SAE Padrão alcançam uma Variância Explicada (VE) de 0,965 e recuperação de CE de 0,975.
• Llama-3.1-8B: Tanto o VASAE quanto o SAE Padrão alcançam uma VE de 0,931 e recuperação de CE de 0,906.
• Linha de Base de Vinculação Rígida: Desempenha mal (ex: VE negativa no GPT-2-small), confirmando que forçar as características a serem exatamente os embeddings de tokens destrói a capacidade de reconstrução.

Alinhamento Geométrico

• GPT-2-small: Com $\lambda_{anchor} = 10^{-4}$, 89–94% das características nas camadas 0–10 excedem o corte de alinhamento ($s_i \ge 0,8$). O alinhamento é forte e consistente em camadas rasas e médias.
• Llama-3.1-8B: O alinhamento depende da profundidade. Com um coeficiente mais alto ($\lambda_{anchor} = 5 \times 10^{-3}$), a camada rasa (L0) atinge 92,8% de alinhamento. No entanto, o alinhamento cai significativamente em camadas mais profundas (ex: L31), sugerindo que a ancoragem de vocabulário é mais compatível com estados residuais próximos ao espaço de embedding.

Estudos de Caso

A análise qualitativa dos nomes de tokens intrínsecos revela coerência semântica:

• No GPT-2-small, características ativas em uma frase sobre "Baker Street" são atribuídas a nomes intrínsecos como located (localizado), Street (rua), corner (esquina) e around (ao redor).
• Exemplos adicionais mostram características alinhadas a adjetivos/advérbios (ex: perfect, feasible), termos relacionados a prêmios (ex: winner, praise) e termos de autoapresentação (ex: fan, name).
• No Llama-3.1-8B, a camada rasa mostra alinhamento semelhante relacionado a localização, enquanto as camadas profundas mostram menos consistência, espelhando as métricas de alinhamento quantitativas.

5. Contribuições e Significância

O artigo apresenta três contribuições principais:

1. Introdução do VASAE: Um objetivo de ancoragem de vocabulário suave que atribui nomes de tokens mais próximos às direções de decoder aprendidas do SAE durante o treinamento, em vez de post-hoc.
2. Avaliação da Ancoragem: Demonstração de que a ancoragem alinhada ao vocabulário preserva a qualidade de reconstrução comparável aos SAEs padrão, enquanto prova que vincular rigidamente o decoder aos embeddings de tokens é insuficiente para essa tarefa.
3. Análise Geométrica: Análise do alinhamento de características no espaço de embedding e cobertura do vocabulário, fornecendo evidências de que nomes de tokens intrínsecos atuam como identificadores derivados da geometria.

Significância:
Os autores argumentam que o VASAE complementa os métodos de interpretação post-hoc existentes ao conectar as características aprendidas aos nomes de tokens intrínsecos durante o treinamento. Isso fornece um "identificador indexado pelo vocabulário" para as características, oferecendo um ponto de partida para inspeção.

Crucialmente, o artigo mantém uma postura modesta sobre as implicações desses nomes:

• Não são Explicações: Um nome de token intrínseco não é uma explicação completa da característica. Ele não explica por que uma característica ativa ou seu papel causal na computação do modelo.
• Ferramenta de Indexação: Os nomes servem como ferramentas preliminares de indexação. A compreensão real requer análise adicional de como as características nomeadas participam de circuitos, como elas se conectam a outros componentes e como intervenções nelas afetam o comportamento.
• Limitações: O alinhamento é mais robusto no GPT-2-small e em camadas rasas/médias do Llama-3.1-8B. O método atualmente ancora apenas aos embeddings de entrada; para modelos com embeddings de entrada/saída não unificados (untied), a escolha da matriz de âncora permanece uma questão aberta.

Em resumo, o VASAE consegue unir com sucesso o hiato entre o espaço geométrico das características aprendidas e o espaço do vocabulário do modelo, permitindo que as características carreguem nomes de tokens intrínsecos sem comprometer sua utilidade para reconstrução.