Resumen Técnico: Leer la Traza, Dirigir el Camino (CAPR)

Planteamiento del Problema

Los Modelos de Lenguaje de Difusión (dLLMs) generan texto mediante la eliminación iterativa de ruido sobre una secuencia completamente enmascarada a través de pasos paralelos. A diferencia de los modelos autorregresivos, los dLLMs exponen una distribución predictiva para cada posición enmascarada en cada paso, creando una rica "traza de eliminación de ruido" que revela qué tokens se vuelven confiables y estables frente a aquellos que permanecen inestables u oscilan.

Los métodos actuales de post-entrenamiento de Aprendizaje por Refuerzo (RL) para dLLMs no aprovechan esta traza de manera efectiva, cayendo en dos extremos:

1. Rollouts Planos (Flat Rollouts): Asignan una única recompensa de resultado a todo el trayecto. Aunque son computacionalmente económicos, no proporcionan señal alguna sobre qué pasos de eliminación de ruido o bloques de tokens fueron críticos para el resultado final.
2. Rollouts de Árbol (Tree Rollouts): Ramifican continuaciones alternativas para comparar recompensas y generar ventajas a nivel de segmento. Si bien proporcionan una granularidad más fina, el costo computacional crece rápidamente con el factor de ramificación y la profundidad, lo que a menudo los hace prohibitivamente costosos.

La pregunta central abordada por este trabajo es: ¿Puede la propia traza de eliminación de ruido proporcionar señales de supervisión similares a un árbol sin el costo computacional de una expansión completa del árbol?

Metodología: CAPR

Los autores introducen CAPR (Refinamiento de Trayectoria Amortizado por Caché), un algoritmo de RL basado en bloques diseñado para extraer ventajas por bloque de la traza de eliminación de ruido. CAPR opera a través de tres mecanismos principales:

1. Cache & Steer (Caché y Dirección)

CAPR registra los primeros pasos de eliminación de ruido para construir un estado de trayectoria ($h_j$) compacto para cada posición $j$. Este estado resume:

• La confianza en el token superior actual.
• La entropía de la distribución de probabilidad.
• La estabilidad del token superior a través de los pasos.

Este estado se actualiza mediante un mecanismo tipo GRU donde la alta entropía reinicia la memoria (olvidando evidencia obsoleta) y la estabilidad la preserva. El estado de la trayectoria cumple dos funciones:

• Dirección (Steering): Sesga los logits para el siguiente paso inverso. Las posiciones estables y confiables reciben un impulso positivo, mientras que las posiciones oscilantes son amortiguadas.
• Control de Desenmascarado (Unmasking Control): Influye en qué posiciones enmascaradas se revelan (desenmascaran) basándose en su certeza, asegurando que el modelo se concentre en compromisos estables.

2. Branch & Prune (Ramificación y Poda)

Para generar señales de entrenamiento sin una expansión completa del árbol, CAPR realiza un único evento de ramificación en el punto medio de la trayectoria de eliminación de ruido ($f_b = 0.5$).

• Caché (Caching): El modelo almacena en caché el prefijo (logits, tokens predichos y estado de la trayectoria) hasta el punto de ramificación.
• Ramificación (Forking): Genera un pequeño número de continuaciones hermanas ($G_b$, por ejemplo, 2) aplicando perturbaciones Gumbel a las probabilidades del vocabulario. Estos hermanos reutilizan el prefijo en caché, evitando pasadas hacia adelante redundantes.
• Poda (Pruning): En los límites de los bloques, las ramas se podan basadas en la calidad de su estado de trayectoria (estabilidad media). Solo las ramas más prometedoras continúan, evitando que el espacio de búsqueda se expanda hacia un árbol completo.

3. Block Critic (Crítico de Bloque)

Una cabeza de valor ligera (Block Critic) lee el estado de la trayectoria en cada límite de bloque.

• Redistribución de Recompensa: Después de que el rollout termina, la recompensa de resultado final se redistribuye entre los bloques basándose en el número de tokens revelados en cada bloque y las características del estado de la trayectoria.
• Cálculo de Ventaja: Utilizando TD/GAE (Estimación de Ventaja Generalizada), el crítico convierte la recomposa dispersa en una ventaja por bloque ($\hat{A}_b$).
• Actualización PPO: Estas ventajas por bloque reponderan la ventaja estándar a nivel de rollout utilizada en PPO. Esto permite al modelo aprender qué bloques específicos contribuyeron más al éxito, recuperando efectivamente la granularidad de la búsqueda en árbol sin el costo del árbol.

