Resumen Técnico: SAEs de Promedio de Turno para el Descubrimiento de Características y la Atribución de Contexto Largo

1. Planteamiento del Problema

Los Autoencoders Dispersos (SAEs, por sus siglas en inglés) han emergido como una herramienta primaria para extraer características interpretables de los modelos de lenguaje, facilitando el análisis de circuitos, la dirección conductual y el monitoreo de seguridad. Sin embargo, las arquitecturas de SAE estándar operan sobre activaciones de tokens individuales. Esta granularidad por token crea dos cuellos de botella significativos para la interpretabilidad a escala:

1. Limitaciones de Escalamiento: El número de características activas escala linealmente con la longitud del contexto. En escenarios de contexto largo o conversaciones de múltiples turnos, el volumen de datos resultante hace que el análisis posterior (como los grafos de atribución) sea engorroso y difícil de inspeccionar.
2. Pérdida de Contexto de Alto Nivel: Muchos objetivos de interpretabilidad implican comprender por qué un modelo produjo una respuesta específica, lo que requiere la identificación de propiedades de alto nivel (p. ej., tema, persona, rasgos de seguridad) que están distribuidas a través de muchos tokens. Los SAEs estándar por token suelen aprender detectores de bajo nivel para palabras específicas o patrones sintácticos. Debido a que el objetivo de entrenamiento reconstruye cada posición de forma independiente, no hay presión para que el modelo separe estas características de alto nivel y de amplio espectro del ruido a nivel de token.

2. Metodología

2.1 SAEs de Promedio de Turno

La propuesta central es entrenar SAEs sobre el estado oculto promedio a través de un intervalo de tokens (específicamente, un solo turno de Humano o de Asistente) en lugar de sobre activaciones de tokens individuales.

• Transformación de Entrada: Dado un turno $T$ con tokens $x_t$, la entrada es el estado oculto promedio $\bar{x} = \frac{1}{|T|} \sum_{t \in T} x_t$.
• Arquitectura: La arquitectura del SAE permanece idéntica a las implementaciones estándar (usando BatchTopK para la dispersión). Dado que el promedio de vectores en $\mathbb{R}^{d_{model}}$ conserva la misma dimensionalidad que los vectores individuales, no se requieren cambios estructurales.
• Intuición: El promedio actúa como un filtro de paso bajo. Suaviza los detalles de alta frecuencia específicos del token (como la elección de palabras específicas o marcadores sintácticos) mientras preserva los componentes de baja frecuencia que corresponden a conceptos amplios como el tema, el estilo y los rasgos de comportamiento.

2.2 Arquitecturas Híbridas

Para abordar la limitación de que el promedio descarta el detalle por token, los autores proponen dos enfoques híbridos para capturar tanto la estructura de alto nivel como los residuales de grano fino:

1. SAE Desacoplado: Entrena dos SAEs independientes. Uno reconstruye el promedio del turno ($\bar{x}$), y el otro reconstruye el residual ($x_t - \bar{x}$). Aunque esto separa las señales, los modelos están causalmente desconectados, lo que complica el análisis de atribución.
2. SAE Anidado (estilo Matryoshka): Entrena un único SAE con un codificador y decodificador compartidos pero con una función de pérdida especializada. El espacio de características se divide en una partición interna $[0:h]$ y una partición externa $[h:d_{sae}]$.
   • Función de Pérdida: $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$.
   • La partición interna se ve forzada a reconstruir el promedio del turno, mientras que el conjunto completo de características reconstruye las activaciones por token. Esto permite que un solo modelo aprenda simultáneamente características de alto nivel del turno y detalles a nivel de token.

2.3 Adaptación del Grafo de Atribución

Los autores adaptan los grafos de atribución al nivel de turno. En lugar de que los nodos representen características en posiciones de tokens específicos, los nodos representan características activas en un turno y capa específicos.

• Pesos de las Aristas: Calculados promediando la atribución del gradiente a través del turno de origen y normalizando por la longitud del turno ($1/|T_A|$ y $1/|T_B|$) para asegurar que los pesos sean agnósticos a la longitud del turno.
• Validación de Intervención: Los pesos de las aristas se validan perturbando las características de origen a lo largo de sus direcciones de decodificación en todos los tokens de un turno y midiendo el cambio en las pre-activaciones de la característica objetivo.

3. Contribuciones Clave

1. SAEs de Promedio de Turno: Demostraron que el entrenamiento en estados ocultos promedio produce características que capturan las características de alto nivel del turno (tema, estilo, rasgos de seguridad) de manera más completa que las características por token.
2. Arquitectura de SAE Anidado: Introdujeron una arquitectura de modelo único que aprende conjuntamente características de promedio de turno y por token usando una pérdida tipo Matryoshka, logrando una alta calidad de reconstrucción para ambos componentes de la señal.
3. Atribución Escalable: Adaptaron los grafos de atribución al nivel de turno, permitiendo un análisis compacto de conversaciones de múltiples turnos. Los autores validaron estos grafos mediante experimentos de intervención y métricas de suficiencia (completitud y reemplazo).

