Resumen Técnico: ARKD – Destilación de Divergencia KL Bidireccional Guiada por Aprendizaje por Refuerzo Adaptativo

1. Planteamiento del Problema

La Destilación de Conocimiento (KD) es una técnica primordial para comprimir Modelos de Lenguaje Extensos (LLMs) con el fin de reducir los costes de inferencia manteniendo el rendimiento. Sin embargo, los métodos tradicionales de KD que dependen de un único objetivo de divergencia de Kullback-Leibler (KL) enfrentan un compromiso fundamental en la generación de texto de código abierto:

• KL Directa (FKL): Fomenta un comportamiento de "cobertura de modos" (mode-covering), donde el estudiante intenta cubrir todos los modos de la distribución del profesor. Esto a menudo conduce a la sobreestimación de probabilidades en regiones de baja confianza y de cola larga, degradando la calidad de la generación.
• KL Inversa (RKL): Fomenta un comportamiento de "búsqueda de modos" (mode-seeking), centrándose en los modos dominantes e ignorando las regiones de baja probabilidad. Si bien esto evita la sobreestimación de la cola larga, compromete la diversidad y la cobertura, lo que puede conducir al colapso de modo.

Los métodos de interpolación estática existentes (por ejemplo, ponderación fija de 0.5:0.5) no logran capturar la naturaleza dinámica de la brecha entre estudiante y profesor a medida que el entrenamiento evoluciona. Un coeficiente estático no puede adaptarse a los requisitos cambiantes de la alineación de distribuciones, que requieren diferentes equilibrios de FKL y RKL en diferentes etapas del entrenamiento.

2. Metodología: ARKD

Los autores proponen ARKD (Destilación de Divergencia KL Bidireccional Guiada por Aprendizaje por Refuerzo Adaptativo), un marco que asigna dinámicamente pesos a FKL y RKL utilizando un agente de aprendizaje por refuerzo (RL).

Fundamento Teórico

El artículo establece primero la consistencia teórica de FKL y RKL. El análisis muestra que ambos objetivos convergen a la misma solución ($q_\theta = p$) cuando los gradientes son cero. Esto confirma que ambas divergencias son válidas para el emparejamiento de distribuciones pero exhiben comportamientos complementarios durante la trayectoria de optimización, justificando su uso conjunto.

Marco Basado en RL

ARKD replantea el problema de la asignación de pesos como un proceso de toma de decisiones secuenciales (Proceso de Decisión de Markov):

• Estado ($s$): Un vector de 6 dimensiones que captura la dinámica actual del entrenamiento, incluyendo las entropías del profesor y del estudiante ($H_T, H_S$), las varianzas ($\sigma^2_T, \sigma^2_S$) y los valores actuales de pérdida de FKL/RKL ($L_{FKL}, L_{RKL}$).
• Acción ($\alpha$): Una red de política ($\pi_\phi$) genera un peso escalar $\alpha \in (0, 1)$, donde $\alpha$ pondera el término FKL y $(1-\alpha)$ pondera el término RKL.
• Recompensa ($r$): Definida como la pérdida de destilación negativa ($r = -L_{distil}$) para el lote actual. El objetivo es maximizar la recompensa acumulada (minimizar la pérdida total) a lo largo de la trayectoria de entrenamiento.
• Optimización:
  • Actualización del Estudiante: Minimiza la pérdida de destilación ponderada más una pérdida de modelado de lenguaje.
  • Actualización de la Política: Utiliza el algoritmo REINFORCE con una base de Media Móvil Exponencial (EMA) y regularización de entropía para estabilizar el entrenamiento y evitar la saturación de los límites.

Dinámica de Entrenamiento

La red de política aprende a navegar el compromiso de forma adaptativa:

