Resumen Técnico: La Geometría de los Vectores de Tarea Subyace a los Modos Duales de Inferencia de Tarea en Transformadores

Enunciado del Problema

Los Modelos de Lenguaje Grandes (LLM) demuestran dos modos distintos de inferencia de tarea: recuperación de tarea bayesiana (reconocer y recuperar tareas vistas durante el entrenamiento) y aprendizaje extrapolativo de tarea (adaptarse a tareas novedosas, fuera de distribución (OOD)). Aunque estudios recientes de interpretabilidad han identificado "vectores de tarea"—direcciones en el espacio de estados ocultos que dirigen el comportamiento del modelo—, faltan fundamentos teóricos que conecten estas representaciones internas con los comportamientos de generalización externos. Específicamente, el trabajo existente no logra explicar rigurosamente:

1. Cómo la geometría del vector de tarea está moldeada por la distribución de los datos de entrenamiento.
2. Qué estructuras geométricas específicas permiten la generalización OOD.
3. Cómo las representaciones internas median el vínculo entre la estructura estadística y el comportamiento del modelo sin depender de métodos de extracción heurísticos y específicos del modelo.

Metodología

Los autores abordan estas brechas entrenando transformadores pequeños, solo decodificadores, desde cero sobre distribuciones sintéticas de secuencias de tareas latentes cuidadosamente controladas. Este enfoque de "caja gris" permite una caracterización matemática rigurosa de los estados ocultos, libre de los factores de confusión del preentrenamiento en lenguaje natural.

Configuración Experimental

Se utilizan tres experimentos sintéticos (E1–E3), donde las secuencias se generan a partir de variables latentes no observadas $z$ (tareas):

• E1 (Dados Sesgados): Las latentes son vectores de probabilidad muestreados de una prior Dirichlet; los tokens se muestrean i.i.d.
• E2 (Regresión Lineal en Contexto): Las latentes son vectores de peso muestreados de una prior Gaussiana; las secuencias consisten en pares de entrada-salida.
• E3 (Mezcla de Cadenas de Markov): Las latentes son matrices de transición muestreadas de una prior Dirichlet; las secuencias son cadenas de Markov.

Un cuarto experimento (E4, Lenguaje Dyck Plantado) sirve como contraejemplo no markoviano para probar los límites del marco propuesto.

Marco Teórico

El artículo define los vectores de tarea $\theta_k$ como la diferencia límite entre el estado oculto para una tarea específica y la media global. Establece cuatro propiedades geométricas (P0–P3) que caracterizan la recuperación exitosa de tareas:

• P0 (Estabilidad): Los estados ocultos se vuelven funciones deterministas de la tarea latente y el token actual.
• P1 (Desacoplamiento): Los estados ocultos se descomponen aditivamente en vectores de tarea y codificaciones de tokens (Hipótesis de Representación Lineal).
• P2 (Interpolación): En contexto finito, los estados ocultos son combinaciones convexas de vectores de tarea.
• P3 (Alineación): Los coeficientes de estas combinaciones convexas aproximan la posterior bayesiana sobre las tareas.

Los autores contrastan esto con el Aprendizaje Extrapolativo de Tarea (M2), hipotetizando que la inferencia OOD requiere representaciones que se encuentran fuera del subespacio de vectores de tarea.

Contribuciones y Resultados Clave

1. Unificación de la Geometría de Representación y la Generalización

El artículo proporciona un marco matemático riguroso que muestra que la recuperación de tarea bayesiana se implementa internamente mediante combinaciones convexas de vectores de tarea aprendidos.

