Tuning Language Models by Mixture-of-Depths Ensemble

Resumen Técnico: Ajuste de Modelos de Lenguaje mediante el Ensamble de Mezcla de Profundidades (MoDE)

Planteamiento del Problema
Los Modelos de Lenguaje Grandes (LLMs) tradicionales basados en Transformers dependen exclusivamente de las representaciones de la última capa tanto para la predicción como para el cálculo de la pérdida durante el ajuste fino (fine-tuning). Aunque las herramientas de interpretabilidad como el "logit lens" han demostrado que los estados ocultos de las capas tardías ya contienen predicciones mayoritariamente formadas y relevantes para la tarea, estas representaciones intermedias son típicamente tratadas como subproductos implícitos en lugar de señales de entrenamiento explícitas. En consecuencia, el potencial predictivo de las capas tardías se desaprovecha y el ajuste fino estándar optimiza solo la última capa, dejando las representaciones intermedias optimizadas solo de forma indirecta. Además, los métodos existentes para aprovechar las capas intermedias (por ejemplo, salida temprana para eficiencia o direccionamiento de representaciones) a menudo no tratan los logits específicos de la capa como un componente directamente entrenable de la predicción final.

Metodología: Ensamble de Mezcla de Profundidades (MoDE)
Los autores proponen el Ensamble de Mezcla de Profundidades (MoDE), un marco de ajuste que trata las salidas de las últimas $k$ capas de un LLM como un ensamble de "expertos" dentro de un mismo modelo. A diferencia de la Mezcla de Expertos (MoE) que enruta diferentes tokens hacia diferentes modelos, MoDE enruta la predicción final hacia una combinación ponderada de logits de diferentes profundidades dentro del mismo modelo.

El marco consta de tres componentes principales:

1. Extracción de Logits de Capas Tardías: Para un modelo con $n$ capas, se seleccionan las últimas $k$ capas ($K = \{L_{n-k+1}, \dots, L_n\}$). Se aplica la cabeza de vocabulario $\phi(\cdot)$ a los estados ocultos de cada una de estas $k$ capas para generar logits $\ell_i(x)$.
2. Red de Enrutamiento Aprendida: Un enrutador ligero $G(x)$, implementado como una única capa lineal seguida de un softmax, calcula los pesos de enrutamiento basados en el estado oculto justo antes de las últimas $k$ capas ($x = H_{n-k}$). Los logits finales se computan como una suma ponderada:
   $$\ell(x_t | x_{<t}) = \sum_{i=0}^{k-1} G(x)_i \cdot \ell_i(x)$$
   El enrutador puede hacerse disperso (Top-K) para permitir la salida temprana (early exit) durante la inferencia.
3. Adaptación mediante Normalización y Destilación: Aplicar directamente la cabeza preentrenada a las capas intermedias suele resultar en una pérdida inicial alta. Para abordar esto sin añadir parámetros pesados (a diferencia de las matrices afines por capa), los autores introducen:
   • Normalización por Capa ($N_k$): Se añaden módulos de normalización entrenables (que coinciden con la arquitectura preentrenada, por ejemplo, LayerNorm) a cada capa tardía para adaptar las representaciones al espacio de entrada de la cabeza.
   • Pérdida de Destilación ($L_{distill}$): La última capa actúa como un "maestro". Se aplica una pérdida de divergencia KL para atraer las distribuciones de salida de las capas tardías anteriores hacia la distribución de la última capa, estabilizando el entrenamiento y alineando el ensamble con el objetivo de la tarea.

La pérdida total es $L_{MoDE} = L_{task} + \lambda L_{distill}$. MoDE está diseñado para ser un mecanismo plug-in que puede superponerse a cualquier método de ajuste fino eficiente en parámetros (PEFT) existente (por ejemplo, LoRA, DoRA) o utilizarse de forma independiente.

Contribuciones Clave

• Señales de Capas Tardías Entrenables: El artículo demuestra que los logits de las capas tardías, cuando se optimizan individualmente con destilación, convergen a pérdidas comparables a las de la última capa y resuelven un conjunto complementario de problemas.
• Marco MoDE: Un mecanismo de ajuste simple y de baja sobrecarga que ensambla los logits de las capas tardías mediante un enrutador aprendido. Añade parámetros entrenables insignificantes (aproximadamente +0.04% del tamaño del modelo) mientras permite al modelo sustituir módulos entrenables sustancialmente mayores (como LoRA de capa completa) con un rendimiento comparable.
• Eficiencia y Flexibilidad: El marco admite el enrutamiento disperso (Top-K), lo que permite la salida temprana y aceleraciones de generación de hasta 1.6$\times$ con una pérdida mínima de precisión. Se compone con los métodos PEFT existentes, mejorando su rendimiento.

Resultados Experimentales
Los experimentos se realizaron en LLaMA2-7B y LLaMA3-8B en tareas de razonamiento aritmético, razonamiento de sentido común y seguimiento de instrucciones.

• Razonamiento Aritmético: MoDE mejoró consistentemente el rendimiento de los métodos base (LoRA, DoRA). La variante ligera (LoRA aplicado a todas las capas excepto las últimas $k$, más MoDE) alcanzó un rendimiento dentro de 0.3–0.5 puntos de LoRA de capa completa utilizando significativamente menos parámetros.
• Razonamiento de Sentido Común: MoDE proporcionó ganancias aditivas sobre los métodos basados en pesos (LoRA, DoRA). Mientras que los métodos basados en representaciones (ReFT) funcionaron fuertemente por sí solos, MoDE ofreció mejoras fiables al combinarse con enfoques basados en pesos.
• Seguimiento de Instrucciones: MoDE produjo ganancias pequeñas pero consistentes en MT-Bench, aunque las mejoras fueron más modestas que en las tareas de razonamiento, lo que sugiere que el seguimiento de instrucciones depende fuertemente de un formateo de capas tardías que el ensamble captura solo parcialmente.
• Ablación y Análisis:
  • El enrutador naturalmente concentra el peso en una o dos últimas capas, pero los componentes de destilación y normalización permiten que las capas tardías anteriores contribuyan efectivamente.
  • Un tamaño de ensamble de $k \approx 3\text{--}4$ es óptimo; un $k$ mayor ayuda en conjuntos de datos fáciles pero perjudica las tareas de razonamiento complejo.
  • El enrutamiento disperso Top-K mantiene la mayor parte de la precisión permitiendo aceleraciones significativas.

Significancia y Reivindicaciones
Los autores presentan a MoDE como un estudio de mecanismo que demuestra que los logits de las capas tardías pueden hacerse directamente útiles para el ajuste. El artículo afirma que:

1. Las representaciones de las capas tardías poseen un potencial predictivo significativo que puede desbloquearse mediante una simple adaptación (normalización y destilación).
2. Estas señales intermedias pueden sustituir a módulos entrenables sustancialmente mayores (como LoRA de capa completa) con un rendimiento comparable, ofreciendo una alternativa eficiente en parámetros.
3. MoDE proporciona una dirección complementaria para convertir las representaciones de las capas intermedias en señales de entrenamiento accionables, distintas de los enfoques de sondeo en tiempo de inferencia o de edición de pesos.

El artículo mantiene la modestia respecto a su alcance, señalando limitaciones tales como la dependencia de una $k$ fija (en lugar de una selección dinámica), las ganancias menores en el seguimiento de instrucciones en comparación con el razonamiento, y una evaluación limitada a LLMs de solo decodificador. Se sugiere como trabajo futuro la selección dinámica de $k$, familias de modelos más amplias y modelos más grandes.