Resumen Técnico: LESS: Muestreo de Estabilidad Mutua para Modelos de Lenguaje de Difusión

Planteamiento del Problema

Los Modelos de Lenguaje de Difusión (dLLMs) ofrecen una alternativa prometedora a la decodificación autorregresiva (AR) al permitir actualizaciones de tokens en paralelo y condicionamiento bidireccional mediante el desruido iterativo de secuencias enmascaradas. Sin embargo, su eficiencia práctica está actualmente limitada por procedimientos de muestreo que dependen de esquemas de desruido inverso fijos. Los muestreadores estándar ejecutan un número predeterminado de pasos, desperdiciando a menudo computación en posiciones donde las predicciones ya se han estabilizado, mientras que corren el riesgo de comprometerse prematuramente con tokens inestables. Una vez que un token es desmascarado en un proceso de difusión, típicamente se trata como fijo; si se compromete un token incorrecto prematuramente debido a un esquema rígido, los pasos inversos subsiguientes no pueden revisarlo, lo que conduce a una degradación de la calidad. Los muestreadores adaptativos existentes que no requieren entrenamiento abordan esto parcialmente mediante el uso de criterios basados en confianza o margen, pero a menudo no logran capturar la trayectoria predictiva completa y la estabilidad de la distribución de un token.

Metodología

Los autores proponen LESS (Less is More), un muestreador adaptativo, sin entrenamiento, agnóstico al modelo, que enmarca el compromiso del token como un problema de parada en línea por posición. LESS implementa el muestreo de estabilidad mutua, donde una posición enmascarada es elegible para ser desmascarada solo cuando tres señales de estabilidad específicas indican conjuntamente que la predicción se ha asentado.

La Regla de Estabilidad Conjunta

Para cada posición enmascarada $i$ en el paso inverso $t$, LESS evalúa tres señales derivadas de la distribución predictiva del modelo $p_{t,i}$:

1. Confianza Predictiva del Top-1: La probabilidad asignada al token actual top-1 ($w_{t,i}$). Una posición debe exceder un umbral de confianza $c$.
2. Persistencia del Token Top-1: El token actual top-1 debe permanecer sin cambios durante una ventana reciente de $P$ pasos inversos. Esto evita comprometerse con un token que acaba de cambiar, incluso si su confianza actual es alta.
3. Divergencia de Jensen–Shannon (JSD) Inter-paso Top-K: Esta es la contribución novedosa de LESS. Mide la estabilidad de la distribución predictiva misma entre pasos inversos consecutivos.
   • En lugar de utilizar la divergencia de Kullback-Leibler (KL) asimétrica, LESS utiliza la JSD simétrica y acotada.
   • Para asegurar la eficiencia computacional y centrarse en los candidatos relevantes, la JSD se calcula sobre un vocabulario simplificado que consiste en la unión de los tokens top-$K$ de los pasos actual y anterior, colapsando la masa restante en un símbolo residual.
   • La JSD debe caer por debajo de un umbral $d$.

Una posición es aceptada para el desmascarado si y solo si:
$$\text{accept}_{t,i} = (\text{conf}_{t,i} \ge c) \land \text{persist}_{t,i} \land (\text{JSD}_{t,i} \le d)$$

Política de Desmascarado

LESS emplea una política de prioridad de frontera (frontier-first policy) para modelos de difusión de secuencia completa. Intenta desmascarar la posición enmascarada más a la izquierda si es elegible. Si no lo es, escanea una ventana de mirada hacia adelante (look-ahead window) para encontrar la primera posición elegible. Si no hay posiciones elegibles, el paso se omite. Para garantizar la terminación, se activa un mecanismo de respaldo de desmascarado forzado tras $\Pi$ pasos omitidos consecutivos, aunque este respaldo no se cuenta como un compromiso "estable".

Contribuciones Clave

