reward-lens: A Mechanistic Interpretability Library for Reward Models

1. Declaração do Problema

O Aprendizado por Reforço a partir de Feedback Humano (RLHF) depende de Modelos de Recompensa (RMs) para moldar o comportamento de Grandes Modelos de Linguagem (LLMs). Se um modelo de recompensa codificar preferências incorretas, a política resultante herdará e amplificará esses erros. Embora a avaliação comportamental de RMs (por exemplo, RewardBench) seja madura, a interpretabilidade mecanística para RMs é virtualmente inexistente.

As ferramentas de interpretabilidade existentes (por exemplo, TransformerLens, Logit Lens, Atribuição Direta de Logit) são projetadas para LLMs generativos. Elas dependem de projetar estados intermediários do modelo na matriz de desincorporação de vocabulário ($W_U$) para analisar distribuições de tokens. No entanto, modelos de recompensa substituem $W_U$ por uma cabeça de regressão escalar ($w_r$). Essa diferença estrutural torna os primitivos padrão inaplicáveis, criando uma "lacuna estrutural" onde os mecanismos internos de modelos de recompensa críticos para a segurança permanecem uma caixa preta.

2. Metodologia: A Visão Centrada em $w_r$

O artigo apresenta o reward-lens, uma biblioteca de código aberto que adapta primitivos de interpretabilidade mecanística para modelos de recompensa. A principal insight teórica é que o vetor de pesos da cabeça de recompensa, $w_r$, serve como o eixo natural para todas as questões de interpretabilidade.

A biblioteca é construída sobre a observação de que a recompensa final $r$ é um funcional linear do fluxo residual final $h_{final}$:
$$r = w_r^\top h_{final} + b_r$$
Como o fluxo residual se decompõe aditivamente através das camadas ($h_{final} = h_{embed} + \sum (attn^{(\ell)} + mlp^{(\ell)})$), a recompensa pode ser exatamente decomposta como:
$$r = w_r^\top h_{embed} + \sum (w_r^\top attn^{(\ell)} + w_r^\top mlp^{(\ell)}) + b_r$$

Com base nisso, a biblioteca implementa os seguintes primitivos substituindo $W_U$ por $w_r$:

• Reward Lens: Projeta fluxos residuais intermediários na camada $\ell$ sobre $w_r$ para rastrear a "cristalização" de sinais de preferência através das camadas.
• Atribuição de Componentes: Calcula a contribuição de cada subcamada (Atenção/MLP) para a recompensa escalar final via $w_r^\top \cdot \text{contribuição}$.
• Patching de Ativação Contrastiva: Troca ativações entre conclusões preferidas e não preferidas e mede a mudança em $w_r^\top h_{final}$. Isso é nativo de RMs devido à estrutura contrastiva inerente dos dados de preferência.
• Atribuição de Características SAE: Usa Autoencoders Esparsos (SAEs) para decompor fluxos residuais em características, calculando o alinhamento de características como $w_r^\top d_i$ (onde $d_i$ é um vetor de dicionário).

Agnosticismo de Arquitetura: A biblioteca utiliza uma abordagem nativa do HuggingFace com um protocolo de adaptador de 10 métodos. Isso permite envolver modelos padrão (Llama, Mistral, Gemma-2) e modelos complexos multi-objetivo (como ArmoRM com gateamento MoE) sem cirurgia de pesos.

Extensões Fundamentadas na Teoria: A biblioteca operacionaliza resultados recentes da teoria de alinhamento em código:

• Índice de Distorção: Prevê quais dimensões de qualidade estão subcobertas por sondas de avaliação (risco de exploração de recompensa).
• Patching Consciente de Divergência: Detecta quando o patching induz ativações fora da distribuição.
• Detector de Cascata de Desalinhamento: Identifica modos de falha correlacionados entre diferentes dimensões de desalinhamento.
• Analisador de Conflito de Termos de Recompensa: Verifica se diferentes objetivos de recompensa estão alinhados, ortogonais ou em conflito.
• Análise de Vetor de Conceito: Extrai direções de conceito lineares e mede sua resposta dose-resposta causal na recompensa.

3. Principais Contribuições

