Resumen Técnico: Mejora de la Decodificación de Modelos de Lenguaje de Difusión mediante la Búsqueda Conjunta en el Espacio de Orden y de Tokens

Planteamiento del Problema

Los Modelos de Lenguaje de Difusión (DLM), particularmente los Modelos de Difusión con Máscaras (MDM), ofrecen una ventaja distintiva sobre los modelos autorregresivos (AR): la generación independiente del orden. A diferencia de los modelos AR que generan tokens estrictamente de izquierda a derecha, los MDM pueden denotar múltiples posiciones en paralelo, lo que permite órdenes de generación adaptativos. Sin embargo, los métodos de decodificación actuales no logran explotar plenamente esta flexibilidad.

Las estrategias de decodificación existentes suelen comprometerse con una única trayectoria por ejecución, regidas por reglas heurísticas de remascarado:

• Remascarado de Baja Confianza: Fija de forma codiciosa los tokens de alta confianza de forma temprana. Aunque esto suele producir una fuerte precisión de muestra única (pass@1), restringe la búsqueda a un conjunto estrecho de trayectorias, lo que conduce a un crecimiento lento en la cobertura de múltiples muestras (pass@k) y al potencial atrapamiento en óptimos locales incorrectos.
• Remascarado Aleatorio: Muestrea patrones de remascarado de forma uniforme para maximizar la diversidad de las trayectorias. Si bien esto mejora el pass@k al explorar un espacio de soluciones más amplio, sufre de un débil pass@1 porque carece de un mecanismo basado en principios para seleccionar la trayectoria correcta entre los candidatos competidores.

Esto crea un intercambio entre exploración y explotación: los métodos actuales o bien sobre-explotan la confianza local (perdiendo soluciones válidas) o exploran ciegamente (fallando en la selección de la mejor solución). El artículo argumenta que el potencial de los DLM para explorar diversos órdenes de generación y valores de tokens permanece sin realizar porque la decodificación estándar no realiza una búsqueda explícita en este espacio conjunto.

Metodología: Búsqueda de Orden-Token

Para abordar esto, los autores proponen la Búsqueda de Orden-Token, un algoritmo de decodificación que realiza una búsqueda de haz (beam search) en el espacio conjunto de órdenes de generación y valores de tokens. El algoritmo alterna entre expandir candidatas trayectorias y podar aquellas que no son prometedoras.

1. Proceso de Búsqueda (Expansión)

La Búsqueda de Orden-Token mantiene un haz de $K$ candidatas parcialmente denotadas. A diferencia del muestreo estándar que sigue una secuencia fija de decisiones, el algoritmo se ramifica explícitamente sobre dos fuentes de incertidumbre en intervalos especificados por el usuario:

• Orden de Generación: Qué posiciones (tokens) desmascarar a continuación.
• Valores de Token: Qué tokens asignar a esas posiciones.

Para gestionar la complejidad computacional, la búsqueda utiliza difusión por bloques. En lugar de ramificarse en cada paso de denotación individual (lo que incurriría en un costo prohibitivo), el algoritmo expande las candidatas en los límites de los bloques de tokens contiguos. Dentro de cada bloque, ocurre una denotación independiente. Esto reduce los puntos de decisión desde la longitud total de la secuencia al número de bloques, permitiendo una diversidad sustancial tanto en el orden como en la asignación de tokens, manteniendo los costos de inferencia manejables.

2. Proceso de Poda (Estimación de Verosimilitud)

La innovación central reside en el estimador de verosimilitud utilizado para puntuar y podar las trayectorias parciales. Un objetivo estándar de MDM (verosimilitud de reconstruir la secuencia completa limpia desde un estado ruidoso) es inestable para estados intermedios con muchas máscaras restantes, ya que obliga al modelo a resolver problemas de relleno difíciles para los cuales no fue optimizado durante el entrenamiento.

Los autores introducen una función de puntuación incremental a nivel de acción:

• En lugar de puntuar el estado completo, el algoritmo puntúa las acciones de denotación específicas tomadas para transicionar de un estado más corrupto $x_t$ a uno menos corrupto $x_s$.
• La puntuación $s(x_t; x_s)$ mide con qué confianza el modelo puede repredicción los nuevos tokens revelados en $x_s$, condicionado a la predicción de secuencia completa del modelo $x_0$ (que sirve como contexto).
• Este enfoque evalúa si las decisiones pasadas de una trayectoria (elección de orden y token) están respaldadas por el modelo, proporcionando una señal de menor varianza y más discriminativa para la poda.

