$\mathcal{X}$-KD: General Experiential Knowledge Distillation for Large Language Models

1. El Problema

La Distilación de Conocimiento (KD) para Modelos de Lenguaje Grandes (LLMs) es crucial para transferir capacidades de modelos grandes (maestros) a modelos más pequeños (estudiantes) debido a las limitaciones de recursos computacionales. Sin embargo, los enfoques actuales presentan limitaciones fundamentales:

• Imitación de comportamiento vs. Entorno de aprendizaje: La mayoría de los métodos existentes (como la clonación de comportamiento o BC) se centran únicamente en imitar las salidas o comportamientos del modelo maestro. Ignoran el entorno de aprendizaje original (la función de recompensa subyacente) que moldeó ese conocimiento.
• Desajuste de distribución y complejidad: Los métodos de caja blanca sufren de desajuste de distribución y costos elevados en la generación de datos. Los métodos de caja negra a menudo se basan en demostraciones o retroalimentación de prompts finos que no maximizan directamente la recompensa del entorno original.
• Limitaciones de enfoques recientes: Trabajos concurrentes como MiniLLM intentan abordar esto mediante aprendizaje por refuerzo inverso (IRL), pero utilizan gradientes de política que se desvían del paradigma de aprendizaje supervisado, introduciendo inestabilidad en el entrenamiento y mayor complejidad de integración.

2. Metodología: X-KD (Distilación de Conocimiento Experiential)

Los autores proponen X-KD, un marco general que permite a los modelos estudiantes aprender en el entorno original de aprendizaje del maestro, en lugar de solo imitar su comportamiento.

• Fundamento Teórico: Se inspira en la Teoría del Aprendizaje Experiential y el Aprendizaje por Refuerzo Inverso Bayesiano (BIRL). Utiliza el marco AVRIL (Approximated Variational Reward Imitation Learning) para modelar conjuntamente la función de recompensa original del maestro y la política del estudiante.
• Formulación del Problema:
  • Se trata la generación de lenguaje natural (NLG) como un Proceso de Decisión de Markov (MDP).
  • El objetivo es estimar la función de recompensa oculta $R_\pi$ del maestro y asegurar que el estudiante maximice esta recompensa.
  • Se formula como un problema de inferencia variacional donde se aprende una distribución de recompensas $q_\phi(R)$ que aproxima la posterior $p(R|\pi)$.
• Objetivo de Entrenamiento:
  La función de pérdida total combina el objetivo de distilación estándar con un término de regularización experiential:
  $$L_{total} = L_{distilación} + L_{experiential}$$
  Donde $L_{experiential}$ asegura que la política del estudiante sea consistente con la función de recompensa aprendida del entorno original, utilizando errores de diferencia temporal (TD-error) y una distribución de recompensas parametrizada.
• Versatilidad del Marco:
  X-KD es lo suficientemente flexible para integrarse en tres tipos de distilación:
  1. Distilación a Nivel de Secuencia: Basada en la probabilidad de secuencias generadas por el maestro.
  2. Distilación Supervisada (Forward-KL): Minimiza la divergencia KL a nivel de token.
  3. Distilación Generalizada (GKD): Utiliza divergencias generalizadas (como JS) y datos de política en línea.
  • Simplicidad: A diferencia de MiniLLM, X-KD mantiene el paradigma de aprendizaje supervisado, lo que facilita su integración en pipelines de entrenamiento existentes sin necesidad de técnicas de estabilización complejas de RL.
  • Caja Negra y Blanca: Funciona tanto con acceso a la distribución de salida del maestro (caja blanca) como solo con sus respuestas finales (caja negra).

3. Contribuciones Clave

1. Propuesta de X-KD: Un nuevo marco basado en IRL Bayesiano que permite a los estudiantes aprender en el entorno de recompensa original del maestro, superando la mera imitación de comportamiento.
2. Flexibilidad y Simplicidad: Demuestra que se puede integrar con métodos de KD existentes (Secuencial, Supervisado, Generalizado) manteniendo la estructura de entrenamiento supervisado, evitando la complejidad de los métodos basados en gradientes de política pura.
3. Evidencia Empírica: Resultados que muestran superioridad sobre líneas base (GKD, MiniLLM, KD estándar) en múltiples tareas, con mejoras en el equilibrio rendimiento-diversidad y eficiencia de datos.

4. Resultados Experimentales

Los experimentos se realizaron en tres tareas de generación de lenguaje natural (NLG): Resumen Abstractivo (XSum), Traducción Automática (WMT14 en-de) y Razonamiento Matemático (GSM8k).

• Rendimiento General: X-KD superó consistentemente a las líneas base (GKD, MiniLLM, KD secuencial) en la mayoría de las configuraciones de tamaño de modelo (de grande a pequeño/base).
• Equilibrio Rendimiento-Diversidad:
  • X-KD logró un mejor compromiso entre la calidad de la generación (ROUGE/BLEU) y la diversidad de las salidas (medida por SelfBLEU).
  • A medida que aumentaba la diversidad (temperatura de muestreo más alta), X-KD degradaba su rendimiento más lentamente que los métodos basales, sugiriendo que aprende estrategias de generalización más robustas en lugar de sobreajustarse al comportamiento del maestro.
• Eficiencia de Datos:
  • X-KD alcanzó un rendimiento comparable a los métodos basales utilizando significativamente menos datos de entrenamiento (ej. 50% o 75% de los datos), lo que indica una mayor eficiencia en la captura de principios fundamentales de la tarea.
• Distilación de Caja Negra: En experimentos usando Llama-2-7b como maestro de caja negra, la variante SeqXKD obtuvo una mayor tasa de victoria frente a SeqKD según evaluaciones de GPT-4.

5. Significado e Impacto

• Cambio de Paradigma: X-KD marca un paso adelante al tratar la distilación no solo como una transferencia de distribución de probabilidad, sino como la transferencia del entorno de aprendizaje y la función de recompensa que guió al maestro.
• Aplicabilidad Práctica: Al mantenerse dentro del paradigma de aprendizaje supervisado, X-KD es más fácil de implementar en ingeniería de LLMs que los métodos de RL puros, ofreciendo una vía práctica para mejorar la eficiencia y calidad de los modelos pequeños.
• Limitaciones y Futuro: Los autores reconocen que el rendimiento es sensible a la sintonización del peso de la regularización experiential ($\lambda$) y que el enfoque actual prioriza el escenario de caja blanca, dejando el escenario de caja negra para estudios más sistemáticos futuros.

En resumen, X-KD ofrece una solución elegante y general para la distilación de LLMs, demostrando que recuperar el "contexto de experiencia" del maestro es clave para entrenar estudiantes más robustos, eficientes y diversos.