Resumo Técnico: Aprendizado por Reforço de Controle (CRL)

1. Declaração do Problema

Autoencoders Esparsos (SAEs) emergiram como uma ferramenta poderosa para decompor ativações de modelos de linguagem em características interpretáveis e legíveis por humanos. No entanto, os métodos de análise existentes são amplamente estáticos: eles identificam quais características ativam em entradas específicas, mas falham em revelar quais características, quando amplificadas, realmente alteram as saídas do modelo. Técnicas atuais de engenharia de ativação frequentemente dependem de conjuntos de dados contrastivos para seleção de características ou armazenamento extensivo de ativações para otimização de coeficientes, carecendo de mecanismos de feedback adaptativo para ajustar estratégias de direcionamento com base na qualidade da geração.

O artigo aborda a lacuna entre a identificação estática de características e a intervenção dinâmica e interpretável. O desafio central é desenvolver um framework que não apenas selecione características de SAE para direcionar um modelo, mas também forneça registros de intervenção ao nível de token que expliquem por que uma característica foi selecionada e como ela afetou a saída, transformando assim o direcionamento em uma ferramenta de interpretabilidade mecânica.

2. Metodologia: Aprendizado por Reforço de Controle (CRL)

Os autores propõem o Aprendizado por Reforço de Controle (CRL), um framework que trata o controle de representações de transformadores como um Processo de Decisão de Markov (MDP) sobre características de SAE.

2.1. Formulação do MDP

• Estado ($S$): A observação é a ativação do fluxo residual ($x \in \mathbb{R}^d$) em uma camada específica $\ell$ para a posição do token atual.
• Ação ($A$): Um vetor binário que seleciona uma ou mais características de SAE para amplificar. A política seleciona as top-$k$ características (tipicamente $k=1$ para interpretabilidade).
• Transição ($P$): Determinística, governada pela passagem forward do transformador com a intervenção de característica selecionada aplicada.
• Recompensa ($R$): Recompensas específicas da tarefa baseadas na qualidade da saída final (por exemplo, correção em MMLU, recusa em tarefas de segurança).

2.2. CRL-Token e Direcionamento de Características

Em cada etapa de geração $t$, o framework realiza intervenções ao nível de token:
$$\tilde{x}_t = x_t + c \cdot a_t W_{dec}$$
Onde $a_t$ é o vetor de seleção binária esparsa, $c$ é o coeficiente de direcionamento e $W_{dec}$ é o decodificador do SAE. A rede de política $\pi_\theta$ mapeia ativações do fluxo residual para logits de seleção de características, implementada como um MLP de 2 camadas.

2.3. Mascaramento Adaptativo de Características (AFM)

Para evitar que a política colapse em um conjunto estreito de características frequentemente ativas, o CRL introduz o Mascaramento Adaptativo de Características.

• Uma máscara booleana por amostra restringe as escolhas de seleção de características.
• A máscara começa com um pequeno subconjunto de características e expande progressivamente à medida que a geração avança, incorporando características que foram naturalmente ativadas em etapas anteriores.
• Este mecanismo incentiva a descoberta diversificada de características enquanto mantém a interpretabilidade de intervenções de característica única.

2.4. Otimização

A política e uma rede crítica (estimando valores de estado) são treinadas usando Otimização de Política Próxima (PPO).

• Coeficiente de Direcionamento ($c$): Determinado adaptativamente pela média das magnitudes de ativação de amostras de treinamento previstas corretamente, eliminando a necessidade de ajuste manual entre camadas.
• Rede Crítica: Um MLP de 2 camadas que estima a função valor $V_\phi(s)$, permitindo a separação das limitações da política dos erros de estimativa de valor durante a análise.

3. Contribuições Principais

O artigo delineia quatro contribuições primárias:

