reward-lens: A Mechanistic Interpretability Library for Reward Models

Aquí se presenta un resumen técnico detallado del artículo "reward-lens: A Mechanistic Interpretability Library for Reward Models."

1. Planteamiento del Problema

El Aprendizaje por Refuerzo a partir de Retroalimentación Humana (RLHF) depende de Modelos de Recompensa (RMs) para moldear el comportamiento de los Modelos de Lenguaje Grandes (LLM). Si un modelo de recompensa codifica preferencias incorrectas, la política resultante hereda y amplifica esos errores. Si bien la evaluación conductual de los RMs (por ejemplo, RewardBench) es madura, la interpretabilidad mecánica para los RMs es prácticamente inexistente.

Las herramientas de interpretabilidad existentes (por ejemplo, TransformerLens, Logit Lens, Direct Logit Attribution) están diseñadas para LLMs generativos. Se basan en proyectar estados intermedios del modelo sobre la matriz de desincrustación del vocabulario ($W_U$) para analizar las distribuciones de tokens. Sin embargo, los modelos de recompensa reemplazan $W_U$ con una cabeza de regresión escalar ($w_r$). Esta diferencia estructural hace que los primitivos estándar sean inaplicables, creando una "brecha estructural" donde los mecanismos internos de los modelos de recompensa críticos para la seguridad permanecen como una caja negra.

2. Metodología: La Perspectiva Centrada en $w_r$

El artículo introduce reward-lens, una biblioteca de código abierto que adapta los primitivos de interpretabilidad mecánica para modelos de recompensa. La idea teórica central es que el vector de pesos de la cabeza de recompensa, $w_r$, sirve como el eje natural para todas las preguntas de interpretabilidad.

La biblioteca se basa en la observación de que la recompensa final $r$ es un funcional lineal del flujo residual final $h_{final}$:
$$r = w_r^\top h_{final} + b_r$$
Dado que el flujo residual se descompone aditivamente a través de las capas ($h_{final} = h_{embed} + \sum (attn^{(\ell)} + mlp^{(\ell)})$), la recompensa puede descomponerse exactamente como:
$$r = w_r^\top h_{embed} + \sum (w_r^\top attn^{(\ell)} + w_r^\top mlp^{(\ell)}) + b_r$$

Basándose en esto, la biblioteca implementa los siguientes primitivos reemplazando $W_U$ con $w_r$:

• Reward Lens: Proyecta los flujos residuales intermedios en la capa $\ell$ sobre $w_r$ para rastrear la "cristalización" de las señales de preferencia a través de las capas.
• Atribución de Componentes: Calcula la contribución de cada subcapa (Atención/MLP) a la recompensa escalar final mediante $w_r^\top \cdot \text{contribución}$.
• Parcheo de Activación Contrastivo: Intercambia activaciones entre completaciones preferidas y no preferidas y mide el cambio en $w_r^\top h_{final}$. Esto es nativo de los RMs debido a la estructura contrastiva inherente de los datos de preferencia.
• Atribución de Características SAE: Utiliza Autoencoders Esparsos (SAE) para descomponer los flujos residuales en características, calculando la alineación de características como $w_r^\top d_i$ (donde $d_i$ es un vector de diccionario).

Agnosticismo Arquitectónico: La biblioteca utiliza un enfoque nativo de HuggingFace con un protocolo de adaptador de 10 métodos. Esto le permite envolver modelos estándar (Llama, Mistral, Gemma-2) y modelos complejos multi-objetivo (como ArmoRM con enrutamiento MoE) sin cirugía de pesos.

Extensiones Fundamentadas en Teoría: La biblioteca operacionaliza resultados recientes de la teoría de alineación en código:

• Índice de Distorsión: Predice qué dimensiones de calidad están subcobiertas por las sondas de evaluación (riesgo de hacking de recompensas).
• Parcheo Consciente de la Divergencia: Detecta cuándo el parcheo induce activaciones fuera de distribución.
• Detector de Cascadas de Desalineación: Identifica modos de fallo correlacionados a través de diferentes dimensiones de desalineación.
• Analizador de Conflictos de Términos de Recompensa: Verifica si diferentes objetivos de recompensa están alineados, ortogonales o en conflicto.
• Análisis de Vectores de Concepto: Extrae direcciones de concepto lineales y mide su respuesta dosis-causal en la recompensa.