Objetivo de Entrenamiento: La pérdida total combina la pérdida PPO recortada (usando ventajas ponderadas por bloque), la divergencia KL (para permanecer cerca del modelo de referencia), un regularizador $\chi^2$ (para suprimir picos de razón de logaritmo poco comunes), pérdida de autodestilación (usando rollouts resueltos como objetivos supervisados) y una pérdida de la cabeza de valor.

Contribuciones Clave

1. RL Basado en la Traza de Eliminación de Ruido: El artículo propone una visión novedosa donde la confianza y la estabilidad a lo largo de la trayectoria de eliminación de ruido se convierten en ventajas por bloque, eliminando la necesidad de una expansión completa del árbol.
2. Algoritmo CAPR: Una implementación práctica que almacena en caché salidas tempranas, amortiza los rollouts hermanos, dirige los pasos de eliminación de ruido mediante estados de trayectoria y utiliza un crítico de bloque para generar pesos PPO granulares.
3. Rendimiento de Vanguardia (SOTA): CAPR logra nuevos resultados de vanguardia para dLLMs ajustados con RL en cuatro benchmarks de razonamiento (Sudoku 4x4, Countdown, GSM8K, Math500) tanto en backbones densos como de Mezcla de Expertos (MoE) con presupuestos de 256 y 512 tokens.

Resultados Experimentales

Los autores evaluaron CAPR contra líneas base que incluyen Diffu-GRPO, VRPO, wd1, SAPO, d2-stepMerge, TraceRL y el d-TreeRPO estructurado en árbol.

• Precisión: Con el presupuesto de 256 tokens, CAPR superó o empató con todas las líneas base en todas las tareas tanto para LLaDA-8B-Instruct como para LLaDA-MoE-7B-A1B-Instruct.
  • Sudoku: CAPR alcanzó un 93.0% (LLaDA-8B) y 84.8% (LLaDA-MoE), igualando o superando a la línea base de estructura de árbol más fuerte (d-TreeRPO: 92.9% y 78.3%).
  • Countdown: Se observaron ganancias significativas, con CAPR alcanzando 72.7% (8B) y 68.8% (MoE), superando a d-TreeRPO.
• Eficiencia: CAPR reduce significativamente el costo computacional en comparación con los rollouts de árbol.
  • Costo de Generación: CAPR requiere aproximadamente 0.75× el costo de los rollouts planos y 0.6× el costo de los rollouts de árbol (en configuraciones estándar).
  • Tiempo de Ejecución (Wall-Clock Time): En Sudoku, CAPR tomó 165 segundos por paso de optimización comparado con 596 segundos para d-TreeRPO, logrando una precisión comparable con menos de un tercio del cómputo por paso.
• Estudios de Ablación:
  • El Block Critic proporcionó la mayor ganancia intermedia, cambiando sustancialmente la curva de aprendizaje.
  • Branch & Prune añadió diversidad útil pero fue menos efectivo por sí solo.
  • Self-Distillation fue crucial para estabilizar el entrenamiento en grupos de varianza cero (rollouts resueltos) y reducir los picos de gradiente.

Significado y Reivindicaciones

El artículo afirma que CAPR ofrece un "punto medio" práctico para el RL de dLLM. Recupera con éxito la granularidad de la búsqueda en árbol (comparación local de alternativas) mientras evita la explosión computacional de la expansión completa del árbol.

• Eficiencia vs. Granularidad: CAPR demuestra que la propia traza de eliminación de ruido contiene información suficiente para asignar recompensas a nivel de bloque, eliminando la necesidad de costosos rollouts basados en árboles.
• Generalización: El método se transfiere eficazmente de backbones densos a dispersos (MoE) sin necesidad de reajuste específico de la tarea.
• Practicidad: Al reducir el cómputo por paso a aproximadamente 0.6× de los rollouts de árbol mientras iguala o supera su precisión, CAPR hace que el entrenamiento de RL de alta granularidad sea factible para modelos más grandes y contextos más largos.

Los autores señalan limitaciones, reconociendo que el método se evalúa actualmente en tareas de razonamiento verificables automáticamente y puede requerir adaptación para el seguimiento de instrucciones de código abierto o dominios con recompensas ruidosas y no locales. Además, el crítico de bloque se aprende en línea, lo que puede presentar desafíos de fiabilidad para cadenas de razonamiento muy largas.