Rethinking Reward Supervision: Rubric-Conditioned Self-Distillation

Resumen Técnico: Replanteando la Supervisión de Recompensas: Auto-destilación Condicionada por Rúbricas

Declaración del Problema

El post-entrenamiento de modelos de lenguaje de razonamiento (LLMs) depende actualmente de dos paradigmas, ambos de los cuales sufren limitaciones significativas al proporcionar retroalimentación detallada:

1. Aprendizaje por Refuerzo con Recompensas Verificables (RLVR): Métodos como la Optimización de Política de Grupo Relativa (GRPO) típicamente comprimen la retroalimentación evaluativa en una señal escalar dispersa (por ejemplo, corrección binaria o una puntuación única). Aunque es efectivo para tareas verificables, este enfoque oscurece el porqué un trayecto tuvo éxito o falló, creando un "cuello de botella de asignación de crédito" donde el modelo no puede distinguir entre pasos de razonamiento intermedios correctos e incorrectos.
2. Auto-destilación On-Policy (OPSD): Estos métodos reemplazan las recompensas dispersas con guía a nivel de token densa proveniente de un modelo profesor. Sin embargo, los enfoques existentes condicionan al profesor a una única trayectoria de referencia privilegiada (una solución de oro específica). Esto vincula la supervisión a un camino de razonamiento particular, lo que puede sobre-especificar cómo debe producirse una respuesta e ignorar el espacio de derivaciones alternativas válidas. Además, si el estudiante se desvía ligeramente de este único camino, el profesor puede fomentar una revisión global en lugar de una corrección localizada.

Además, mientras que las rúbricas (criterios de evaluación estructurados) ofrecen una representación de la calidad más rica e interpretable que las puntuaciones escalares, el trabajo previo ha descartado mayoritariamente esta estructura durante la optimización al agregar los juicios de la rúbrica en una única recompensa escalar.

Metodología: Auto-destilación Condicionada por Rúbricas (RCSD)

Los autores proponen la Auto-destilación Condicionada por Rúbricas (RCSD), un marco que trata las rúbricas como supervisión estructurada y privilegiada del lado del profesor para la auto-destilación on-policy. En lugar de colapsar la retroalimentación de la rúbrica en un escalar, RCSD condiciona al modelo profesor en la información de la rúbrica a nivel de criterio para proporcionar guía densa a nivel de token sobre las propias trayectorias muestreadas por el estudiante.

El marco opera mediante un pipeline de dos etapas:

Etapa I: Aprendizaje de un Generador de Rúbricas

Para abordar el alto costo de obtener rúbricas específicas de cada instancia, los autores entrenan un generador de rúbricas dedicado.

• Política del Profesor: Recibe la pregunta ($x$) y una respuesta de referencia ($y^\star$) como entrada privilegiada para inferir los criterios de evaluación apropiados.
• Política del Estudiante: Observa únicamente la pregunta ($x$) y aprende a generar una rúbrica ($r$) que aproxime los criterios del profesor.
• Objetivo: El estudiante es entrenado mediante destilación on-policy para coincidir con la distribución de tokens del profesor a lo largo de la propia trayectoria de la rúbrica muestreada por el estudiante. Esto amortiza la generación de criterios específicos de la instancia en un modelo reutilizable.

Etapa II: Razonamiento Condicionado por Rúbricas

El generador de rúbricas entrenado se utiliza para guiar el entrenamiento del modelo de razonamiento.

• Política del Profesor: Recibe la pregunta ($x$) y la rúbrica generada ($r$) como contexto privilegiado.
• Política del Estudiante: Observa únicamente la pregunta ($x$) y genera una trayectoria de razonamiento ($\hat{y}$).
• Objetivo: El estudiante minimiza la divergencia de Kullback-Leibler (KL) entre su distribución y la distribución del profesor, donde el profesor está condicionado por la rúbrica.
  $$\mathcal{L}_{reason} = \mathbb{E}_{\hat{y} \sim p_S} \left[ \sum_{t=1}^{|\hat{y}|} D_{KL} \left( p_T(\cdot \mid x, r, \hat{y}_{<t}) \parallel p_S(\cdot \mid x, \hat{y}_{<t}) \right) \right]$$
