The Distributional View of Knowledge Distillation

Resumen Técnico: La Visión Distribucional de la Destilación de Conocimiento

Planteamiento del Problema
La Destilación de Conocimiento (KD) estándar para modelos de lenguaje autorregresivos típicamente empareja la distribución de salida de un estudiante con la de un profesor en cada posición de token utilizando una divergencia puntual (por ejemplo, la divergencia de Kullback-Leibler). Los autores identifican dos limitaciones estructurales en este paradigma:

1. Representación de Vista Única: El profesor se resume en una única distribución suavizada (una temperatura), ignorando la trayectoria completa de "recocido" (annealing path) de las distribuciones definidas por sus logits, la cual contiene información que va desde distribuciones uniformes hasta modos codiciosos (greedy modes).
2. Ceguera Puntual: Los gradientes de KL estándar tratan todos los tokens incorrectos por igual. No distinguen entre un "error cercano" (un token semánticamente similar) y un valor atípico semántico, aplicando la misma penalización para ambos.

El artículo propone una visión distribucional donde el profesor no se representa mediante una única salida, sino mediante una familia de vistas de múltiples temperaturas. El estudiante es entonces entrenado contra un agregado consciente de la geometría de estas vistas, utilizando el Transporte Óptimo (OT) para dar cuenta de las distancias semánticas entre tokens.

Metodología
Los autores formalizan un marco donde los logits del profesor $z_T$ definen una familia exponencial de un parámetro de distribuciones $\{p_\beta\}$ mediante un softmax templado ($\beta = 1/\tau$).

• Vistas: El profesor se representa mediante $K$ vistas templadas $\mu_k = \sigma_{\tau_k}(z_T)$.
• Agregación: En lugar de una simple mezcla aritmética, los autores exploran la agregación de estas vistas utilizando Transporte Óptimo Entrópico. Definen una matriz de costo base $C$ basada en las distancias de coseno de los embeddings de los tokens normalizados.
• Objetivos:
  • Barycenters de Objetivo: Precomputar el baricentro de Wasserstein entrópico de las vistas y entrenar al estudiante mediante la divergencia KL hacia este objetivo.
  • Divergencia de Sinkhorn (Flagship): Minimizar directamente la divergencia de Sinkhorn debiasada ($S_\epsilon$) entre el estudiante y el conjunto de vistas del profesor. Esto se implementa en dos modos:
    • Hub: El estudiante utiliza una única vista de temperatura fija emparejada contra todas las vistas del profesor.
    • Path: La propia curva templada del estudiante se empareja contra la curva del profesor en las temperaturas correspondientes.
• Restricciones Teóricas: Los autores demuestran que el agrupamiento log-lineal (media geométrica) de las vistas templadas de un solo profesor colapsa en una única vista templada (Proposición 1). Por lo tanto, cualquier beneficio de la agregación de múltiples vistas debe provenir de mezclas aritméticas o baricentros basados en transporte, no de un simple promedio geomético.

Contribuciones Clave

1. Marco Formal: Un sistema de notación unificado para vistas templadas, familias de divergencia y la tricotomía de medias (log-lineal, aritmética, transporte).
2. Demostraciones Teóricas:
   • Prueba de que el agrupamiento log-lineal de las vistas templadas es equivalente a una sola temperatura, descartándolo como fuente de beneficio de múltiples vistas.
   • Una interpretación de puente de Schrödinger multimarinal que sugiere que la dispersión efectiva de las vistas (distancia ponderada desde $\tau=1$), en lugar del número de vistas, es la variable operativa.
3. Estudio Empírico Controlado: Una comparación sistemática de ocho métodos (incluyendo SFT, diversas divergencias KL, mezclas aritméticas, baricentros entrópicos y el método insignia basado en Sinkhorn) utilizando pares de Pythia con ajuste de instrucciones en el conjunto de datos Dolly-15k.
4. Salvaguardas de Diagnóstico: Introducción de un "guardián de fuerza del profesor sensible al margen" y una sonda de capacidad zero-shot para determinar si un experimento de KD es teóricamente capaz de producir ganancias, evitando la confusión de un "fallo de KD" con una "capacidad insuficiente del profesor".

Resultados y Leyes Empíricas
Los experimentos revelan que la efectividad de una pérdida de destilación no es una propiedad intrínseca de la función de pérdida, sino que depende del margen de techo (gap de techo) $\Gamma = \text{PPL}_{\text{SFT}} - \text{PPL}_{\text{Profesor}}$.

• Dos Regímenes:

  • Régimen de Techo Delgado ($\Gamma \lesssim 0.1$): El profesor ajustado con instrucciones apenas supera a un estudiante supervisado. En este régimen, ningún método de KD supera al SFT. Todos los métodos de KD incurren en un "impuesto" (peor rendimiento que SFT). Sin embargo, el objetivo basado en transporte (BaryOT) minimiza este impuesto y mantiene una mejor calibración. La fidelidad (acuerdo con el profesor) es anti-correlacionada con la generalización.
  • Régimen de Techo Real ($\Gamma \approx 1.5$): El profesor es genuinamente más fuerte. Aquí, el ranking se invierte: los métodos de KL puntuales fieles (JSKD, VanillaKD) superan al SFT, mientras que los métodos basados en transporte se quedan atrás. La fidelidad se correlaciona positivamente con la generalización.
  • Cruce: Los autores delimitan el punto de cruce $\Gamma^*$ entre 0.10 y 1.49.

• Tres Leyes Empíricas:

  1. Ley de Dispersión: El beneficio de la agregación de múltiples temperaturas crece monótonamente con la dispersión de temperatura efectiva de las vistas, no con el número de vistas.
  2. Las Vistas Desbloquean el Operador: La brecha de rendimiento entre el baricentro de transporte y la mezcla aritmética crece exactamente cuando las vistas están suficientemente dispersas. Las vistas agrupadas hacen que el baricentro y la mezcla sean casi idénticos.
  3. Imagen de Dos Regímenes: El signo de la correlación fidelidad-generalización cambia dependiendo de $\Gamma$.

• Hub vs. Path: En el objetivo insignia de Sinkhorn, el modo path (emparejar la curva del estudiante con la del profesor) supera consistentemente al modo hub. La sensibilidad al conjunto de vistas observada en el modo hub se identifica como un "artefacto de suavizado" causado por la aplicación de una temperatura fija al estudiante; el modo path elimina esta asimetría y mejora el rendimiento.

Significancia y Reivindicaciones
El artículo argumenta que la pregunta "¿qué pérdida de destilación es la mejor?" no es una propiedad fija de la función de pérdida, sino una función de las condiciones experimentales, específicamente del margen de techo $\Gamma$.

• Honestidad en Resultados Negativos: Los autores enfatizan que en techos delgados, cada método de KD paga un impuesto sobre el SFT. El sesgo de la divergencia de Sinkhorn no produce ningún beneficio mensurable en este régimen.
• Refinamiento de la Paradoja de la Fidelidad: El trabajo refina el hallazgo de Stanton et al. de que "estudiantes más fieles no siempre generalizan mejor". Los autores muestran que esta inversión es condicional: ocurre cuando el profesor es débil (techo delgado), pero se revierte cuando el profesor es fuerte (techo real).
• Alcance Práctico: La visión distribucional es valiosa, pero sus ingredientes (vistas de múltiples temperaturas, geometría de transporte) solo rinden frutos cuando el profesor tiene una ventaja de capacidad genuina y las vistas están suficientemente dispersas. El artículo concluye que el objetivo de transporte es la elección robusta cuando el profesor tiene "poco que enseñar", mientras que el KL puntual es superior cuando el profesor es un "techo real".