Towards Isolated Interventions via Almost Orthogonal Features in Language Models

Resumen Técnico: Hacia intervenciones aisladas mediante características casi ortogonales en modelos de lenguaje

Planteamiento del problema
Una premisa central de la interpretabilidad mecánica es que los conceptos significativos en los modelos de lenguaje (LM) corresponden a características lineales en el espacio de activación. Para que estas características permitan intervenciones causales fiables, la manipulación de una característica no debería alterar sustancialmente los efectos de otras. Sin embargo, en la práctica, las características aprendidas suelen estar entrelazadas. Este "interferencia de características" ocurre cuando las características están fuertemente alineadas en el espacio de representación, lo que provoca que las intervenciones localizadas (por ejemplo, dirigir un modelo) tengan efectos secundarios no deseados, desdibujando las contribuciones causales y limitando la capacidad de intervención. Este problema se ve exacerbado por el fenómeno de la superposición, donde los modelos representan más características de las que tienen dimensiones, lo que genera patrones de interferencia difíciles de controlar.

Metodología
Motivados por el principio de Mecanismos Causales Independientes (ICM), que postula que los sistemas complejos pueden descomponerse en módulos autónomos que no se influyen entre sí, los autores proponen restringir las características internas para que sean casi ortogonales. La metodología consiste en un proceso de ajuste fino de dos pasos:

1. Entrenamiento de SAE Ortogonal: Los autores entrenan un Autoencoder Disperso (SAE) sobre las activaciones del flujo residual de un modelo de lenguaje. A diferencia de los SAE estándar, introducen una penalización de regularización de ortogonalidad explícita en la matriz de decodificación $D$. La función de pérdida incluye un término $\lambda \| \text{tril}(D^\top D) \|_F^2$, que penaliza los elementos fuera de la diagonal de la matriz de Gram, imponiendo así una baja coherencia propia (cercanía a la ortogonalidad) entre los átomos de las características.
2. Adaptación del LM: Una vez entrenado el SAE y fijados sus pesos, los autores lo insertan en un Transformer pre-entrenado (específicamente un modelo Gemma 2 2B) en una capa intermedia (capa 12 de 26). Posteriormente, realizan una adaptación de bajo rango (LoRA) en el modelo de lenguaje para adaptarlo a la nueva geometría de las características preservando el rendimiento general.

Marco Teórico
El artículo formaliza la relación entre la interferencia de características y la capacidad de intervención. Al modelar la representación del LM como una combinación lineal dispersa de características ($X = DZ$) y asumir una hipótesis lineal para el mapeo de salida, los autores derivan una cota superior en la discrepancia esperada entre la salida real y la salida del modelo tras una intervención.

• Demuestran que el desajuste post-intervención está controlado por un término de línea base y un término específico de la intervención proporcional a la coherencia propia del diccionario de características, denotado como $\mu(D)$.
• Específicamente, la discrepancia esperada está acotada superiormente por un término que involucra $\delta^2 (d-1)\mu(D)$, donde $\delta$ es la magnitud de la intervención. Esto demuestra teóricamente que reducir la coherencia propia (aumentar la ortogonalidad) limita directamente la propagación de la interferencia, mejorando así el aislamiento de las intervenciones.

Contribuciones clave

1. Cota Teórica: Los autores proporcionan una cota superior teórica que relaciona la capacidad de intervención en el flujo residual con la interferencia de características, vinculando explícitamente la discrepancia en las salidas del modelo con la coherencia propia del diccionario de características.
2. Pipeline de Ajuste Fino Ortogonal: Demuestran un pipeline que ajusta un LM alrededor de un SAE con un decodificador casi ortogonal. Muestran que este enfoque mantiene el rendimiento y los puntajes de interpretabilidad comparables a los SAE sin penalización en tareas de razonamiento matemático.
3. Intervenciones Aisladas: Los autores realizan las primeras intervenciones aisladas en conceptos específicos dentro de un SAE al limitar explícitamente los efectos de desbordamiento (spill-over). Logran sustituir con éxito conceptos (por ejemplo, cambiar el nombre "Jerry" por "Aquaman") mientras preservan las capacidades de razonamiento del modelo y minimizan los cambios no deseados en otras características.

Resultados

• Ortogonalidad vs. Rendimiento: La evaluación empírica en el conjunto de datos GSM8K muestra que aumentar la penalización de ortogonalidad $\lambda$ disminuye monótonamente la pérdida de ortogonalidad (mejorando la separación de características) sin degradar significativamente la precisión del modelo en tareas de razonamiento matemático. El rendimiento de los modelos con penalizaciones de ortogonalidad estrictas sigue siendo competitivo con el de la línea base sin penalización.
• Eficacia de la Intervención: En experimentos que implican la sustitución de nombres de pila (por ejemplo, "Jerry" por "Aquaman") y otros conceptos (animales, lugares), los modelos con mayores penalizaciones de ortogonalidad demostraron un aislamiento significativamente mejor.
  • Sustitución Correcta: A medida que la penalización aumentaba, el modelo incluía con mayor frecuencia el nombre sustituido correcto (por ejemplo, recuperando "Aquaman" ~70.9% de las veces con $\lambda=10^{-4}$ frente a ~60.1% sin penalización).
  • Reducción del Desbordamiento: Los modelos con penalizaciones de ortogonalidad fueron más exitosos al eliminar el concepto original sin retener erróneamente el mismo en la salida.
• Interpretabilidad: La interpretabilidad de las características (medida por la capacidad de un LLM para emparejar correctamente las explicaciones de las características con fragmentos de texto) se mantuvo constante a través de las diferentes penalizaciones de ortogonalidad, lo que indica que imponer la ortogonalidad no sacrifica la comprensibilidad humana de las características.

Significado y Reivindicaciones
El artículo afirma que la entrelazada de características es un inhibidor primario de las intervenciones aisladas en los modelos de lenguaje. Al relacionar formalmente la capacidad de intervención con el principio ICM y demostrar que la regularización de la ortogonalidad reduce la interferencia de características, los autores sugieren una ruta fundamentada para mejorar la identificabilidad y atomicidad de las características de los SAE sin sacrificar la utilidad descendente.

Los autores enmarcan modestamente su trabajo como un paso hacia sistemas de IA más fiables y predecibles. Argumentan que su enfoque permite modificaciones precisas y localizadas de las representaciones internas, lo cual es crucial para entender cómo componentes específicos contribuyen a los resultados. Reconocen que, si bien su método mejora la estructura y el control de las representaciones aprendidas, introduce preocupaciones de doble uso relacionadas con la dirección precisa de las salidas del modelo. Sin embargo, señalan que su enfoque requiere acceso directo a los componentes internos del modelo y experiencia técnica, lo que limita la accesibilidad y reduce la probabilidad de uso indebido. El trabajo no pretende resolver todos los problemas de superposición o robustez adversarial, sino que sugiere que fomentar la ortogonalidad puede actuar como un sesgo arquitectónico fuerte hacia modelos más robustos.