El bucle de capas recorta hasta 3,3 veces el entrenamiento de las LLM de difusión

Los modelos de lenguaje de difusión con ruido uniforme (UDLM) llevan tiempo considerándose una alternativa interesante a la autorregresión, aunque cara. Un nuevo trabajo muestra cómo reducir ese coste de forma notable. Si dentro de cada paso de eliminación de ruido se ejecutan de forma recursiva varias de las primeras capas del transformer, looped UDLM alcanza la validation NLL de una UDLM convencional entrenada con aproximadamente 0,5 billones de tokens, con hasta 2,8 veces menos coste en OpenWebText y hasta 3,3 veces menos en FineWeb-Edu. El objetivo de entrenamiento y el sampler no se modificaron.

Por qué es caro entrenar UDLM y por qué merece la pena

Los modelos de difusión enmascarados (MDLM) trabajan con un contexto limpio. Una parte de los tokens se sustituye por una máscara, el resto permanece correcto y el modelo reconstruye los huecos. Los modelos uniformes son más exigentes. Cualquier token de la secuencia puede sustituirse por un token aleatorio del vocabulario normal y corregirse en los pasos siguientes. Durante el entrenamiento, el modelo debe limpiar el ruido del texto apoyándose en un contexto que a su vez puede ser poco fiable.

La tarea resulta más difícil, así que las UDLM requieren más FLOPs que las MDLM. A cambio, según los autores, obtienen mejores resultados en tareas downstream relacionadas con el razonamiento. Poder revisar un token ya escrito da al modelo la oportunidad de corregir su propio error en lugar de arrastrarlo hasta el final de la generación. La cuestión es el precio, y es precisamente lo que este trabajo intenta reducir.

Bucle en lugar de nuevos parámetros

La idea del selective looping ya demostró su utilidad en MDLM con un presupuesto de FLOPs equivalente. Varias de las primeras capas del transformer se aplican de forma recursiva dentro de un único paso de eliminación de ruido. Los pesos de las capas son los mismos, pero la profundidad efectiva del cálculo aumenta. Los autores trasladaron el enfoque a las UDLM y no tocaron nada más: ni la función de pérdida ni el procedimiento de muestreo.

Principales cifras del trabajo:

• Entrenamiento: looped UDLM alcanza la validation NLL de la UDLM base entrenada con ~0,5 billones de tokens, con hasta 2,8x menos coste en OpenWebText.
• En FineWeb-Edu la ganancia es mayor: hasta 3,3x menos FLOPs.
• Generación: con el mismo presupuesto de NFE (número de evaluaciones de la red por muestra), la generation perplexity es hasta un 9,8% menor que la del modelo sin bucles.
• La mayor parte de la mejora se logra con un número moderado de bucles; aumentarlo más aporta bastante menos.
• Distribución del número de bucles en el entrenamiento: sesgarla hacia valores bajos reduce el número de bucles en inferencia un 12% respecto al muestreo uniforme, sin que caiga la precisión en tareas downstream.

Conviene subrayar que lo que cuenta es la combinación de métricas. La validation NLL refleja la calidad del entrenamiento, mientras que la generation perplexity con NFE fijo refleja la calidad de la propia generación. La mejora se observa en ambas, y la comprobación en tareas downstream confirma que no se limita a métricas favorables para el modelo.

Cuántos bucles usar y a quién le sirve

La conclusión práctica es muy terrenal. El rendimiento de añadir bucles no es lineal, así que no tiene sentido perseguir la máxima recursión. Un valor moderado aporta la mayor parte de la mejora, y a partir de ahí el coste de cada paso crece más rápido que el beneficio. En los modelos de difusión, donde la generación ya exige muchas iteraciones de eliminación de ruido, esto pesa especialmente: cada bucle adicional se multiplica por el número de pasos.

El segundo punto práctico tiene que ver con la distribución del número de bucles. El muestreo uniforme durante el entrenamiento parece la opción natural, pero usar valores bajos con más frecuencia hace que el modelo aguante mejor un modo de funcionamiento recortado. Como resultado, en inferencia hacen falta un 12% menos de bucles con la misma precisión. Para los equipos que construyen LLM de difusión, es una receta lista para aplicar. No exige cambiar nada en el sampler, de modo que puede incorporarse a un pipeline existente.

El principal interés del trabajo está en otro lado. Las UDLM razonaban mejor, pero perdían en coste de entrenamiento. Si el bucle recorta de forma estable los costes entre 2,8 y 3,3 veces, la brecha de precio con las MDLM se estrecha y el argumento a favor del ruido uniforme resulta más convincente. Queda abierta la cuestión de la escala más allá de OpenWebText y FineWeb-Edu. Pero la dirección está clara: la calidad de los modelos de lenguaje de difusión puede conseguirse con arquitectura, no únicamente con más tokens y más horas de GPU.