Who Should Lead Decoding Now? Tracking Reliable Trajectories for Ensembling Masked Diffusion Language Models

Resumen Técnico: ¿Quién debe liderar la decodificación ahora? Seguimiento de trayectorias fiables para el ensamblaje de modelos de lenguaje de difusión con máscara

Planteamiento del problema
Los Modelos de Lenguaje de Difusión con Máscara (MDLM, por sus siglas en inglés) han surgido como un paradigma distinto para la generación de secuencias, utilizando el re-enmascaramiento iterativo para eliminar el ruido de las secuencias en paralelo en lugar de hacerlo de forma autorregresiva. A medida que la familia de los MDLM se diversifica con modelos que exhiben diferentes fortalezas y dinámicas de decodificación, surge un desafío crítico: cómo orquestar o fusionar eficazmente el conocimiento de MDLMs heterogéneos. Las estrategias de ensamblaje tradicionales para modelos autorregresivos —como el promedio de las probabilidades del siguiente token o la derivación hacia el modelo más confiable— no son directamente aplicables a los MDLM. En los MDLM, las secuencias se generan en un orden flexible y no de izquierda a derecha, lo que significa que diferentes modelos pueden operar sobre secuencias parcialmente eliminadas de ruido distintas en cada paso. Esta falta de una posición compartida de "siguiente token" dificulta la aplicación de las definiciones estándar de ensamblaje, dejando el ensamblaje de MDLM en gran medida inexplorado.

Metodología: TIE (Ensamblaje Iterativo basado en Trayectorias)
Para abordar esto, los autores proponen TIE, un marco de fusión de conocimiento diseñado específicamente para los MDLM. TIE opera sobre dos observaciones centrales derivadas del análisis de la dinámica de decodificación:

1. Estabilidad de las trayectorias correctas: Las muestras respondidas correctamente exhiben dinámicas más estables y confiables sobre los tokens relevantes para la respuesta (menor entropía, mayor probabilidad del top-1 y menos cambios de tokens) en comparación con las muestras incorrectas, incluso mientras los tokens aún están enmascarados.
2. Corrección mediante inyección: Los modelos que generan respuestas incorrectas a menudo pueden ser dirigidos hacia respuestas correctas si reciben estados intermedios parcialmente eliminados de ruido que resultan prometedores de modelos que siguen trayectorias más fiables.

Basándose en estos hallazgos, TIE implementa un ciclo recurrente de tres pasos durante el proceso de generación:

1. Generación de trayectoria: Cada modelo constituyente ($M$) decodifica de forma independiente durante un número fijo de pasos (el intervalo de ensamblaje, $n$), produciendo una trayectoria de decodificación parcial.
2. Evaluación de la trayectoria: El marco evalúa la fiabilidad de cada trayectoria puntuando la dinámica de confianza sobre las posiciones de los tokens de la respuesta (tokens que aparecen después de un sufijo específico como "Respuesta:"). La puntuación utiliza cuatro métricas:
   • Conteo de cambios de tokens negativos: Rastrea el historial de cambios del top-1 sobre las posiciones de la respuesta enmascaradas. Un conteo menor indica mayor estabilidad.
   • Probabilidad Top-1: La probabilidad del token más probable.
   • Entropía negativa: Mide la certeza de la distribución.
   • Margen de probabilidad: La brecha entre las probabilidades del top-1 y el top-2.
   • Puntuación entre modelos: Para las métricas basadas en logits, el marco procesa cada trayectoria a través de todos los modelos constituyentes y promedia las puntuaciones. Esto mitiga las diferencias de calibración entre modelos, asegurando que la trayectoria seleccionada sea consistentemente confiable en todo el ensamblaje en lugar de ser simplemente favorecida por su modelo de origen.