1. Registros de Intervenção Interpretáveis ao Nível de Token: Um framework que produz registros identificando exatamente qual característica de SAE foi selecionada em cada token e seu efeito na saída, permitindo análise comportamental de alta granularidade.
2. Ferramentas de Diagnóstico:
   • Rastreamento de Pontos de Ramificação: Identifica pontos críticos de decisão onde prefixos idênticos produzem resultados diferentes com base na seleção de características.
   • Análise de Trajetória da Crítica: Utiliza a função valor aprendida para distinguir entre falhas da política (seleção de características) e falhas da crítica (estimativa de recompensa).
3. Mascaramento Adaptativo de Características (AFM): Um mecanismo que equilibra exploração e exploração para garantir descoberta diversificada de características sem sacrificar a interpretabilidade de característica única.
4. Caracterização Empírica de Semântica de Características: Demonstra que camadas iniciais codificam estrutura sintática enquanto camadas posteriores capturam semântica abstrata, e que diferentes tarefas exibem padrões de gargalo distintos.

4. Resultados Experimentais

O framework foi avaliado no Gemma 2 2B (e LLaMA-3.1 8B) usando SAEs Gemma Scope e LLaMA Scope em cinco benchmarks: MMLU, BBQ, GSM8K, HarmBench e XSTest.

• Ganhos de Desempenho: O CRL alcançou melhorias em todos os benchmarks. Ganhos notáveis incluem:
  • HarmBench: +7,66 pontos percentuais (41,46% $\to$ 49,12%).
  • BBQ (Ambíguo): +5,69 pontos percentuais.
  • MMLU: +3,31 pontos percentuais.
  • GSM8K: +1,03 pontos percentuais.
• Complementaridade: O CRL fornece ganhos aditivos quando aplicado sobre modelos Ajustados por Fine-Tuning Supervisionado (SFT) (por exemplo, a precisão do MMLU aumentou de 55,2% para 56,1% com SFT + CRL).
• Semântica Específica por Camada:
  • Camada 10: Características capturam propriedades concretas e sintáticas (por exemplo, "notação matemática", "referências pronominais").
  • Camada 20: Características capturam semântica abstrata e contextual (por exemplo, "estrutura de derivação lógica", "provas matemáticas").
  • As camadas de intervenção ótimas variam conforme a tarefa (por exemplo, Camada 5 para BBQ, Camada 24 para MMLU).
• Análise da Crítica:
  • BBQ: Mostra coordenação efetiva entre política e crítica com discriminação clara de amostras.
  • MMLU: Exibe um "gargalo da crítica" onde amostras corrigidas e equivocadas são quase indistinguíveis, sugerindo que a crítica luta para estimar valor em tarefas de raciocínio complexo.
  • HarmBench vs. XSTest: HarmBench mostra lacunas claras entre amostras corretas/incorretas, enquanto XSTest mostra padrões de estimativa invertidos, indicando desafios dependentes da tarefa na estimativa de recompensa.

5. Significância e Alegações

O artigo posiciona o CRL como um avanço significativo na interpretabilidade mecânica. Sua significância primária reside na mudança de paradigma da análise estática de características (observar o que ativa) para probes de intervenção dinâmica (observar o que altera o modelo).

• Interpretabilidade: Ao impor seleção de característica única e registrar intervenções, o CRL fornece uma visão "mecânica" do comportamento do modelo, revelando que características como "operações matemáticas" podem generalizar lexicalmente para unidades quantitativas como "bebês" em problemas de palavras.
• Capacidade Diagnóstica: O framework serve como ferramenta de diagnóstico tanto para o modelo quanto para o SAE. Se uma política colapsar para características semanticamente incoerentes apesar do AFM, isso sinaliza uma deficiência no SAE (por exemplo, conceitos ausentes ou polissemia) em vez de uma falha do método de controle.
• Praticidade: O CRL opera com baixa sobrecarga de inferência (8–9% de latência) e não requer retreinamento dos pesos do modelo base, oferecendo um mecanismo de aprimoramento pós-hoc e leve.

Os autores concluem que o direcionamento de características aprendido, quando acoplado ao aprendizado por reforço, estabelece uma ferramenta prática para entender e controlar modelos de linguagem, complementando a análise estática com insights dinâmicos e específicos da tarefa.