Una política de apenas 856K parámetros acelera el LLM de difusión LLaDA-8B

Los modelos de lenguaje de difusión prometen lo que los autorregresivos clásicos no pueden: generar muchos tokens en una sola pasada. En la práctica, todo depende de una pregunta: cuáles de los tokens predichos se pueden fijar ahora sin riesgo y cuáles conviene dejar bajo la máscara hasta el siguiente paso. Un nuevo estudio controlado ha descompuesto esa pregunta en sus partes y ha acabado entrenando una política de solo 856 mil parámetros. En LLaDA-8B-Instruct, esa política logró la mejor frontera entre precisión y número de evaluaciones del modelo (NFE) y el mayor rendimiento (throughput) en GSM8K.

La política como destilación del decodificado

Las políticas de unmasking entrenables ya existían, y mejoran de verdad el equilibrio entre velocidad y calidad. Pero los trabajos previos diferían en varios ejes a la vez: con qué se etiqueta el conjunto de entrenamiento, qué entradas ve la política y cuál es su tamaño. Como resultado, comparar los enfoques entre sí era casi imposible, y no estaba claro cuál de esas decisiones pesaba realmente.

Los autores proponen un marco unificado. En él, una política entrenada con supervisión destila un procedimiento de decodificado a través de un «profesor», es decir, del etiquetado. A partir de ahí, estudiaron por separado cuatro cuestiones:

• qué considera el profesor un token que se puede fijar;
• en qué estados de la generación imita la política al profesor;
• qué observa la política como entrada;
• cuánto hay que reducir el NFE para compensar su propia sobrecarga.

Dónde falla el etiquetado ingenuo

La forma más obvia de etiquetar es esta: un token se considera listo para fijarse si coincide con el que acabará en la respuesta final. El estudio mostró que, con un NFE bajo, ese criterio es demasiado laxo. La política aprende a aceptar tokens que coinciden por casualidad con el resultado final, pero que en realidad desvían la generación del camino de referencia. Cuando hay pocas evaluaciones, el coste de esos errores crece.

En cuanto a las entradas, el panorama tampoco es claro. Los estados ocultos del modelo, por sí solos, aportan una señal débil para decidir si fijar un token. En cambio, funcionan como complemento de la confianza (confidence) de la predicción: juntas, las dos fuentes de información aportan más que cada una por separado.

Por otro lado, los autores calcularon el coste de la propia política. Cualquier capa añadida sobre el modelo suma cómputo, y la ganancia de usar menos pasos puede perderse en su mantenimiento. Según sus estimaciones, basta una reducción del NFE de aproximadamente un 10% para compensar la sobrecarga de políticas de hasta 7,5 millones de parámetros en todas las configuraciones evaluadas.

Trayectoria en lugar de coincidencia con el resultado final

De estas observaciones surgió el esquema final. La política se entrena con objetivos de trajectory-contraction: una fijación solo se acepta cuando permanece en la trayectoria de decodificado de referencia. Como entrada, la política recibe la confianza de las predicciones y los estados ocultos.

Principales resultados:

• tamaño de la política: 856K parámetros, menos del 0,02% de los 8 mil millones de parámetros de LLaDA-8B-Instruct;
• GSM8K: mejor frontera entre precisión y NFE y mayor rendimiento entre los enfoques comparados;
• MATH-500 y HumanEval: resultados competitivos con objetivos de aceleración moderados;
• Dream-7B-Instruct: la política se transfiere en gran medida a otro modelo sin un rediseño completo.

Una salvedad que se desprende de los propios resultados: en MATH-500 y HumanEval, la ventaja se declara solo con objetivos moderados, es decir, con una reducción agresiva de pasos no está garantizado el liderazgo.

Qué llevarse a la práctica

Para los equipos que ejecutan LLM de difusión en producción, la conclusión principal es de ingeniería. Acelerar el decodificado no requiere necesariamente una capa pesada: una política de menos de un millón de parámetros se amortiza si reduce el NFE al menos en torno a un 10%. Es una referencia útil para el presupuesto: si la ganancia en pasos es menor, la sobrecarga de la política se la comerá por completo.

La segunda recomendación atañe al etiquetado. Si se destila una estrategia de decodificado, no conviene dar por suficiente el criterio «coincide con la respuesta final», sobre todo con pocas evaluaciones. Anclarse a la trayectoria de referencia da señales de entrenamiento más estrictas. Y los estados ocultos conviene añadirlos a la confianza, no usarlos como única entrada.

Por último, la transferencia a Dream-7B-Instruct sugiere que estas políticas pueden hacerse relativamente universales dentro de la familia de modelos de difusión enmascarados. Para una industria en la que los modelos de lenguaje de difusión pugnan por un lugar junto a los autorregresivos en tareas con requisitos estrictos de latencia, es un argumento a favor de capas añadidas baratas y transferibles en lugar de reentrenar los propios modelos.