1. Fase de Exploración: Al principio del entrenamiento, la política explora un rango más amplio de $\alpha$ (por ejemplo, 0.3–0.6), enfatizando FKL para asegurar una cobertura amplia y prevenir el colapso de modo.
2. Fase de Convergencia: A medida que la distribución del estudiante se aproxima a la del profesor, la política se desplaza hacia valores de $\alpha$ más bajos (dominancia de RKL) para centrarse en una alineación precisa y evitar la sobreestimación de la cola larga.
3. Fase de Estabilidad: La política se estabiliza en un $\alpha$ específico (por ejemplo, $\approx 0.18$) para el refinamiento final.

3. Contribuciones Clave

1. Análisis Teórico: El artículo proporciona un análisis teórico que establece la complementariedad y la consistencia de convergencia de FKL y RKL, formando la base para la destilación adaptativa.
2. Marco ARKD: Introduce un novedoso marco de optimización conjunta basado en RL donde una red de política pondera dinámicamente FKL y RKL basándose en las características dependientes del estado, yendo más allá de las heurísticas estáticas.
3. Validación Empírica: Experimentos extensos demuestran que ARKD mejora consistentemente la calidad de la generación, el ajuste de la distribución y la generalización a través de múltiples evaluaciones y escalas de modelos.

4. Resultados Experimentales

Los autores evaluaron ARKD en tareas de ajuste de instrucciones utilizando configuraciones de GPT-2 (1.5B $\to$ 120M/340M) y LLaMA (13B $\to$ 7B).

• Rendimiento de Distribución Interna (ID): En el benchmark DollyEval, ARKD logró el BertScore más alto y puntuaciones de Rouge-L competitivas o superiores en comparación con las líneas base (SFT, SeqKD, FKL, RKL, FKL+RKL y MiniLLM). Por ejemplo, en GPT-2 1.5B $\to$ 120M, ARKD logró un Rouge-L de 24.5 y un BertScore de 0.815, superando la fuerte línea base de MiniLLM.
• Generalización de Distribución Externa (OOD): ARKD demostró una generalización superior en conjuntos de datos no vistos (SelfInst, VicunaEval, S-NI, UnNI), superando consistentemente a todas las líneas base en ambas familias de modelos.
• Estudios de Ablación:
  • vs. Programación Estática/Lineal: ARKD superó significativamente a la ponderación estática ( $\alpha$ fija) y a las estrategias de programación lineal (decaimiento/incremento basado en el tiempo), las cuales fallaron en capturar las complejas dinámicas de alineación.
  • vs. Heurísticas Voraces (Greedy): ARKD superó a las heurísticas voraces (que seleccionan la divergencia de menor pérdida en cada paso) por 0.4–0.6 puntos. Se encontró que el enfoque voraz era miope, degenerando a menudo en un RKL puro, mientras que ARKD mantuvo un componente beneficioso de FKL mediante la mezcla continua.
  • Análisis de Trayectoria: La trayectoria aprendida de $\alpha$ confirmó una estrategia inteligente: énfasis temprano en FKL para la cobertura, transicionando hacia RKL para la precisión, un patrón no alcanzable mediante diseño manual.

5. Significado y Reivindicaciones

El artículo afirma que ARKD hace avanzar el estado del arte de la destilación de conocimiento adaptativa al abordar las limitaciones de los métodos basados en KL estáticos. Al aprovechar el aprendizaje por refuerzo para equilibrar dinámicamente los comportamientos de cobertura de modos y búsqueda de modos, ARKD logra un mejor alineamiento de la distribución que mejora tanto la calidad de la generación como la generalización.

Los autores posicionan a ARKD como un enfoque fundamentado que valida la necesidad de una ponderación dinámica y consciente del estado por encima de las heurísticas fijas. Señalan que, si bien la política actual depende de un vector de estado diseñado a mano y recompensas de lote inmediatas, el marco es compatible con familias de $f$-divergencia más amplias y paradigmas de política en línea (on-policy). El trabajo proporciona una base sólida para la compresión eficiente de LLMs al demostrar que la optimización adaptativa guiada por RL puede superar consistentemente a las líneas base estáticas y voraces en diversas tareas y escalas de modelos.