• Diseño Clave: Al condicionarse en la rúbrica en lugar de en una única respuesta de referencia, el profesor proporciona una guía consciente de los criterios. Esto permite que la señal de optimización distinga diferentes dimensiones de corrección parcial (por ejemplo, un paso correcto seguido de un error local) en lugar de comprimirlas en un solo puntaje o en un único camino de referencia. Los autores utilizan la divergencia KL Forward para fomentar que el estudiante cubra el soporte de la distribución del profesor condicionada por la rúbrica, preservando la diversidad entre caminos de razonamiento válidos.

Contribuciones Clave

1. Nueva Interfaz de Supervisión: RCSD replantea las rúbricas no solo como herramientas de puntuación, sino como interfaces de supervisión estructuradas que moldean el aprendizaje a nivel de token durante la optimización.
2. Pipeline de Dos Etapas: Los autores introducen un método escalable para aprender rúbricas específicas de la tarea automáticamente (Etapa I) y luego usarlas para guiar el razonamiento (Etapa II), eliminando la dependencia de la anotación manual de rúbricas o de recompensas costosas de "LLM-as-a-Judge" durante la fase de destilación.
3. Aprendizaje On-Policy Consciente de Criterios: A diferencia de la OPSD estándar que se condiciona en una trayectoria única, RCSD se condiciona en criterios, permitiendo una asignación de crédito fina que penaliza errores locales específicos mientras preserva los pasos correctos, incluso en las propias trayectorias muestreadas del estudiante.

Resultados Experimentales

El método fue evaluado en un conjunto diverso de benchmarks de razonamiento científico, incluyendo tareas verificables (GPQA-DIAMOND, SciBench, PIQA) y tareas abiertas no verificables (ResearchQA, RubricHub).

• Desempeño: RCSD logró la mejor puntuación promedio general (70.6) a través de los benchmarks. Superó al modelo base Qwen3-8B por 4.7 puntos, superando a GRPO por 1.0 punto y a OPSD por 0.9 puntos en promedio.
• Especificidad de Dominio: Las ganancias fueron más pronunciadas en tareas de razonamiento científico y basadas en rúbricas (por ejemplo, +8.2 puntos en ResearchQA, +4.9 puntos en RubricHub), donde las recompensas escalares a menudo fallan en capturar la calidad de la respuesta.
• Generalización: El método demostró una robusta generalización a benchmarks médicos fuera de dominio (MedMCQA, PubMedQA) sin olvido catastrófico, sugiriendo que los criterios aprendidos capturan estructuras de razonamiento generales.
• Estudios de Ablación:
  • Tipo de Pérdida: Se encontró que la KL Forward era el objetivo más efectivo, produciendo el mayor desempeño promedio y manteniendo una mayor entropía (diversidad) del estudiante en comparación con Reverse KL o la divergencia JS.
  • Fuente de la Rúbrica: El uso de rúbricas generadas por el generador aprendido en la Etapa I produjo resultados competitivos con el uso de rúbricas de referencia de verdad de campo (ground-truth), validando la escalabilidad del enfoque de dos etapas.
  • Robustez: El método se mantuvo efectivo incluso con entradas de rúbrica degradadas (genéricas, ruidosas o reducidas), aunque las rúbricas aprendidas específicas de la instancia proporcionaron la supervisión más fuerte.

Significación y Reivindicaciones

El artículo afirma que RCSD ofrece un enfoque libre de verificadores para el post-entrenamiento de modelos de razonamiento de lenguaje abierto. Al preservar la estructura a nivel de criterio durante la optimización y convertir las rúbricas en guía densa a nivel de token, el método aborda el cuello de botella de la asignación de crédito inherente a la optimización de recompensa escalar y el sesgo de trayectoria específica de la destilación condicionada por referencia.

Los autores postulan que RCSD es particularmente significativo para tareas donde las respuestas de alta calidad no pueden ser capturadas completamente por la verificación automática (respuestas de coincidencia exacta o pruebas unitarias) o por recompensas de resultado escalares. Al permitir que el modelo aprenda qué propiedades debe satisfacer una respuesta (vía rúbricas) en lugar de cómo imitar un camino de solución específico, RCSD facilita mejoras de razonamiento más flexibles y robustas. El trabajo sugiere que la retroalimentación estructurada y multidimensional es un signo de supervisión superior para tareas de razonamiento complejo en comparación tanto con las recompensas dispersas como con la destilación de trayectoria única.