1. Formulación del Compromiso del Token: El artículo enmarca el desmascarado adaptativo en dLLMs como un problema de parada en línea inducido por la irreversibilidad del desmascarado.
2. JSD Inter-paso Top-K: Los autores introducen la divergencia de Jensen–Shannon inter-paso top-$K$ como una señal de estabilidad simétrica y acotada. Esto aborda las limitaciones del trabajo previo donde la confianza por sí sola no puede detectar el movimiento distributivo continuo, y la divergencia KL es direccional e ilimitada.
3. Muestreador LESS: Un muestreador sin entrenamiento que combina confianza, persistencia del token y JSD inter-paso en una regla de estabilidad conjunta. No requiere el reentrenamiento del modelo y es compatible con los checkpoints existentes de dLLM.
4. Validación Empírica: Evaluación exhaustiva a través de tres familias de dLLM (Dream-7B, LLaDA-8B, LLaDA-1.5-8B) y siete benchmarks de matemáticas, código y conocimiento general.

Resultados

La evaluación demuestra que LESS mejora significativamente la relación precisión-eficiencia en comparación tanto con la decodificación de presupuesto fijo como con otros fuertes muestreadores adaptativos sin entrenamiento (Prophet y KLASS).

• Eficiencia: A través de tres familias de modelos, LESS utiliza un 72.1% menos de pasos inversos en promedio en comparación con la decodificación de presupuesto fijo (Base).
• Precisión: LESS logra una ganancia de precisión promedio de +0.91 puntos sobre la línea base de presupuesto fijo mientras utiliza menos pasos. También supera a Prophet y KLASS en la precisión media en todas las familias de modelos.
• Rendimiento de Tareas: Las ganancias son más pronunciadas en tareas de generación de matemáticas y código. Por ejemplo, en Dream-7B, LESS logra una precisión media del 65.18% (frente al 63.30% de Base) con solo 64.3 pasos (frente a 256 de Base).
• Ahorro de Computación: Dado que cada paso inverso requiere una pasada hacia adelante (forward pass) del Transformer, la reducción de pasos se traduce directamente en una menor computación de inferencia. En Dream-7B GSM8K, LESS reduce la latencia de tiempo de ejecución de 19.45s a 5.16s y reduce el número de evaluaciones forward en aproximadamente un 75% en presupuestos de pasos altos.
• Estudios de Ablación:
  • Complementariedad de Señales: Ninguna señal individual (confianza, persistencia o JSD) funciona tan bien como la regla conjunta. La confianza es el motor principal, pero la persistencia y la JSD proporcionan salvaguardas esenciales contra el compromiso prematuro y la deriva distributiva.
  • Robustez: El método funciona de manera consistente a través de diferentes modelos y regímenes de decodificación (secuencia completa vs. bloque semiautoregresivo) utilizando una única configuración global $(c=0.75, d=0.040)$.

Significancia y Reivindicaciones

El artículo afirma que LESS proporciona un mecanismo práctico y sin entrenamiento para acelerar la decodificación de modelos de lenguaje de difusión, comprometiendo los tokens solo cuando sus predicciones son confiables, temporalmente persistentes y distributivamente estables.

Los autores enfatizan que la estabilidad es más que solo confianza. Una predicción de alta confianza aún puede ser poco fiable si el token top-1 ha cambiado recientemente o si la masa de probabilidad aún se está desplazando entre las alternativas. Al tratar la estabilidad como una propiedad conjunta de la prediccción actual y su trayectoria reciente, LESS evita las trampas del compromiso prematuro.

El trabajo posiciona a LESS como ortogonal y complementario a otros métodos de aceleración:

• Complementa los métodos de tiempo de entrenamiento (ej. destilación) que acortan el proceso inverso.
• Complementa los métodos de sistemas (ej. caching, decodificación especulativa) que reducen la latencia por paso.
• Actúa como una política de tiempo de inferencia que se puede combinar con estos métodos sin modificar el modelo subyacente.

El artículo concluye que la estabilidad mutua es un criterio viable para un compromiso temprano de tokens robusto, permitiendo que los dLLM alcancen una mayor eficiencia sin sacrificar, y a menudo mejorando, la calidad de la generación. Los autores reconocen que el método expone una relación de compromiso calidad-eficiencia controlada por los umbrales de estabilidad, que puede requerir ajuste para modelos o tareas específicas, aunque una configuración global fija resultó efectiva en su diversa suite de evaluación.