Resumen Técnico: Sondeo Causal en Tiempo de Inferencia en LLMs (HDMI)

1. Declaración del Problema

El artículo aborda las limitaciones de los métodos actuales de sondeo causal utilizados para comprender cómo los Modelos de Lenguaje Grandes (LLMs) utilizan representaciones internas para la predicción.

• Limitación del Sondeo Correlacional: Los métodos tradicionales entrenan clasificadores auxiliares ("sondeos") sobre estados ocultos para detectar propiedades lingüísticas (por ejemplo, número del sujeto). Aunque estos muestran qué está codificado, no prueban cómo el modelo utiliza esa información para generar el siguiente token.
• Limitación del Sondeo Causal: Los enfoques causales existentes intentan intervenir en estados ocultos para alterar propiedades (por ejemplo, cambiar singular a plural) y observar cambios en la salida. Sin embargo, estos métodos dependen típicamente de entrenar un sondeo separado para determinar la dirección de la intervención (por ejemplo, utilizando el gradiente del sondeo para invertir una etiqueta).
  • Desventajas: Esto introduce costos adicionales de entrenamiento, requiere supervisión específica de la propiedad y conlleva riesgos de desalineación. El límite de clasificación del sondeo puede no coincidir con la "geometría predictiva" interna del modelo, lo que lleva a intervenciones que no dirigen eficazmente la generación del modelo o alteran inadvertidamente propiedades no relacionadas.

Los autores proponen el Sondeo Causal en Tiempo de Inferencia, un marco donde las intervenciones se aplican completamente en tiempo de inferencia sin reentrenar el modelo ni entrenar sondeos auxiliares.

2. Metodología: Intervención de Margen Impulsada por Estado Oculto (HDMI)

La contribución central es HDMI, una técnica basada en gradientes, libre de sondeos, que aprovecha directamente la cabeza de salida nativa del modelo generador para dirigir los estados ocultos.

2.1 Mecanismo Central

HDMI opera definiendo un objetivo de margen entre un token "fuente" (la continuación preferida original del modelo) y un token "objetivo" (la continuación contrafáctica deseada).

• Objetivo: Maximizar el margen de logits $L(x) = \phi(x)_\tau - \phi(x)_\sigma$, donde $\phi(x)$ es el vector de logits, $\tau$ es el índice del token objetivo y $\sigma$ es el índice del token fuente.
• Intervención: El método realiza un ascenso de gradiente sobre el estado oculto $h_\ell(x)$ en una capa específica $\ell$ para aumentar la probabilidad del token objetivo mientras disminuye la del token fuente.
• Eficiencia Computacional: Para la capa final, el gradiente es un producto matriz-vector de forma cerrada: $\nabla_{h_L} L(x) = W_U^\top (e_\tau - e_\sigma)$. Esto evita la necesidad de entrenar sondeos iterativos.
• Extensión a Múltiples Tokens: El método admite conjuntos de tokens (por ejemplo, múltiples verbos plurales válidos) sumando logits sobre conjuntos objetivo ($T^+$) y conjuntos fuente ($T^-$), definidos como $L_{set}(x) = \sum \phi(x)_{\tau_i} - \sum \phi(x)_{\sigma_j}$.

2.2 HDMI con Perspectiva (LA-HDMI) para Edición de Texto

Para abordar la edición de texto donde un usuario proporciona una secuencia modificada, los autores introducen LA-HDMI.

• Desafío: La generación autoregresiva estándar no permite que los estados ocultos en el paso $t$ sean conscientes de futuras ediciones en $t' > t$.
• Solución: LA-HDMI introduce un mecanismo de "perspectiva" que retropropaga gradientes a través de las incrustaciones de tokens esperadas.
  • Desacopla el paso hacia adelante (usando una temperatura baja $\beta_f$ para mantener la fluidez) de la trayectoria del gradiente (usando una temperatura alta $\beta_g$ para calcular las incrustaciones esperadas).
  • La función objetivo $J_t(h_t)$ acumula objetivos de margen para futuras posiciones de edición ($S_{max}$ pasos hacia adelante), ajustando el estado oculto actual $h_t$ para dirigir las generaciones futuras hacia las ediciones del usuario mientras preserva la fluidez en las regiones no editadas.

3. Marco de Evaluación

El artículo evalúa las intervenciones utilizando dos métricas principales, siguiendo el marco de Canby et al. [2024]:

1. Completitud: Mide si la intervención desplaza con éxito la propiedad objetivo (por ejemplo, cambiar el número del sujeto) en el estado oculto, verificado por un sondeo de validación separado.
2. Selectividad: Mide si las propiedades no relacionadas (variables de molestia, por ejemplo, número de sustantivos no sujetos) permanecen estables después de la intervención.
3. Fiabilidad: La media armónica de Completitud y Selectividad, que sirve como medida general de la calidad de la intervención.

4. Resultados Clave

Los autores evaluaron HDMI en Meta-Llama-3-8B-Instruct y Pythia-70M utilizando el corpus de acuerdo LGD y la prueba CausalGym.

• Rendimiento: HDMI logra consistentemente una mayor Fiabilidad que las líneas base impulsadas por sondeos (AlterRep, FGSM, PGD) en la mayoría de las tareas.
  • En suites de concordancia (por ejemplo, concordancia sujeto-verbo), HDMI a menudo logra Completitud perfecta y alta Selectividad.
  • A diferencia de las líneas base, el rendimiento de HDMI es independiente de la calidad de un sondeo intervencionista, ya que no utiliza uno. Líneas base como AlterRep y PGD mostraron caídas significativas de rendimiento en modelos más pequeños (Pythia-70M) cuando la precisión del sondeo era menor.
• Robustez: HDMI mantiene una fiabilidad consistente a través de diferentes tamaños de modelo y tareas, mientras que los métodos basados en sondeos mostraron resultados mixtos sensibles a la propiedad específica y al modelo.
• Edición de Texto: LA-HDMI demostró con éxito la capacidad de dirigir la generación hacia ediciones especificadas por el usuario (por ejemplo, cambiar "chica" por "búho" o "era" por "fueron") manteniendo la fluidez gramatical y ajustando el contexto necesario (por ejemplo, cambiar "una" por "un").

5. Significado y Afirmaciones

El artículo afirma que HDMI representa un cambio del sondeo causal dependiente de sondeos al sondeo causal nativo del modelo.

• Alineación: Al utilizar la propia cabeza de salida del modelo para definir la dirección de la intervención, HDMI asegura que la perturbación se alinee con la geometría predictiva real del modelo, evitando los riesgos de desalineación de los sondeos externos.
• Eficiencia: El método es computacionalmente ligero, requiriendo solo un ascenso de gradiente sobre estados ocultos sin la sobrecarga de entrenar clasificadores auxiliares.
• Generalidad: El enfoque se extiende naturalmente a propiedades de múltiples tokens y ofrece un mecanismo para la edición fluida de texto mediante gradientes con perspectiva.

Los autores reconocen limitaciones, incluida la sensibilidad a los hiperparámetros (tamaño de paso, temperatura, horizonte) y la posible atenuación de gradientes en tareas de edición de texto de largo alcance. Sugieren trabajos futuros en la selección adaptativa de capas y la programación del tamaño de paso para mejorar aún más la selectividad sin sacrificar la completitud.