Resumen Técnico: Mapas de Flujo Discretos de Núcleo Latente para la Generación de Pocos Pasos

Planteamiento del Problema

Los modelos de difusión discreta y de ajuste de flujo (flow-matching) generan texto mediante el denoise de una secuencia a lo largo de muchos pasos. Para mantener la eficiencia computacional, estos modelos suelen factorizar la probabilidad de transición a través de las posiciones de los tokens, decidiendo cada token de forma independiente dada la secuencia actual. Si bien esta factorización hace que los pasos individuales sean económicos, crea un cuello de botella fundamental para la generación de pocos pasos.

Cuando una secuencia objetivo requiere correlaciones fuertes entre posiciones (por ejemplo, concordancia sujeto-verbo o paréntesis coincidentes), las actualizaciones independientes se comprometen con los tokens por separado. Esto suele resultar en combinaciones localmente consistentes pero globalmente inconsistentes (por ejemplo, "dog barks" frente a "dogs bark"). Corregir estos desajustes requiere muchos pasos adicionales de denoise, lo que anula las ganancias de eficiencia de la generación de pocos pasos.

Las soluciones existentes para acelerar la generación se dividen en dos categorías:

1. Modelos de verosimilitud (ej. MDLM, SEDD): Entrenan denoisers factorizados robustos desde cero, pero requieren muchos pasos en la inferencia para resolver las correlaciones.
2. Aceleradores de pocos pasos (ej. Di4C, SDTT, ReDi): Destilan o rectifican un modelo "maestro" entrenado hacia un estudiante de pocos pasos. Sin embargo, estos métodos heredan el techo de calidad de su maestro y no pueden superar el rendimiento del modelo al que imitan.

El desafío central es si un modelo puede expresar pasos correlacionados de forma nativa sin depender de un maestro, rompiendo así el supuesto de factorización desde su raíz manteniendo un bajo costo computacional.

Metodología: Mapas de Flujo Discretos de Núcleo Latente (LKF)

Los autores proponen los Mapos de Flujo Discretos de Núcleo Latente (LKF), un modelo de mapa de flujo entrenado desde cero que introduce correlación intra-paso mediante una mezcla de componentes factorizados vinculados por una única variable latente compartida.

La Arquitectura del Modelo

El núcleo de transición de tiempo finito $\Psi_{t \to s}(x_s | x_t)$ se define como una mezcla de $M$ componentes factorizados:
$$\Psi_{t \to s}(x_s | x_t) = \sum_{k=1}^{M} w_k(x_t, t, s) \prod_{i=1}^{L} P_i(x_s^i | x_t, k, t, s)$$
Donde:

• $k \in [M]$ es un índice de componente latente compartido muestreado una vez por secuencia.
• $w_k$ son pesos del enrutador que predicen la probabilidad del componente $k$.
• $P_i$ son distribuciones de tokens factorizadas para la posición $i$ condicionada al latente $k$.

Mecanismo Clave:

1. Latente Compartido: Al inicio de una trayectoria de denoise (o en cada paso), se extrae un latente $k$ de la función de enrutamiento $w(\cdot | x_t)$.
2. Independencia Condicional: Dado $k$, los tokens todavía se muestrean de forma independiente entre las posiciones. Sin embargo, debido a que $k$ es compartido en toda la secuencia, la marginalización sobre $k$ induce correlaciones entre las posiciones.
3. Eficiencia: El modelo utiliza una arquitectura de inyección de latente tardía. Un tronco de Transformer compartido procesa la entrada para crear una representación agnóstica al latente. Solo los últimos $L_k$ bloques se replican para cada uno de los $M$ componentes, condicionados a $k$. Esto mantiene bajo el costo computacional de la marginalización ($O(M \cdot L_k)$ en lugar de $O(M \cdot \text{profundidad})$).

Entrenamiento e Inferencia

• Entrenamiento: El modelo se entrena marginalizando el latente $k$ de forma exacta. El objetivo es el logaritmo de la verosimilitud de la mezcla:
  $$\log \Psi = \log \sum_{k} \exp(\log w_k + \ell_k)$$
  donde $\ell_k$ es la log-probabilidad asignada por el componente $k$ a los tokens objetivo. Esto obliga al enrutador a aprender una mezcla donde el latente captura las correlaciones de la secuencia. Un regularizador del enrutador asegura que los componentes de la mezcla sean utilizados (evitando el colapso a un solo componente) y que el latente transporte suficiente información.
