WDL-OPD: Weak-Driven On-Policy Distillation via Mixture-Constrained Co-Training

Resumen Técnico: WDL-OPD

Planteamiento del Problema
La destilación on-policy (OPD) aborda el desajuste entre el entrenamiento y la prueba (train–test state mismatch) inherente a la destilación offline, entrenando un modelo estudiante mediante trayectorias que este genera por sí mismo, utilizando un profesor congelado para la supervisión a nivel de token. Sin embargo, la OPD introduce una inestabilidad dinámica: cada actualización de la política altera el comportamiento del estudiante, lo que a su vez desplaza la distribución de los estados visitados en las ejecuciones (rollouts) subsiguientes. Este bucle de retroalimentación puede conducir a un entrenamiento inestable, crecimiento de la entropía o degradación de la trayectoria, particularmente cuando los patrones de razonamiento del profesor y del estudiante son incompatibles o cuando el conjunto de tokens de alta probabilidad compartidos es pequeño. Los métodos existentes como OPD2 y W2S-OPD intentan estabilizar el aprendizaje utilizando señales "delta" estáticas (diferencias entre un profesor y su modelo base) o proxies contrastivos congelados, pero dependen de una única política entrenable.

Metodología: WDL-OPD
Los autores proponen WDL-OPD (Weak-Driven On-Policy Distillation), un marco de co-entrenamiento con mezcla restringida que utiliza dos políticas entrenables: una política ancla ($\pi_A$) y una política auxiliar ($\pi_U$).

1. Ejecución (Rollout) y Construcción de Soporte: Solo la política ancla genera trayectorias ($y \sim \pi_A$). Para cada estado visitado (prefijo) $s_t$, se construye un conjunto de soporte compartido $K_t$ que consiste en los top-$k$ tokens con los logits más altos del ancla. Tanto la política ancla como la auxiliar se evalúan sobre este mismo soporte para asegurar la consistencia distributiva.
2. Mezcla Geométrica: El método define una puntuación de log-score mixta $m_t$ como una media geométrica ponderada de las log-probabilidades del ancla y la auxiliar sobre el soporte compartido:
   $$m_t(v) = \lambda \ell_A(v) + (1-\lambda) \ell_U(v)$$
   Esta mezcla se renormaliza para formar una distribución $q_M$.
3. Función de Pérdida: El objetivo de entrenamiento minimiza la divergencia de Kullback-Leibler inversa (KL) entre esta distribución de mezcla $q_M$ y la distribución del profesor congelado $\pi_T$:
   $$L = D_{KL}(q_M \parallel \pi_T)$$
   Crucialmente, los gradientes de esta pérdida actualizan tanto la política ancla como la política auxiliar simultáneamente.
4. Límite Teórico: El artículo demuestra que, si la política auxiliar permanece congelada, el método se reduce algebraicamente a un objetivo de proxy "ancla-más-contraste", estrechamente relacionado con OPD2 y W2S-OPD. Sin embargo, al mantener la auxiliar entrenable, el método introduce "grados de libertad a nivel de rama", permitiendo que las dos políticas dividan la señal de corrección del profesor de acuerdo con sus distintos Jacobianos y estados del optimizador.

Contribuciones Clave

• Formulación: Los autores formulan un algoritmo de OPD de dos políticas con mezcla restringida donde la distribución auxiliar es parte del estado de optimización en lugar de un objetivo fijo.
• Caracterización Teórica: Caracterizan la relación entre este método y los enfoques de delta estático, mostrando que el entrenamiento conjunto crea un espacio nulo a nivel de rama que un delta estático no puede expresar.
• Consolidación Empírica: El artículo proporciona un registro experimental consolidado a través de dos escalas de modelos (1.7B y 4B) y dos dominios (razonamiento matemático y generación de código), incluyendo ejecuciones negativas y reevaluaciones independientes para resaltar los modos de fallo.

Resultados Experimentales
Los experimentos se realizaron utilizando modelos Qwen3 (1.7B y 4B) en razonamiento matemático (MATH500, GSM8K, etc.) y generación de código (HumanEval+, MBPP+, etc.).

• Rendimiento: En los cuatro entornos de escala–dominio, WDL-OPD produjo el checkpoint del estudiante más fuerte en comparación con los baselines de política única registrados.
  • Matemáticas (4B): La precisión en MATH500 mejoró de 0.630 (mejor de política única) a 0.685 (auxiliar de WDL-OPD).
  • Matemáticas (1.7B): La precisión en MATH500 mejoró de 0.521 a 0.585.
  • Código (4B): Las configuraciones de OPD de política única exhibieron inestabilidad (crecimiento de entropía o repetición), fallando en producir un candidato final saludable. WDL-OPD logró un puntaje Dev500 reevaluado de 0.637.
  • Código (1.7B): El rendimiento de la política única se degradó (0.106 frente al base 0.184), mientras que WDL-OPD alcanzó 0.375.
• Dinámica de Ramas: Las dos ramas no se comportaron como copias redundantes. En el entorno de matemáticas 4B, la auxiliar (inicialmente más débil) superó al ancla para convertirse en el mejor artefacto final. En el entorno de matemáticas 1.7B, el ancla se mantuvo superior. Esto sugiere que la identidad y especialización de la rama son empíricamente significativas.
• Estabilidad: En la generación de código, donde la OPD de política única sufrió frecuentemente de crecimiento de entropía y degradación de la trayectoria, el enfoque de co-entrenamiento evitó estas firmas de fallo específicas el tiempo suficiente para producir checkpoints útiles, aunque no eliminó todas las patologías (por ejemplo, la repetición apareció eventualmente en la ejecución de 4B).

Significancia y Reivindicaciones
El artículo posiciona a WDL-OPD no como una solución causal universal para todas las inestabilidades de la OPD, sino como un método que respalda una hipótesis de estabilización. Los autores declaran explícitamente que sus resultados son descriptivos más que una reivindicación causal controlada, señalando confusiones como diferencias en currículos, estados de inicialización y presupuestos de cómputo entre su método y los baselines de política única.

La significancia principal reside en la hipótesis del buffer de optimización: la rama auxiliar entrenable puede absorber parte de una señal de corrección aguda, evitando el cambio inmediato y potencialmente desestabilizador de la distribución del estado de la ejecución que ocurre cuando una única política debe manejar tanto la corrección como la generación del estado. El artículo concluye que el siguiente paso científico no son reivindicaciones más amplias, sino una matriz de comparación controlada que aísle el acoplamiento de la mezcla de los factores de inicialización, capacidad y diseño del objetivo.