Resumen Técnico: Representaciones de Relé Aprendidas para Modelos de Difusión Discreta con Pensamiento Prospectivo

1. Declaración del Problema: El "Reinicio Duro" en Modelos de Difusión Enmascarada

Los Modelos de Difusión Enmascarada (MDM) generan secuencias discretas mediante desruido iterativo. Comenzando desde una secuencia completamente enmascarada, el modelo desenmascara iterativamente una fracción de posiciones. Sin embargo, existe una ineficiencia crítica en la inferencia estándar de MDM: después de cada paso hacia adelante, el modelo descarta los ricos cálculos internos continuos (estados ocultos) generados para las posiciones enmascaradas. El siguiente paso comienza únicamente desde los tokens discretos recién comprometidos.

Los autores denominan a esto el problema del reinicio duro. Aunque los MDM realizan múltiples pasos hacia adelante por generación, la incapacidad de reutilizar el cálculo continuo intermedio impide que el modelo amortice eficazmente el razonamiento entre pasos. Esto contrasta con el cálculo recurrente, donde se ha demostrado que desenrollar un modelo de parámetros fijos a través de los pasos expande la clase de funciones que el modelo puede aproximar, ayudando al rendimiento en tareas de razonamiento complejo. La pregunta central abordada es cómo permitir que los MDM soporten un cálculo recurrente que transporte información más rica a través de los pasos secuenciales de desenmascaramiento sin descartar los estados latentes intermedios.

2. Metodología: Representaciones de Relé Aprendidas (Relay)

La solución propuesta, Relay, introduce un estado latente continuo y diferenciable que se transmite hacia adelante a través de los pasos de decodificación, permitiendo que el modelo sea "con pensamiento prospectivo".

Mecanismo Central

• Trajectoria de Estado Aumentada: En cada paso de inferencia $k$, el modelo mantiene un estado aumentado $s_k = (x_{t_k}, h_k)$, donde $x_{t_k}$ es la secuencia parcialmente enmascarada actual y $h_k$ es un estado de relé aprendido (vector latente continuo).
• Arquitectura: El modelo consiste en una base $f_\theta$, incrustaciones de tokens $Emb_\theta$, una cabeza de desenmascaramiento $UnEmb_\theta$ y un módulo de relé específico $R_\theta$.
  • El siguiente estado oculto se calcula como: $h_{k+1} = f_\theta(Emb_\theta(x_{t_k}) + R_\theta(h_k))$.
  • El estado de relé $h_k$ se proyecta y se suma a los tokens incrustados antes de entrar a la base.
  • Los logits de salida $\ell_k$ se derivan de $h_{k+1}$ mediante $UnEmb_\theta$.
• Entrenamiento mediante BPTT Truncado: A diferencia de los MDM estándar entrenados con pasos independientes, Relay se entrena utilizando Retropropagación a través del Tiempo Truncada (BPTT).
  • El modelo realiza una simulación de $K$ pasos (por ejemplo, $K=2$) comenzando desde una secuencia completamente enmascarada.
  • Durante la simulación, las posiciones se desenmascaran basándose en una política (por ejemplo, umbrales de confianza), y los tokens comprometidos se "fuerzan con el maestro" a valores de verdad para evitar la acumulación de errores durante la simulación.
  • La pérdida es la suma de las pérdidas de entropía cruzada sobre los $K$ pasos.
  • Los gradientes se retropropagan a través de la ruta del relé ($h_k \to h_{k+1}$), permitiendo que el modelo aprenda no solo a predecir el paso actual, sino a codificar información en $h_k$ que sea máximamente útil para los pasos futuros.

Inferencia

En el momento de la inferencia, el proceso permanece en gran medida sin cambios en términos del cronograma de desenmascaramiento o la política de muestreo. La única adición es que el modelo transporta el estado de relé $h_k$ hacia adelante junto con los tokens comprometidos. El estado de relé sirve como un canal de memoria diferenciable para la información que aún no se ha comprometido a tokens discretos.

3. Contribuciones Clave