4. Resultados

4.1 Calidad de la Característica y Reconstrucción

• Reconstrucción: El SAE Anidado logra un Error Cuadrático Medio Normalizado (nMSE) por token de 0.167, igualando de cerca al SAE puro por token (0.162), mientras que también logra un nMSE de promedio de turno de 0.116. En contraste, un SAE estándar por token tiene un desempeño deficiente en entradas de promedio (nMSE 0.355), indicando que el promedio está fuera de la distribución para los modelos entrenados por token.
• Descubrimiento Cualitativo: En ejemplos que involucran errores matemáticos, estilos técnicos entusiastas y cumplimiento regulatorio, las características de promedio de turno identificaron con éxito propiedades de comportamiento de alto nivel (p. ej., "respuestas incorrectas", "estilo técnico entusiasta") que las características por token pasaron por alto o capturaron solo parcialmente.
• Generalización de Contexto Largo: Los SAEs de promedio de turno entrenados en turnos de chat (longitud promedio ~150 tokens) se generalizaron con éxito a documentos 150 veces más largos (24k–32k tokens) de LongBench v2, identificando conceptos de alto nivel como "macroeconomía" y "gráficos por computadora" que las características por token no pudieron detectar.

4.2 Evaluación Cuantitativa

Usando un marco de LLM-como-juez para comparar las descripciones de las características contra resúmenes estructurados de los turnos:

• Discriminación (emparejamiento de 10 vías): Las características por token (max-pooling) superaron a las características de promedio de turno (95.0% vs. 73.9%), ya que estas últimas dependen de detalles específicos de nivel de token para distinguir textos.
• Cobertura (ranking de 5 vías y comparaciones por pares): Las características de promedio de turno superaron significativamente a las por token (ganando el 87.9% de las comparaciones por pares). Proporcionaron descripciones más completas de las características de alto nivel de un turno.
• Métrica de Embedding: Las características de promedio de turno alcanzaron la mayor similitud de coseno (0.663) con los campos de resumen estructurado, confirmando su superior cobertura de contenido semántico.

4.3 Análisis de Atribución

• Compacidad del Grafo: Los grafos de promedio de turno redujeron el conteo de nodos de ~128K (por token) a ~5K para una conversación de 10 turnos, haciendo que la inspección sea factible.
• Casos de Estudio:
  • Atribución de Prompt de Sistema: Los grafos de promedio de turno identificaron correctamente qué párrafos del prompt de sistema influyeron en consultas de usuario específicas, mostrando cambios de atribución comparables a los grafos por token.
  • Rastreo de Degeneración: En una conversación de 14 turnos donde el modelo degeneró en texto repetitivo, el grafo de promedio de turno trazó una cadena causal coherente: "Descargos de seguridad de IA" $\to$ "Armas nucleares" $\to$ "Detección de prompts adversarios" $\to$ "Salida degenerada". Los grafos por token fallaron en encontrar la característica de detección adversaria, trazando en su lugar hacia conceptos no relacionados como "compuestos químicos".
• Intervención y Suficiencia: Los pesos de las aristas mostraron una correlación significativa con los efectos de intervención. El SAE Anidado logró los puntajes de suficiencia de "completitud" y "reemplazo" más altos, superando tanto al SAE puro como a los agregados post-hoc de características por token.

5. Significado y Reivindicaciones

El artículo afirma que los SAEs de promedio de turno hacen que las técnicas de interpretabilidad sean prácticas en contextos largos. Al cambiar la unidad de análisis de tokens a intervalos, el método:

• Produce características que describen características de alto nivel (tema, persona, seguridad) de manera más completa que las características a nivel de token.
• Permite la construcción de grafos de atribución compactos e interpretables para conversaciones de múltiples turnos que de otro modo serían demasiado grandes para analizar.
  les permite demostrar que modificar la señal de entrenamiento (mediante el promedio) puede obligar a los SAEs a aprender características a la granularidad deseada (conceptos de alto nivel) independientemente del conteo de tokens.
• Demuestra que estas características se generalizan a intervalos mucho más largos que aquellos vistos durante el entrenamiento, lo que sugiere robustez para el análisis de modelos de frontera.

Los autores señalan las limitaciones, incluyendo la escala actual de los experimentos (modelo 7B, corpus LMSYS-Chat-1M) y la dependencia de LLMs para las métricas de evaluación, identificando la aplicación a modelos de frontera y la creación de benchmarks de verdad fundamental como trabajos futuros necesarios.