We Think, Therefore We Align LLMs to Helpful, Harmless and Honest Before They Go Wrong

Resumen Técnico: Dirección Adaptativa de Múltiples Ramas (AMBS) para la Alineación de LLM con Múltiples Objetivos

Enunciado del Problema
Los Modelos de Lenguaje Grandes (LLM) requieren alineación con la intención humana, típicamente operacionalizada a través de tres objetivos en competencia: Utilidad, Inocuidad y Honestidad (HHH). Aunque trabajos anteriores han logrado la alineación mediante Ajuste Fino Supervisado (SFT) o Aprendizaje por Refuerzo a partir de Retroalimentación Humana (RLHF), estos métodos a menudo luchan por satisfacer conjuntamente los objetivos HHH durante la inferencia, particularmente en entornos complejos.

Los enfoques existentes de alineación en tiempo de inferencia, como los vectores de dirección, típicamente tratan los objetivos de forma aislada. Cuando una única representación compartida es dirigida hacia un objetivo, puede sobrescribir o interferir con otros. Extensiones recientes a un entorno de Transformer "1-a-N" (replicando una representación compartida en múltiples vías) intentan mitigar esto aplicando transformaciones independientes. Sin embargo, estas actualizaciones independientes carecen de restricciones sobre las relaciones entre las vías, lo que lleva a respuestas inconsistentes y divergencia entre objetivos. Además, los métodos de compensación basados en entrenamiento (por ejemplo, RLHF seguro, Mezcla de Expertos) no restringen las transformaciones específicas de objetivos dentro de una representación compartida durante la inferencia, resultando en "colapso de ejes" o representaciones disjuntas donde los objetivos conflictivos no pueden satisfacerse simultáneamente.

Metodología: Dirección Adaptativa de Múltiples Ramas (AMBS)
Los autores proponen AMBS, un marco de dos etapas diseñado para un entorno de Transformer "1-a-N" que parametriza transformaciones específicas de objetivos relativas a una representación compartida. La innovación central consiste en definir las actualizaciones como desviaciones de una referencia común en lugar de modificaciones independientes.

1. Etapa I: Cálculo de Representación Compartida
   El Transformer calcula una representación oculta compartida $h^{(0)}_\ell$ en la capa $\ell$ (típicamente la capa final $L$) para una entrada dada. Esta representación sirve como base común y se reutiliza en todos los objetivos.

2. Etapa II: Transformaciones Guiadas por Referencia
   La representación compartida se replica en $N$ vías ($h^{(i)}_\ell = h^{(0)}_\ell$). Para cada objetivo $i$, se aplica una transformación $\Delta h^{(i)}_\ell$ relativa a la referencia compartida $h^{ref}_\ell = h^{(0)}_\ell$.

   • Actualización Basada en Desviación: En lugar de actualizaciones independientes, la transformación se define como una función de la desviación desde la referencia: $\Delta h^{(i)}_\ell = f_i(\tilde{h}^{(i)}_\ell - h^{ref}_\ell; \phi_i)$.
   • Parametrización de Baja Rango: La política $f_i$ se implementa como una proyección de baja rango con un sesgo aditivo: $\Delta h^{(i)}_\ell = ((\tilde{h}^{(i,0)}_\ell - h^{ref}_\ell)A_i)B_i + b_i$. Aquí, $A_i$ y $B_i$ son matrices de baja rango aprendibles, y $b_i$ es un sesgo aprendible. Crucialmente, el sesgo $b_i$ asegura una desviación inicial no nula, activando la actualización de baja rango inmediatamente dentro de un solo paso hacia adelante sin computación iterativa de punto fijo.
   • Generación de Salida: Cada vía produce una distribución de salida específica del objetivo $p^{(i)}_\theta$. Estas se combinan durante la decodificación en una única distribución conjunta $p^{joint}_\theta$ utilizando una media geométrica ponderada.