• Evidencia: En los experimentos E1–E3, los estados ocultos del modelo satisfacen las propiedades P0–P3. Los coeficientes de la mezcla de vectores de tarea ($\beta_{t,k}$) rastrean estrechamente la posterior bayesiana real ($\alpha_{t,k}$) a medida que se acumula el contexto.
• Validación Causal: La intervención en el subespacio de vectores de tarea (dirigiendo el estado oculto a un punto específico en el simplex de tareas) dirige con éxito la salida del modelo para que coincida con la predicción de la mezcla correspondiente, confirmando que los vectores de tarea controlan causalmente la inferencia para tareas dentro de la distribución (ID).

2. Descubrimiento de Modos Duales de Inferencia mediante Subespacios Casi Ortogonales

El artículo demuestra que el aprendizaje extrapolativo de tarea (M2) para tareas OOD opera en un subespacio casi ortogonal al subespacio de vectores de tarea ID.

• Mecanismo: A medida que aumenta la diversidad de tareas (entrenando sobre una mezcla de unas pocas tareas "mayores" y muchas tareas "menores"), el modelo transita de depender exclusivamente de la recuperación de vectores de tarea a utilizar una vía computacional separada.
• Evidencia Geométrica: Cuando la diversidad de tareas es baja, los estados ocultos OOD proyectan fuertemente sobre el subespacio de tareas ID. A medida que aumenta la diversidad, las representaciones OOD se desvían hacia la ortogonalidad (proyección $R^2 \to 0$).
• Validación Causal: Las intervenciones confirman una "doble disociación":
  • Suprimir el subespacio de vectores de tarea devastará el rendimiento ID pero dejará el rendimiento OOD mayormente intacto.
  • Suprimir el subespacio ortogonal (identificado mediante optimización en tareas menores) destruye el rendimiento OOD mientras deja el rendimiento ID cerca de la línea base.
• Base Teórica: El Teorema 2 proporciona una cota inferior para la distancia de las representaciones OOD desde el subespacio de tareas, mostrando que si la dimensión intrínseca del mapa de predicción OOD excede la dimensión del subespacio de vectores de tarea, el modelo debe utilizar representaciones fuera de ese subespacio para generalizar con precisión.

3. Límites del Marco (Caso No Markoviano)

En el experimento del lenguaje Dyck Plantado (E4), donde los datos carecen de la propiedad de Markov de primer orden requerida para una compresión compacta, el modelo no logra satisfacer la Propiedad P0. En lugar de comprimir el contexto en un resumen de baja dimensión, el modelo retiene el historial completo del prefijo en sus estados ocultos. Esto destaca que el marco geométrico propuesto depende de que el proceso generador de datos permita estadísticas suficientes.

Significado y Afirmaciones

El artículo afirma proporcionar la primera caracterización rigurosa y fundamentada matemáticamente de cómo los Transformadores implementan modos duales de inferencia:

• Inferencia ID: Gobernada por la recuperación de tarea bayesiana dentro de un subespacio de vectores de tarea, donde el modelo interpola entre tareas memorizadas.
• Inferencia OOD: Gobernada por el aprendizaje extrapolativo de tarea en un subespacio casi ortogonal, que codifica estadísticas de contexto (por ejemplo, unigramas empíricos, soluciones de ridge, bigramas) en lugar de identidades de tareas específicas.

Los autores argumentan que la geometría de los vectores de tarea, las distribuciones de entrenamiento y los comportamientos de generalización están estrechamente relacionados. Específicamente, aumentar la diversidad de tareas obliga al modelo a desarrollar subespacios ortogonales para manejar la generalización, explicando la coexistencia de modos de recuperación y aprendizaje. El trabajo sugiere que la "hipótesis de representación lineal" se cumple para tareas ID, pero requiere una extensión (a través de subespacios ortogonales) para explicar la generalización OOD.

El artículo se mantiene modesto, señalando que su análisis depende de datos sintéticos con propiedades markovianas específicas y que una cuenta mecánica completa de la generalización OOD en lenguaje natural sigue siendo un desafío abierto. Se posiciona como un paso fundamental hacia la superación de la brecha entre la geometría de la representación interna y las capacidades de generalización externas.