DistillLens: Symmetric Knowledge Distillation Through Logit Lens

Resumen Técnico: DISTILLLENS

1. El Problema

La Destilación de Conocimiento (KD) estándar comprime Modelos de Lenguaje Grandes (LLMs) optimizando únicamente la distribución de salida final (logits). Este enfoque trata el proceso de pensamiento intermedio del modelo "maestro" como una caja negra, ignorando las representaciones de las capas intermedias.

Esto genera dos problemas críticos:

• Desconexión del proceso de pensamiento: Aunque el modelo estudiante converja en la respuesta final correcta, su trayectoria interna (el "proceso de pensamiento" evolutivo) diverge significativamente de la del maestro, como se ilustra en la Figura 1 del artículo. Esto es análogo a memorizar la respuesta final de una prueba matemática sin entender los pasos de la derivación, lo que limita la generalización.
• Limitaciones de los métodos existentes:
  • Los métodos basados en Error Cuadrático Medio (MSE) asumen un espacio de incrustación isotrópico, lo cual es falso en espacios de lenguaje de alta dimensión (anisotrópicos). El MSE es insensible a la divergencia de probabilidad.
  • Los métodos que utilizan Divergencia de Kullback-Leibler (KL) asimétrica (como $KL(p||q)$) sufren de un comportamiento de "búsqueda de la media" (mean-seeking) o "búsqueda de modo" (mode-seeking) desequilibrado. Penalizan fuertemente la subestimación de tokens de alta probabilidad pero ignoran la sobreestimación de las colas de baja probabilidad, o viceversa, lo que impide una alineación precisa de la incertidumbre semántica en las capas ocultas.

2. Metodología: DISTILLLENS

El authors proponen DISTILLLENS, un marco que alinea simétricamente los procesos de pensamiento evolutivos del estudiante y el maestro. La metodología se basa en tres pilares:

• Proyección mediante Logit Lens:
  En lugar de alinear directamente los estados ocultos ($h$), el método proyecta estos estados intermedios al espacio de vocabulario utilizando la matriz de "desincrustación" (unembedding matrix, $W_U$) del modelo. Esto convierte las representaciones internas en distribuciones de probabilidad sobre el vocabulario en cada capa $l$:
  $$p^{(l)}(y|x) = \text{softmax}(W_U h^{(l)})$$
  Esto permite interpretar las capas intermedias como estados de creencia evolutivos.

• Objetivo de Divergencia Simétrica:
  Para alinear estas distribuciones proyectadas, DISTILLLENS utiliza una función de pérdida simétrica, específicamente la Divergencia de Jensen-Shannon (JSD), en lugar de la KL asimétrica.

  • Justificación Teórica: La JSD impone una penalización de doble cara.
    • Penaliza linealmente la sobreconfianza (cuando el estudiante asigna demasiada probabilidad a un token).
    • Penaliza de forma acotada (saturada) la subconfianza (cuando el estudiante ignora la probabilidad del maestro).
  • Esto evita que el estudiante se desvíe hacia estados de entropía excesiva o suprima matices, forzando una alineación perfecta ($c_\theta \to 1$) donde la probabilidad del estudiante coincide con la del maestro.

• Mapeo de Capas y Entrenamiento Modular:

  • Se establece un mapeo uniforme entre las capas del estudiante y las del maestro (basado en la proporción de profundidad) para capturar la evolución de los estados ocultos a través de toda la red.
  • La pérdida total combina la pérdida de tarea estándar (ej. KD final) con la pérdida de alineación intermedia ponderada por un hiperparámetro $\lambda$:
    $$L_{total} = L_{task} + \lambda \cdot L_{inter}$$
  • Donde $L_{inter}$ es el promedio de la JSD sobre las capas mapeadas.

3. Contribuciones Clave

1. Nuevo Marco de Destilación: Introducción de DISTILLLENS, el primer marco que supervisa simétricamente las distribuciones de las capas intermedias proyectadas al espacio de vocabulario mediante Logit Lens.
2. Fundamentación Teórica: Demostración teórica de que el paisaje de pérdida simétrico (JSD) impone una penalización dual que previene tanto la sobreconfianza como la subconfianza, preservando los conductos de información de alta entropía esenciales para la deducción final.
3. Superioridad Empírica: Validación exhaustiva que demuestra que DISTILLLENS supera consistentemente a la KD estándar y a los métodos de transferencia de características (MSE, FDD) en diversas arquitecturas (GPT-2, Llama) y tareas de seguimiento de instrucciones.

4. Resultados Experimentales

Los experimentos se realizaron en familias de modelos GPT-2 (120M, 340M, 1.5B) y Llama (1.1B, 7B) utilizando el conjunto de datos databricks-dolly-15k.

• Rendimiento General: DISTILLLENS superó a todos los baselines (incluyendo SFT, KD estándar, SeqKD, y métodos avanzados como AKL).
  • En GPT-2-120M, el método logró un puntaje promedio Rouge-L de 21.12, superando significativamente a la KD estándar (17.74) y al SFT (16.56).
  • En GPT-2-340M, alcanzó 23.72, superando ligeramente al propio modelo maestro (23.52).
  • En TinyLlama-1.1B, alcanzó 25.48, superando a la KD estándar en más de 1.6 puntos.
• Reducción del Sesgo de Exposición: A diferencia de la destilación on-policy (que es costosa computacionalmente), DISTILLLENS reduce el sesgo de exposición al alinear los procesos internos, logrando un rendimiento competitivo sin la sobrecarga de generación de muestras durante el entrenamiento.
• Análisis de Capas: Se identificó que utilizar 5 capas intermedias con un patrón de selección intercalado (interleaved) ofrece el mejor rendimiento, superando a configuraciones con menos capas o capas consecutivas.
• Comparación con Transferencia de Características: DISTILLLENS superó consistentemente a métodos basados en MSE y FDD, confirmando que proyectar al espacio de vocabulario y usar divergencia simétrica es superior a la regresión vectorial directa.

5. Significado e Impacto

• Cambio de Paradigma: El trabajo demuestra que la calidad de un modelo destilado no depende solo de la coincidencia de la salida final, sino de la fidelidad en la trayectoria de pensamiento interna.
• Eficiencia vs. Rendimiento: Ofrece un camino computacionalmente eficiente para transferir capacidades de razonamiento complejas a modelos pequeños, cerrando la brecha entre el entrenamiento rápido off-policy y la robustez del entrenamiento on-policy.
• Interpretabilidad Activa: Reutiliza la herramienta de "Logit Lens" (originalmente para interpretación post-hoc) como un mecanismo de supervisión activa durante el entrenamiento, asegurando que la trayectoria semántica interna del estudiante siga la del maestro.
• Limitaciones y Futuro: El principal costo es el sobrecarga computacional durante el entrenamiento debido a la proyección de múltiples capas. El trabajo futuro se centrará en optimizar esta eficiencia y extender el marco a la destilación entre arquitecturas heterogéneas (cross-architecture).

En conclusión, DISTILLLENS establece un nuevo estándar en la destilación de LLMs al tratar el proceso de inferencia como un todo estructurado y simétrico, en lugar de un simple mapeo de entrada-salida.