Beyond Linear Steering: Unified Multi-Attribute Control for Language Models

1. El Problema

El control de múltiples atributos de comportamiento en Modelos de Lenguaje Grandes (LLMs) durante la inferencia es un desafío significativo. Los métodos existentes de "dirección" (steering) se basan principalmente en intervenciones lineales, que asumen que los comportamientos se componen de manera aditiva en el espacio de activaciones.

Esto presenta varias limitaciones críticas:

• Interferencia de atributos: Al combinar vectores de dirección lineales para múltiples atributos (ej. tono y estilo de debate), a menudo se produce interferencia, degradación de la fluidez o resultados diluidos.
• Suposición de linealidad: La realidad de las representaciones internas de los LLMs es no lineal; los comportamientos no siempre se suman de forma simple.
• Complejidad de ajuste: Los métodos actuales requieren el ajuste y almacenamiento de vectores separados para cada atributo y a menudo necesitan funciones de "gating" (puerta) aprendidas por token, lo que complica la composición dinámica de comportamientos.

2. Metodología: K-Steering

Los autores proponen K-Steering, un enfoque unificado y flexible que abandona la linealidad para utilizar un clasificador no lineal entrenado sobre las activaciones ocultas del modelo.

Componentes Clave:

1. Entrenamiento de un Clasificador No Lineal (MLP):

   • Se entrena un Perceptrón Multicapa (MLP) como un clasificador de múltiples etiquetas sobre las activaciones de la capa final del modelo.
   • El clasificador mapea las activaciones a $K$ categorías (ej. tonos conversacionales o estilos de debate).
   • A diferencia de los métodos anteriores que usan vectores fijos, este clasificador aprende una representación conjunta de múltiples comportamientos.

2. Intervención Basada en Gradientes:

   • En tiempo de inferencia, no se suman vectores precalculados. En su lugar, se calculan los gradientes de la función de pérdida del clasificador respecto a las activaciones del modelo.
   • La intervención se formula como:
     $$a'_i = a_i - \alpha \nabla_{a_i} \mathcal{L}(g_\phi(a_i))$$
     Donde $a_i$ es la activación, $\alpha$ es un factor de escala, $\mathcal{L}$ es la pérdida (maximizar clases objetivo, minimizar clases a evitar) y $g_\phi$ es el clasificador.
   • Esto permite ajustar el comportamiento del modelo hacia o desde un conjunto específico de atributos sin asumir una estructura lineal.

3. Algoritmos de Intervención:

   • Algoritmo 1 (Dirección Iterativa): Ajusta las activaciones iterativamente basándose en los gradientes para recompensar clases objetivo y penalizar las no deseadas. Utiliza una tasa de aprendizaje decreciente para estabilidad.
   • Algoritmo 2 (Eliminación de Proyección): Un método de un solo paso para eliminar atributos no deseados. Utiliza una reflexión de Householder sobre el gradiente de la pérdida de las clases a evitar. Esto invierte el componente de la activación en la dirección del gradiente, empujando la representación lejos de los límites no deseados en el espacio no lineal. Es computacionalmente más eficiente que la iteración.

3. Contribuciones Principales

1. Nuevos Conjuntos de Datos (Benchmarks):

   • TONEBANK: Un dataset de preguntas diseñadas para ser respondidas en 6 tonos conversacionales distintos (ej. empático, experto, casual).
   • DEBATEMIX: Un dataset de preguntas de debate diseñadas para ser respondidas en 10 estilos retóricos distintos (ej. Reductio ad Absurdum, Apelación al Precedente, Carga de la Prueba).
   • Estos datasets permiten probar la composición real de múltiples atributos en una sola respuesta, algo que los datasets binarios anteriores no lograban.

2. Avance Metodológico:

   • Demostración de que K-Steering puede controlar múltiples direcciones simultáneamente superando o igualando a los baselines más fuertes (CAA y DCT).
   • Validación de que el uso de pasos de gradiente adicionales con tamaños de paso más pequeños mejora el rendimiento, aunque tamaños de paso grandes pueden degradar la coherencia.
   • Introducción de la eliminación de proyección como un método eficiente para suprimir comportamientos específicos.

3. Marco de Evaluación:

   • Desarrollo de un protocolo de evaluación semi-automatizado que utiliza búsqueda binaria para calibrar la magnitud de la dirección ($\alpha$), maximizando la efectividad del control mientras se minimiza el desplazamiento de la distribución (garantizando que el modelo no genere texto incoherente).
   • Uso de clasificadores basados en activaciones y jueces basados en LLMs para validar los resultados.

4. Resultados Experimentales

Los experimentos se realizaron en tres familias de modelos: Llama-3.2-3B, Mistral-7B y OLMo-2-7B.

• Rendimiento Superior: K-Steering superó consistentemente a los métodos de línea base, Contrastive Activation Addition (CAA) y Deep Causal Transcoding (DCT), en la mayoría de las combinaciones de atributos ($K=1, 2, 3$).
• Composición de Atributos: Mientras que CAA y DCT sufrían de degradación al combinar 3 atributos (especialmente en Llama-3.2-3B), K-Steering mantuvo un rendimiento alto, demostrando su capacidad para manejar interacciones no lineales entre atributos.
• Evaluación de Jueces: Un juez basado en LLM (GPT-4o-Mini) confirmó en la mayoría de los casos (10 de 15) que las generaciones de K-Steering exhibían los atributos objetivo con mayor fuerza que las de CAA.
• Preservación de Capacidades: En pruebas de benchmark general (MMLU), K-Steering preservó la precisión del modelo base de manera competitiva, mostrando que la intervención no degrada significativamente las capacidades generales del modelo.
• Análisis de Pasos Múltiples: Se observó que el uso de múltiples pasos de gradiente mejora la efectividad, especialmente con valores de $\alpha$ pequeños, aunque con un costo computacional lineal adicional.

5. Significado e Implicaciones

• Superación de la Linealidad: El trabajo demuestra que la suposición de que los comportamientos de los LLMs se componen linealmente es incorrecta para tareas complejas. El enfoque basado en gradientes de un clasificador no lineal ofrece una vía más robusta para el control.
• Control Dinámico y Unificado: Elimina la necesidad de almacenar y ajustar vectores separados para cada atributo, permitiendo una composición dinámica de comportamientos sin reentrenar el modelo.
• Escalabilidad: El método escala naturalmente a conjuntos de atributos más grandes, lo cual es crucial para aplicaciones en entornos de alto riesgo donde se requiere un control fino y multifacético del comportamiento del modelo.
• Limitaciones y Futuro: El método tiene un mayor costo computacional en comparación con la adición de vectores simples (CAA) debido a los cálculos de gradiente iterativos. Además, la generalización a datasets del mundo real (fuera de los datasets sintéticos creados) sigue siendo una pregunta abierta.

En conclusión, K-Steering representa un avance significativo hacia el control composicional fiable de LLMs, ofreciendo una alternativa más flexible y potente a los métodos de dirección lineal tradicionales.