3. Relevo de trayectoria: Se selecciona la trayectoria con la puntuación agregada más alta. Esta trayectoria "ganadora" se transmite a todos los modelos constituyentes, reemplazando sus estados parciales actuales. Los modelos luego reanudan la decodificación independiente a partir de este estado intermedio compartido y más fiable.

Este ciclo se repite hasta que la generación termina. Finalmente, el sistema selecciona la respuesta final de los $M$ candidatos basándose en la trayectoria con la dinámica de tokens de respuesta más estable (menor conteo de cambios de tokens normalizado).

Contribuciones clave

• Análisis de decodificación novedoso: El artículo identifica que los tokens relevantes para la respuesta en las generaciones correctas de MDLM siguen trayectorias de confianza estables, proporcionando una señal para identificar rutas fiables antes de que se genere la respuesta completa.
• Marco de ensamblaje específico para MDLM: TIE introduce un mecanismo para fusionar el conocimiento mediante el relevo iterativo de trayectorias parciales, permitiendo que los modelos se corrijan mutuamente en sus desviaciones durante el proceso de eliminación de ruido, en lugar de esperar a los resultados finales.
• Mecanismo de puntuación entre modelos: Los autores demuestran la necesidad de evaluar las trayectorias a través de todos los modelos constituyentes para dar cuenta de las disparidades de calibración, evitando la selección sesgada hacia modelos excesivamente confiados.

Resultados experimentales
Los autores evaluaron TIE en ocho bancos de pruebas en cuatro dominios: razonamiento general (MMLU, ARC-Challenge, Winogrande), razonamiento matemático (GSM8K, MATH500), codificación (HumanEval, MBPP) y planificación (Countdown).

• Ganancias de rendimiento: TIE superó consistentemente a los MDLM individuales (LLaDA-1.5, Dream-7B, DreamCoder-7B, DiffuCoder-7B) y a las líneas base de ensamblaje post-generación. Por ejemplo, en GSM8K, el ensamblaje de LLaDA y Dream mediante TIE mejoró la precisión del 78.39% (Dream) al 83.47%.
• Métricas de puntuación: Entre las funciones de puntuación, el Conteo de cambios de tokens (específicamente la versión normalizada) resultó ser el más robusto, logrando los mejores resultados en cuatro de los ocho bancos de pruebas.
• Sinergia de modelos: El marco es más efectivo cuando los modelos constituyentes tienen capacidades individuales comparables y fuertes. Cuando las brechas de rendimiento son grandes (por ejemplo, >15%), el ensamblaje se vuelve menos efectivo ya que los modelos más débiles introducen señales ruidosas.
• Liderazgo dinámico: El análisis revela que el modelo que produce la trayectoria con la puntuación más alta cambia frecuentemente a lo largo del proceso de decodificación (tasas de cambio de modelo del 12 al 44% dependiendo de la tarea). Esto confirma que diferentes modelos aportan fortalezas complementarias en diferentes etapas de la generación.
• Compatibilidad: TIE sigue siendo efectivo cuando se combina con técnicas estándar de aceleración de MDLM, como el umbralización (thresholding) y el desmascaramiento top-$k$.

Significancia y afirmaciones
El artículo afirma que TIE ofrece un enfoque práctico para el problema inexplorado del ensamblaje de MDLM sin requerir entrenamiento adicional. Al aprovechar la dinámica de confianza sobre los tokens relevantes para la respuesta, TIE permite que modelos heterogéneos refinen colaborativamente las respuestas, permitiendo que los modelos más débiles se recuperen de trayectorias subóptimas. Los autores posicionan este trabajo como un paso significativo hacia una orquestación más efectiva de los modelos de lenguaje basados en difusión, sugiriendo que la inferencia colaborativa guiada por la estabilidad de la trayectoria puede mejorar significamente las capacidades de razonamiento en diversas tareas. El artículo señala modestamente las limitaciones, específicamente que la efectividad disminuye con grandes brechas de rendimiento entre modelos y que el estudio actual está limitado a ensamblajes de dos modelos.