Resumen Técnico: Un Estudio sobre Destilación de Capas Ocultas para el Pre-entrenamiento de Modelos de Lenguaje Grandes

Enunciado del Problema

Aunque la Destilación de Conocimiento (KD) es un paradigma establecido para comprimir Modelos de Lenguaje Grandes (LLM) transfiriendo conocimiento de un "maestro" grande a un "estudiante" más pequeño, la mayoría de la investigación actual se basa exclusivamente en los logits de salida. Este enfoque descuida la información semántica contenida dentro de las representaciones intermedias ocultas del maestro. Aunque la Destilación de Capas Ocultas (HLD) ha mostrado promesa en arquitecturas de codificadores y modelos a menor escala, su aplicación al pre-entrenamiento exclusivo de decodificadores a gran escala permanece en gran medida inexplorada. Además, las comparaciones existentes a menudo no tienen en cuenta la sobrecarga computacional introducida por la HLD (por ejemplo, almacenar activaciones completas, costos de regresores) o no controlan rigurosamente el ajuste de hiperparámetros, lo que dificulta determinar si la HLD ofrece ventajas genuinas sobre una KD basada en logits bien ajustada.

Metodología

Los autores realizan una evaluación controlada por cómputo de la HLD aplicada al pre-entrenamiento de LLMs causales utilizando la familia de modelos Gemma.

• Modelos:
  • Maestro: Un modelo Gemma de 3.4B (transformador de decodificador único de 34 capas, dimensión de incrustación de 2560) preentrenado con 106B tokens únicos del conjunto de datos C4.
  • Estudiantes: Dos modelos basados en Gemma con inicialización aleatoria estándar: un modelo de 123M de parámetros (18 capas) y un modelo de 735M de parámetros (26 capas).
• Métodos de Destilación Evaluados:
  1. KD Estándar (Basada en Logits): El estudiante minimiza una pérdida compuesta de Verosimilitud Negativa Logarítmica (NLL) y divergencia Kullback-Leibler (KL) contra los logits de salida del maestro.
  2. Optimización Secuencial (HLDF): Adaptado de FitNets, esto implica un protocolo de dos fases. La Fase 1 alinea el estado intermedio oculto del estudiante (en la capa mediana) con el del maestro utilizando un regresor aprendido (perceptrón denso de una capa) mediante una pérdida de Error Cuadrático Medio (MSE) normalizada. La Fase 2 realiza KD estándar.
  3. Optimización Conjunta (HLDC): Un enfoque de entrenamiento de una sola etapa que minimiza una función objetivo agregada que combina NLL, destilación de logits y una pérdida de alineación de capa oculta (pérdida de incrustación) simultáneamente.
• Rigor Experimental:
  • Ajuste de Cómputo: Las comparaciones se basan estrictamente en presupuestos computacionales equivalentes (FLOPs), no solo en conteos de tokens. Los autores refinan la aproximación estándar de cómputo $6ND$ para incluir el costo significativo de las capas de desincrustación y el mapeo de regresores, lo cual es no despreciable para modelos submilesionarios con vocabularios grandes.
  • Control de Hiperparámetros: El estudio evita "ajustar para ganar" evaluando los métodos a través de una cuadrícula compartida de hiperparámetros (tasa de aprendizaje, temperatura, peso NLL) y comparando resultados punto por punto. También identifica la mejor configuración para cada método para asegurar una comparación justa "de techo a techo".
  • Datos: El entrenamiento ocurre en el subconjunto en inglés del conjunto de datos C4 (hasta 168B tokens para el estudiante de 735M).

Contribuciones Clave