La puntuación acumulada de un candidato es la suma de estas puntuaciones incrementales a lo largo de los intervalos de búsqueda. El algoritmo retiene los $K$ mejores candidatos después de cada expansión, filtrando eficazmente las rutas incorrectas o subóptimas antes de continuar la decodificación.

Contribuciones Clave

1. Identificación del Intercambio: El artículo demuestra empíricamente que los métodos de decodificación de MDM estándar ocupan los extremos opuestos del espectro de exploración–explotación. El remascarado de baja confianza sobre-explota (alto pass@1, bajo crecimiento de pass@k), mientras que el remascarado aleatorio explora ciegamente (alto pass@k, bajo pass@1).
2. Algoritmo de Búsqueda de Orden-Token: La propuesta de un nuevo algoritmo de decodificación que busca en el espacio conjunto de órdenes de generación y valores de tokens. Utiliza un estimador de verosimilitud incremental estable para podar trayectorias parciales, convirtiendo la diversidad de trayectorias en una mejora de la precisión de muestra única.
3. Validación Empírica: Se demuestra que el método supera consistentemente a los modelos de referencia en tareas de razonamiento matemático y programación, igualando o superando las ganancias obtenidas típicamente mediante métodos costosos de post-entrenamiento como diffu-GRPO.

Resultados Experimentales

Los autores evaluaron la Búsqueda de Orden-Token (OTS) en GSM8K, MATH500, Countdown y HumanEval utilizando LLaDA-8B-Instruct y LLaDA-1.5 (una variante de post-entrenamiento por RL).

• Ganancias de Rendimiento: OTS logró consistentemente la mayor precisión promedio en los benchmarks.
  • En LLaDA-8B, OTS mejoró el pass@1 en un +3.1% en GSM8K, +3.8% en MATH500, +7.9% en Countdown y +6.8% en HumanEval comparado con la línea base de remascarado de baja confianza.
  • En LLaDA-1.5, OTS alcanzó mejoras relativas similares, llegando al 40.4% de precisión promedio frente al 37.2% del mejor baseline (Low-conf + MV).
• Comparación con Post-Entrenamiento: OTS igualó o superó el rendimiento de los modelos sometidos a post-entrenamiento por aprendizaje por refuerzo (diffu-GRPO), lo que sugiere que la búsqueda en tiempo de inferencia puede ser tan efectiva como el reentrenamiento del modelo para mejorar las capacidades de razonamiento.
• Estudios de Ablación:
  • Estimador de Verosimilitud: La función de puntuación incremental específica (Ec. 2) fue crítica. Las variantes que utilizaban la verosimilitud conjunta de todos los bloques revelados o verosimilitudes de estilo AR funcionaron significativamente peor, confirmando que el estimador diseñado a medida es necesario para una poda estable en MDMs.
  • Espacio de Búsqueda: OTS superó a "Búsqueda de Orden" (buscando solo el orden) y "Búsqueda de Token" (buscando solo tokens) que utilizaban verosimilitudes de estilo AR, destacando la necesidad de la búsqueda conjunta y el desajuste de las heurísticas AR con los MDM.
• Eficiencia: OTS escala de manera más efectiva con el cómputo en tiempo de prueba (NFE) que el voto de mayoría. Con presupuestos de cómputo igualados, OTS logró una mayor precisión (por ejemplo, 29.3% frente a 19.9% en Countdown) al aprovechar la búsqueda estructurada en lugar del muestreo no estructurado.

Significado y Reivindicaciones

El artículo sostiene que la búsqueda conjunta de orden-token es un componente clave para avanzar en la decodificación de DLM. Al explorar explícitamente el espacio de órdenes de generación y valores de tokens y podar basándose en un estimador de verosimilitud estable, el método convierte el potencial de múltiples trayectorias inherente de los modelos de difusión en un rendimiento fiable de muestra única.

Los autores posicionan este trabajo como una solución al cuello de botella de "exploración–explotación" en la decodificación actual de MDM. Afirman que su enfoque permite a los DLM escapar de los óptimos locales que atrapan a los métodos codiciosos y seleccionar las soluciones correctas que el muestreo aleatorio pierde, todo ello sin requerir entrenamiento adicional del modelo. El trabajo establece que las estrategias de búsqueda en tiempo de inferencia, específicamente adaptadas a la naturaleza paralela y agnóstica al orden de los modelos de difusión, pueden generar mejoras sistemáticas en tareas de razonamiento y programación.