3. Objetivo de Optimización
   El marco minimiza una función de pérdida compuesta:
   $$\min \sum_{i=1}^N L^{(i)}_{obj} + \lambda L_{cons}$$

   • $L^{(i)}_{obj}$: La pérdida específica del objetivo (por ejemplo, entropía cruzada) para la vía $i$.
   • $L_{cons}$: Un término de consistencia que mide la distancia euclidiana al cuadrado entre las representaciones transformadas a través de las vías. Este término penaliza la divergencia excesiva mientras permite las variaciones necesarias específicas del objetivo, previniendo el "colapso" observado en otros métodos conjuntos.

Contribuciones Clave

• Diseño del Marco: Introducción de AMBS, que extiende el entorno de Transformer "1-a-N" parametrizando las transformaciones como desviaciones de una referencia compartida, controlando así la interferencia mientras mantiene la distinción de las vías.
• Formulación Teórica: Una formulación matemática donde las diferencias entre vías dependen únicamente de los términos de desviación ($\Delta h^{(i)}_\ell - \Delta h^{(j)}_\ell$) en lugar de transformaciones independientes del estado base, respaldada por lemas que prueban la inicialización no degenerada y la activación inmediata de los parámetros de baja rango.
• Rendimiento Empírico: Demostración de que AMBS mejora el rendimiento conjunto HHH en cuatro modelos base diferentes (LLaMA-2-7B, Mistral-7B, Gemma-7B, DeepSeek-7B) manteniendo la eficiencia computacional.

Resultados
Se realizaron experimentos en benchmarks estándar: Alpaca (Utilidad), BeaverTails (Inocuidad) y TruthfulQA (Honestidad).

• Ganancias de Rendimiento: AMBS logró las puntuaciones promedio (Avg) más altas entre los modelos base. Para LLaMA-2-7B, AMBS alcanzó un Avg de 56.5%, superando al Estado del Arte anterior (TrinityX) en 55.1% y a AlignX en 51.3%.
• Mejoras Específicas por Objetivo:
  • Utilidad: Aumentó la Tasa de Victoria (WR) al 53.0% (vs. 48.7% para RAHF).
  • Inocuidad: Redujo la Puntuación de Seguridad (SS, donde menor es mejor) al 5.3% (vs. 7.2% para Aligner).
  • Honestidad: Mejoró la Veracidad-Informatividad (TI) al 37.9% (vs. 31.7% para DExperts).
• Análisis de Interferencia: A diferencia de los métodos de objetivo único que muestran grandes caídas de rendimiento cuando cambian los objetivos, o los métodos conjuntos que exhiben "colapso" (puntuaciones idénticas entre objetivos), AMBS mantuvo un rendimiento alto pero no uniforme entre objetivos (por ejemplo, 37.0% / 25.8% / 14.3%), indicando una interferencia reducida sin colapso.
• Eficiencia: El método mantiene un alto rendimiento (por ejemplo, 189 tokens/seg) y un bajo costo de entrenamiento (9 horas de GPU) congelando el modelo base y optimizando solo los parámetros de baja rango.

Significado y Afirmaciones
El artículo afirma que AMBS aborda una brecha crítica en la alineación de múltiples objetivos en tiempo de inferencia. Al restringir las transformaciones específicas de objetivos dentro de una representación compartida, el marco equilibra exitosamente la compensación entre satisfacer objetivos distintos (Utilidad, Inocuidad, Honestidad) y mantener la consistencia entre ellos. Los autores afirman que este enfoque previene el "colapso de ejes" y la inestabilidad observados en métodos anteriores de alineación conjunta, ofreciendo una solución robusta para desplegar LLM confiables sin la necesidad de reentrenamiento costoso o arquitecturas modulares complejas. Los resultados sugieren que la alineación se captura efectivamente a nivel de representación, con la decodificación sirviendo como un paso de consolidación en lugar de la fuente principal de ganancias.

Limitaciones
Los autores reconocen limitaciones que incluyen la escala de la evaluación, la dependencia de jueces proxy (LLM) para las métricas en lugar de una evaluación humana extensa, y el enfoque actual en modelos solo decodificadores. Se sugiere trabajo futuro que incluya estudios humanos más amplios, configuraciones de seguridad más extensas y la extensión a modelos multimodales.