1. A Biblioteca (reward-lens): A primeira ferramenta instalável e agnóstica de arquitetura para interpretabilidade mecanística de modelos de recompensa. Unifica visões observacionais (atribuição) e causais (patching) sob um único framework algébrico.
2. Enquadramento Unificador: Estabelece que todo primitivo de interpretabilidade útil para um RM é uma projeção sobre ou uma decomposição ao longo de $w_r$.
3. Validação Empírica: Aplicação sistemática a dois modelos de produção: Skywork-Reward-Llama-3.1-8B-v0.2 (cabeça escalar única) e ArmoRM-Llama3-8B-v0.1 (cabeça multi-objetivo).
4. Descoberta Empírica Negativa: O artigo relata explicitamente que a atribuição linear não prevê importância causal em modelos de recompensa.

4. Principais Resultados e Descobertas

Os autores avaliaram a biblioteca em ~695 pares de preferência do RewardBench por modelo.

• Profundidade de Cristalização:
  • Skywork: Preferências cristalizam uniformemente tarde (profundidade de camada fracionária $\approx 0,90$), sugerindo um mecanismo unificado montado próximo à saída.
  • ArmoRM: Preferências cristalizam mais cedo e com maior variância, refletindo sua arquitetura multi-objetivo com gateamento MoE.
• Atribuição vs. Causalidade (A Lacuna de Fidedignidade):
  • Atribuição Observacional: As camadas finais de MLP (especificamente mlp_L31) dominam o sinal de recompensa em ambos os modelos.
  • Patching Causal: Componentes de camadas iniciais (por exemplo, mlp_L0, attn_L0) são frequentemente os mais críticos para a recompensa final.
  • Correlação: A correlação de posto de Spearman entre magnitude de atribuição e efeito de patching é negativa ou próxima de zero (Skywork: $\rho = -0,256$; ArmoRM: $\rho = -0,027$).
  • Interpretação: Isso não é um erro, mas uma característica de fluxos residuais redundantes. Camadas tardias podem ser ablatadas porque caminhos paralelos compensam, enquanto camadas iniciais são "estruturalmente essenciais".
• Perfis de Exploração de Recompensa:
  • Skywork penaliza comprimento, confiança e vieses de formatação (diferente de modelos típicos de RLHF que os recompensam).
  • ArmoRM recompensa confiança e formatação, mas penaliza sycophancy (servilismo).
• Alinhamento de Conceitos:
  • Skywork mostra forte alinhamento linear com conceitos como "concordância" (inclinação $\approx 1,0$).
  • ArmoRM mostra alinhamento fraco ($\approx 1\%$ de magnitude) com conceitos lineares, sugerindo uma geometria de recompensa mais complexa e não linear.

5. Significado

• Mudança de Paradigma: O artigo argumenta que os pesquisadores devem parar de tratar a atribuição observacional como um proxy para importância causal em modelos de recompensa. A lacuna entre as duas é estrutural e deve ser medida, não ignorada.
• Ferramentas de Segurança: Ao fornecer ferramentas para detectar "distorção" (dimensões subcobertas) e "cascatas de desalinhamento", a biblioteca permite monitoramento durante o treinamento. Isso muda a auditoria de modelos de recompensa de uma atividade post-hoc para uma salvaguarda em tempo real durante o loop de treinamento de RLHF.
• Editabilidade: O framework identifica características e conceitos específicos que impulsionam sinais de recompensa, abrindo caminho para edição de modelos de recompensa guiada por interpretabilidade (por exemplo, suprimir um conceito de "servilismo" da cabeça de recompensa) antes do treinamento da política a jusante.
• Ciência Aberta: Ao liberar o reward-lens, os autores democratizam o acesso aos mecanismos internos de modelos de recompensa, um passo crítico para garantir a segurança de sistemas de IA de fronteira.

Em conclusão, o reward-lens preenche uma lacuna crítica na segurança de IA ao adaptar a interpretabilidade mecanística para a natureza de saída escalar dos modelos de recompensa, revelando que sua lógica interna é frequentemente contra-intuitiva (atribuição tardia $\neq$ causalidade inicial) e fornecendo as ferramentas necessárias para diagnosticar e corrigir falhas de alinhamento antes que elas se propaguem para modelos de política.