Hybrid Policy Distillation for LLMs

Resumen Técnico: Hybrid Policy Distillation (HPD) para LLMs

1. El Problema

La compresión de Modelos de Lenguaje Grandes (LLMs) mediante Destilación de Conocimiento (KD) es crucial para su despliegue práctico, pero enfrenta desafíos fundamentales derivados de la interacción compleja entre tres ejes de diseño:

1. Dirección de la divergencia:
   • KL Directo (FKLD): Promueve la cobertura de modos (mode coverage), pero tiende a generar predicciones excesivamente suaves (over-smoothing), perdiendo precisión.
   • KL Inverso (RKLD): Se enfoca en buscar modos (mode-seeking), lo que puede llevar a inestabilidad si la brecha entre estudiante y profesor es grande, y a ignorar salidas válidas de baja probabilidad.
2. Estrategia de optimización: Diferencias entre minimizar una pérdida (loss) directa o usar señales de recompensa.
3. Régimen de datos:
   • Off-policy (fuera de política): Usa datos estáticos (generados por el profesor o ground-truth). Es eficiente computacionalmente pero sufre de una desconexión entre el comportamiento de entrenamiento e inferencia.
   • On-policy (en política): El estudiante genera sus propias muestras. Mejora la alineación pero introduce un alto costo computacional y desplazamiento de distribución.

La elección de un solo enfoque (ej. solo SFT, solo FKLD o solo RKLD) suele ser insuficiente para equilibrar la cobertura de modos, la búsqueda de modos y la estabilidad del entrenamiento.

2. Metodología: Hybrid Policy Distillation (HPD)

Los autores proponen HPD, un marco unificado que reformula la destilación como un objetivo de verosimilitud logarítmica reponderada a nivel de token. HPD integra las ventajas complementarias de FKLD y RKLD, combinando datos off-policy con muestreo on-policy ligero.

Mecanismos Clave:

• Visión Unificada: Se formaliza la pérdida como $L(\theta) = -E [w(a_t|s_t) \log q_\theta(a_t|s_t)]$, donde el término de reponderación $w$ captura la discrepancia local entre el estudiante y el profesor.
• Estimador K1 Negativo: Se utiliza un estimador Monte Carlo (K1) de la divergencia KL inversa como señal de recompensa a nivel de token.
  • Si el estudiante subestima un token experto (ground-truth), el estimador es positivo, activando un aprendizaje tipo FKLD (aumentar la probabilidad).
  • Si el estudiante sobreestima un token no experto (muestreado por el estudiante), el estimador es negativo, activando un aprendizaje tipo RKLD (suprimir la probabilidad y redistribuir la masa hacia los tokens expertos).
• Estrategia Híbrida Asimétrica:
  1. Token Experto ($a^*_t$): Se evalúa la brecha KL. Si el estudiante lo subestima, se refuerza con un peso de FKLD. Si el estudiante sobreestima un token no experto, se duplica el peso de refuerzo del token experto para redirigir la masa de probabilidad suprimida hacia la acción correcta.
  2. Token Muestreado ($a_t \sim q_\theta$): Se muestrea un token no experto. Si el estudiante lo sobreestima (es decir, lo considera más probable que al profesor), se penaliza fuertemente (peso negativo) para suprimir comportamientos irracionales.
• Eficiencia: A diferencia de los métodos on-policy puros que requieren rollouts completos, HPD realiza un muestreo ligero de un solo token no experto por estado, manteniendo la eficiencia de la supervisión one-hot mientras corrige la distribución.

3. Contribuciones Clave

1. Visión Unificada: Presentan una perspectiva de "verosimilitud reponderada" que conecta SFT, FKLD y RKLD, permitiendo entender y diseñar nuevos métodos de destilación de manera sistemática.
2. Propuesta HPD: Un algoritmo que hibrida señales de KL directo e inverso utilizando datos off-policy y muestreo on-policy ligero, logrando un equilibrio entre cobertura de modos y búsqueda de modos sin incurrir en el alto costo de los rollouts completos.
3. Validación Empírica: Demuestran que HPD mejora la estabilidad de la optimización, la eficiencia computacional y el rendimiento final en comparación con los pipelines estándar (SFT + OPD) en diversas tareas y escalas de modelos.

4. Resultados Experimentales

Los autores evaluaron HPD en tareas de razonamiento matemático (generación larga), diálogo y código (generación corta), utilizando familias de modelos como Qwen2.5 y LLaMA 3.

• Razonamiento Matemático (Off-policy):
  • HPD superó consistentemente a SFT, SeqKD, RKLD y JSD.
  • Logró mejoras masivas en modelos pequeños: mejoró el rendimiento de Qwen 2.5 3B en un 41% y de LLaMA 3 3B en un 77.9% en promedio, acercándolos al rendimiento de modelos más grandes.
  • Evitó el colapso de entropía (overfitting) observado en SFT, manteniendo una alineación estable con el profesor durante el entrenamiento.
• Personalización y Diálogo (Off-policy):
  • En benchmarks como MT-Bench y AlpacaEval, HPD mostró la mejor capacidad para preservar la coherencia conversacional y el entendimiento contextual en diálogos multi-turno.
• Codificación (Off-policy):
  • Obtuvo el mejor rendimiento promedio en HumanEval y MBPP, demostrando mayor robustez y menor varianza que otros métodos de destilación.
• Datos On-policy (OPD):
  • HPD sirvió como una inicialización superior para la destilación en política (OPD). Al combinar HPD con OPD, se alcanzaron los puntajes más altos, superando significativamente a los baselines que usan solo SFT inicial.
• Ablación: Se confirmó que tanto el muestreo del estudiante como la operación de refuerzo (Reinforce) son esenciales; sin ellos, el modelo converge prematuramente o la alineación es más lenta.

5. Significado e Impacto

• Eficiencia vs. Rendimiento: HPD resuelve el dilema clásico de la destilación: ofrece la estabilidad y cobertura de los métodos off-policy con la capacidad de corrección de errores de los métodos on-policy, pero a una fracción del costo computacional.
• Inicialización Robusta: Proporciona una base sólida para etapas posteriores de alineación (como DPO o RLHF), permitiendo que modelos pequeños alcancen capacidades de razonamiento y generación de alto nivel.
• Generalización: La metodología no solo mejora el rendimiento en dominio, sino que también muestra una mejor generalización fuera de dominio (ej. en GPQA), indicando que la transferencia de señales de decisión es más efectiva que la mera imitación de comportamiento.

En conclusión, Hybrid Policy Distillation representa un avance significativo en la compresión de LLMs, ofreciendo un marco teórico unificado y un algoritmo práctico que supera las limitaciones de las técnicas de destilación tradicionales, permitiendo desbloquear el potencial de modelos pequeños para tareas complejas.