3. Contribuciones Clave

1. La Biblioteca (reward-lens): La primera herramienta instalable y agnóstica a la arquitectura para la interpretabilidad mecánica de modelos de recompensa. Unifica las visiones observacional (atribución) y causal (parcheo) bajo un único marco algebraico.
2. Marco Unificador: Establece que cada primitivo de interpretabilidad útil para un RM es una proyección sobre o una descomposición a lo largo de $w_r$.
3. Validación Empírica: Aplicación sistemática a dos modelos de producción: Skywork-Reward-Llama-3.1-8B-v0.2 (cabeza escalar única) y ArmoRM-Llama3-8B-v0.1 (cabeza multi-objetivo).
4. Hallazgo Empírico Negativo: El artículo reporta explícitamente que la atribución lineal no predice la importancia causal en los modelos de recompensa.

4. Resultados y Hallazgos Clave

Los autores evaluaron la biblioteca en $\sim$695 pares de preferencia de RewardBench por modelo.

• Profundidad de Cristalización:
  • Skywork: Las preferencias cristalizan uniformemente tarde (profundidad de capa fraccional $\approx 0.90$), lo que sugiere un mecanismo unificado ensamblado cerca de la salida.
  • ArmoRM: Las preferencias cristalizan antes y con mayor varianza, reflejando su arquitectura multi-objetivo con enrutamiento MoE.
• Atribución vs. Causalidad (La Brecha de Fidelidad):
  • Atribución Observacional: Las últimas capas MLP (específicamente mlp_L31) dominan la señal de recompensa en ambos modelos.
  • Parcheo Causal: Los componentes de capas tempranas (por ejemplo, mlp_L0, attn_L0) suelen ser los más críticos para la recompensa final.
  • Correlación: La correlación de rango de Spearman entre la magnitud de la atribución y el efecto del parcheo es negativa o cercana a cero (Skywork: $\rho = -0.256$; ArmoRM: $\rho = -0.027$).
  • Interpretación: Esto no es un error, sino una característica de los flujos residuales redundantes. Las capas tardías pueden ser ablatadas porque caminos paralelos las compensan, mientras que las capas tempranas son "estructurales".
• Perfiles de Hacking de Recompensa:
  • Skywork penaliza la longitud, la confianza y los sesgos de formato (a diferencia de los modelos RLHF típicos que los recompensan).
  • ArmoRM recompensa la confianza y el formato pero penaliza la sycophancy (adulación).
• Alineación de Conceptos:
  • Skywork muestra una fuerte alineación lineal con conceptos como "acuerdo" (pendiente $\approx 1.0$).
  • ArmoRM muestra una alineación débil ($\approx 1\%$ de magnitud) con conceptos lineales, lo que sugiere una geometría de recompensa más compleja y no lineal.

5. Significado

• Cambio de Paradigma: El artículo argumenta que los investigadores deben dejar de tratar la atribución observacional como un proxy para la importancia causal en los modelos de recompensa. La brecha entre ambas es estructural y debe medirse, no ignorarse.
• Herramientas de Seguridad: Al proporcionar herramientas para detectar "distorsión" (dimensiones subcobiertas) y "cascadas de desalineación", la biblioteca permite la monitoreo durante el entrenamiento. Esto cambia la auditoría de modelos de recompensa de una actividad post-hoc a una salvaguarda en tiempo real durante el bucle de entrenamiento RLHF.
• Editabilidad: El marco identifica características y conceptos específicos que impulsan las señales de recompensa, allanando el camino para la edición de modelos de recompensa guiada por interpretabilidad (por ejemplo, suprimir un concepto de "adulación" de la cabeza de recompensa) antes del entrenamiento de la política aguas abajo.
• Ciencia Abierta: Al lanzar reward-lens, los autores democratizan el acceso a los mecanismos internos de los modelos de recompensa, un paso crítico para garantizar la seguridad de los sistemas de IA de vanguardia.

En conclusión, reward-lens llena un vacío crítico en la seguridad de la IA al adaptar la interpretabilidad mecánica a la naturaleza de salida escalar de los modelos de recompensa, revelando que su lógica interna a menudo es contra intuitiva (atribución tardía $\neq$ causalidad temprana) y proporcionando las herramientas necesarias para diagnosticar y corregir fallos de alineación antes de que se propaguen a los modelos de política.