• Inferencia (Best-of-M): Para aprovechar la heterogeneidad de la mezcla, los autores emplean una estrategia de decodificación Best-of-M. Se generan $M$ ejecuciones paralelas, cada una congelando un componente $k$ diferente. La salida final se selecciona basada en el propio ELBO negativo de la mezcla (NELBO) de la secuencia del modelo. Esta selección no requiere un modelo de recompensa externo ni un maestro.

Contribuciones Clave

1. Pasos Correlacionados Nativos: Introducción de LKF, un núcleo de mapa de flujo que expresa pasos correlacionados de forma nativa mediante una mezcla de $M$ componentes factorizados vinculados por un único latente compartido, eliminando la necesidad de un modelo maestro.
2. Análisis Teórico: Una descomposición de información exacta que muestra que la correlación total (TC) que puede portar un paso de rango-$M$ está acotada por $(L-1)H(k|xt)$. Los autores demuestran que el error de aproximación del modelo está acotado inferiormente por la correlación total no capturada de la distribución objetivo.
3. Validación Sintética: Experimentos en conjuntos de datos sintéticos (acuerdo oculto y paridad) demuestran que LKF captura información y mejora la verosimilitud de forma monótona con $M$ para correlaciones representables, mientras que falla correctamente en capturar correlaciones que requieren un rango de mezcla exponencial (paridad), alineándose con los resultados teóricos de dificultad.
4. Rendimiento Empírico: En los benchmarks LM1B y WikiText-103, LKF alcanza la perplejidad generativa (gen-PPL) de estado del arte. Con $M=8$, mejora la perplejidad generativa en 2.1× a 3.3× respecto a las bases de verosimilitud (MDLM, SEDD) y supera a los aceleradores destilados/rectificados (Di4C, ReDi) sin utilizar ningún maestro.

Resultados Experimentales

• Benchmarks: Evaluado en LM1B y WikiText-103.
• Rendimiento:
  • LM1B: El modelo LKF con $M=8$ logra una perplejidad generativa (gen-PPL) de 105 con 32 evaluaciones de función (NFE), comparado con 304 para MDLM y 117 para el mejor acelerador rectificado (ReDi).
  • WikiText-103: El modelo con $M=8$ alcanza una gen-PPL de 75, superando significativamente a MDLM (228) y SEDD (247).
• Diversidad: A diferencia de los aceleradores rectificados que a menudo sufren de colapso de entropía (produciendo texto de baja diversidad), LKF mantiene una alta entropía de muestreo mientras mejora la verosimilitud.
• Ablaciones:
  • La ganancia se atribuye a la estructura de la mezcla en sí misma, no solo al reordenamiento. Reemplazar los $M$ componentes con $M$ dibujos de ruido de un solo componente resulta en una caída significativa de rendimiento (ej. un 35% más de PPL con $M=8$).
  • La selección Best-of-M es crucial; una política de compromiso único (commit-M) produce resultados similares al baseline, confirmando que la mezcla proporciona la diversidad de candidatos que la selección explota.
  • Eficiencia del Enrutador: El enrutador actúa como un peso de mezcla más que como una señal de calidad; la selección es más efectiva cuando se basa en la propia verosimilitud del modelo (evidencia en ejecución) en lugar de la predicción inicial del enrutador.

Significado y Reivindicaciones

El artículo afirma que LKF representa un cambio de la aceleración "dependiente de un maestro" hacia la generación de pocos pasos "desde cero". Al modelar el núcleo de transición como una mezcla de bajo rango, LKF permite que un solo paso asigne masa a completaciones correlacionadas (ej. concordancia sujeto-verbo) con la misma complejidad de muestreo que un modelo factorizado.

Los autores enfatizan que:

• El método no hereda el techo de calidad de un maestro.
• La ganancia escala con el tamaño de la mezcla $M$.
• El enfoque está fundamentado teóricamente, con límites sobre la capacidad de correlación y el error de aproximación.
• Logra la mejor perplejidad generativa de pocos pasos entre los métodos de referencia en benchmarks estándar de modelado de lenguaje, ofreciendo una vía viable para la generación discreta eficiente y de alta calidad sin la sobrecarga de la destilación o la rectificación.