1. Marco Relay: Un método general para incorporar cálculo recurrente en MDMs entrenando al modelo para transmitir un relé latente aprendido hacia adelante a través de los pasos de decodificación mediante BPTT truncado. Es agnóstico a la arquitectura y compatible con técnicas de aceleración de MDM existentes como la difusión por bloques y la caché KV.
2. Escalabilidad y Adaptación: Validación de Relay a la escala de Modelos de Lenguaje Grandes (LLM) mediante la adaptación de parámetros completos de Fast-dLLM v2 (1.5B). El método adapta exitosamente un Modelo de Lenguaje Discreto (DLM) de última generación con FLOPs de entrenamiento adicionales mínimos.
3. Mejoras en Rendimiento-Latencia: Demostración empírica de que Relay avanza la frontera de Pareto de rendimiento-latencia, superando al ajuste fino supervisado estándar (SFT) en tareas de codificación mientras reduce significativamente la latencia de inferencia.
4. Análisis del Espacio de Diseño: Estudios de ablación extensos en una tarea de planificación basada en Sudoku desafiante que aíslan las contribuciones del procedimiento de simulación, el canal de relé y el objetivo de entrenamiento BPTT.

4. Resultados Experimentales

Tarea de Planificación de Sudoku

Utilizando el conjunto de datos Sudoku-Extreme, los autores compararon cuatro objetivos de entrenamiento:

• MLM: Difusión enmascarada uniforme estándar (paso único, sin relé).
• Rollout: Simulación de múltiples pasos con fuerza de maestro pero sin canal de relé ($R_\theta \equiv 0$).
• Relay (sg): Incluye la ruta de relé pero utiliza gradientes de parada (sin BPTT a través del relé).
• Relay: Método completo con $K=2$ BPTT a través del relé.

Hallazgos:

• Frontera Precisión-NFE: Relay logró el mejor equilibrio entre precisión y Número de Evaluaciones hacia Adelante (NFE).
• Impacto de los Componentes: Pasar de MLM a Rollout mejoró el rendimiento. Agregar el canal de relé (Relay sg) proporcionó un salto significativo. Finalmente, habilitar BPTT (Relay) produjo los mejores resultados, demostrando que el modelo aprende a mantener el tablero parcial autoconsistente bajo un desenmascaramiento más agresivo.
• Eficiencia: Con un umbral de confianza fijo, Relay comprometió más celdas por paso hacia adelante mientras mantenía la legalidad (sin violaciones de fila/columna/caja), alcanzando la misma precisión en menos pasos hacia adelante totales en comparación con las líneas base.

Adaptación de Modelo Preentrenado (Fast-dLLM v2)

Los autores adaptaron Fast-dLLM v2 (1.5B) utilizando Relay en una mezcla de conjuntos de datos de código y matemáticas.

• Rendimiento en Codificación: En los benchmarks HumanEval y MBPP, Relay superó al SFT estándar.
  • HumanEval: Relay logró una precisión del 42.1% (vs. 38.4% para SFT) con 88.3 NFE (vs. 130.7 para SFT).
  • MBPP: Relay logró una precisión del 46.6% (vs. 43.9% para SFT) con 78.8 NFE (vs. 84.8 para SFT).
• Latencia: Relay redujo la latencia de inferencia hasta en un 32% en comparación con el SFT estándar mientras mejoraba la precisión.
• Sobrecarga de Memoria: El perfilado en una GPU A100 mostró que, aunque Relay aumenta la memoria de activación en vivo durante el segundo paso hacia adelante, el uso máximo de memoria permanece casi idéntico al SFT estándar (aprox. 20.1 GiB vs. 21.2 GiB) porque el cuello de botella es la pasada hacia atrás de la proyección del vocabulario, no el almacenamiento del estado de relé.

5. Significado y Afirmaciones

El artículo afirma que Relay cambia fundamentalmente la naturaleza de la generación no autoregresiva iterativa al permitir que los MDM acumulen cálculo continuo intermedio.

• Pensamiento Prospectivo: El método permite a los modelos optimizar explícitamente las representaciones intermedias para decisiones futuras de desenmascaramiento, en lugar de tratar cada paso como un problema de predicción nuevo.
• Avance de la Frontera de Pareto: Al reducir el número de pasos hacia adelante necesarios para alcanzar una precisión objetivo, Relay avanza la frontera de rendimiento-latencia, haciendo que los modelos de difusión discreta sean más eficientes para tareas complejas de razonamiento y generación.
• Compatibilidad: El enfoque se presenta como una mejora directa que funciona con DLMs de última generación y técnicas de aceleración existentes (por ejemplo, difusión por bloques, caché KV) sin requerir cambios arquitectónicos.

Los autores concluyen que Relay representa un paso significativo hacia una generación no codiciosa y con pensamiento prospectivo, ofreciendo un sustrato continuo para el razonamiento latente que puede explorarse más a fondo para la interpretabilidad y la aplicación en otras modalidades (por ejemplo, imágenes, moléculas).