1. Evaluación Sistemática: El artículo proporciona el primer estudio de código abierto de HLD específicamente durante la fase de pre-entrenamiento de LLMs de decodificador único, escalando a modelos con cientos de millones de parámetros.
2. Líneas Base Rigurosas: Establece un marco que compara métodos basándose en equivalencia computacional estricta (FLOPs), contabilizando explícitamente los costos no despreciables de la desincrustación y el cálculo de pérdidas en modelos pequeños, los cuales a menudo fueron pasados por alto en estudios anteriores.
3. Evaluación Comparativa Exhaustiva: La evaluación abarca múltiples tareas posteriores (Wikitext 103, HellaSwag, WinoGrande, LAMBADA, PIQA, ARC-E) y analiza el rendimiento a través de una amplia cuadrícula de hiperparámetros en lugar de depender de una única configuración "mejor".

Resultados

Los experimentos arrojan hallazgos matizados sobre la eficacia de la HLD en comparación con la KD estándar:

• Rendimiento Posterior: La HLD no supera consistentemente a la KD basada en logits estándar en tareas de evaluación posteriores. En la comparación de "mejor hiperparámetro", ninguna variante de HLD domina a la KD en todas las pruebas; las diferencias son marginales y se dividen entre los métodos.
• Mejoras en Perplejidad: A pesar de la falta de superioridad posterior, HLDF (Optimización Secuencial) produce una mejora sistemática, aunque pequeña, en la perplejidad de C4 a través de todas las configuraciones de hiperparámetros compartidas. La mejora media es de aproximadamente 0.26% para ambos tamaños de estudiante.
• Contexto de las Mejoras: Las ganancias en perplejidad de HLDF son más pronunciadas en regímenes donde la KD estándar tiene un rendimiento deficiente. En regímenes donde la KD ya es eficiente, la ventaja de la HLD desaparece.
• Posterior vs. Perplejidad: La ganancia sistemática en perplejidad no se traduce en mejoras consistentes en las puntuaciones posteriores. Los autores sugieren que esto indica que, aunque existe una señal latente en las capas ocultas, esta puede no alterar fundamentalmente las capacidades de grano grueso capturadas por las pruebas posteriores, o que la señal solo es útil para el ajuste de distribuciones de grano fino.

Significado y Afirmaciones

El artículo concluye que, aunque la HLD es un concepto atractivo para capturar conocimiento interno profundo, no proporciona una mejora consistente y fiable sobre una línea base de KD basada en logits bien ajustada cuando los presupuestos de entrenamiento se igualan estrictamente.

• Resiliencia de la KD de Logits: Los hallazgos subrayan la resiliencia de la destilación basada en logits estándar, sugiriendo que, para el pre-entrenamiento de decodificadores únicos, la complejidad adicional y la sensibilidad a hiperparámetros de la HLD pueden superar sus beneficios en el despliegue práctico.
• Discrepancias Metodológicas: Los autores atribuyen la contradicción con trabajos anteriores (que a menudo reportan ganancias de HLD) a cuatro factores:
  1. Contabilidad de Cómputo: Estudios anteriores a menudo igualaron conteos de tokens en lugar de FLOPs, ignorando el costo extra de la HLD.
  2. Arquitectura/Fase: Trabajos anteriores se centraron en codificadores o post-entrenamiento, donde las representaciones ya están moldeadas por objetivos de pre-entrenamiento. Los decodificadores causales entrenados desde cero presentan un paisaje de optimización diferente.
  3. Ajuste de Línea Base: Estudios anteriores a menudo compararon HLD contra líneas base de KD mal ajustadas.
  4. Escalabilidad: Adaptar la HLD de pequeñas CNNs a transformadores de miles de millones de parámetros requiere decisiones de diseño más profundas (pérdidas, adaptadores) que simples regresores lineales.
• Perspectiva Futura: El artículo postula que la ganancia sistemática en perplejidad observada en HLDF sugiere que se puede extraer una "señal latente", pero lograr un avance probablemente requeriría innovaciones significativas en el diseño o evidencia estadística más sólida a través de escalas mayores y diferentes familias de modelos. Los autores enfatizan la necesidad de una contabilidad exacta de FLOPs al evaluar nuevas técnicas de supervisión.