Resumen Técnico: No Reentrenar—Alinear: Adaptación de Modelos de Lenguaje Autoregresivos a Modelos de Lenguaje de Difusión mediante Alineación de Representaciones

Enunciado del Problema

Los Modelos de Lenguaje de Difusión (DLM) ofrecen ventajas distintas sobre los modelos autoregresivos (AR) estándar, particularmente en la generación no secuencial, la edición bidireccional y el refinamiento iterativo. Sin embargo, escalar los DLM es computacionalmente costoso. Teóricamente, los DLM deben aprender $L!$ caminos de generación en comparación con el único camino de izquierda a derecha de los modelos AR, lo que a menudo requiere significativamente más capacidad de cómputo.

Los enfoques recientes intentan mitigar este costo convirtiendo checkpoints preentrenados de AR en DLM. Si bien estos métodos logran inicializar DLM a partir de pesos de AR, dependen principalmente del entrenamiento continuo de desruido con modificaciones en los objetivos, patrones de atención o convenciones de muestreo. Crucialmente, las recetas existentes no restringen explícitamente al DLM convertido para preservar la geometría de representación interna aprendida durante el preentrenamiento original de AR. Esto plantea una pregunta central: ¿Es necesario tratar a los AR y a los DLM como paradigmas disjuntos que requieren reaprender representaciones del lenguaje, o puede preservarse explícitamente la estructura semántica aprendida por el preentrenamiento de AR durante la conversión?

Metodología: REPR-ALIGN

Los autores proponen REPR-ALIGN, un método de adaptación minimalista que trata el entrenamiento de DLM como un problema de alineación en lugar de un problema de aprendizaje de representaciones. La hipótesis central es que la "parte difícil" de la generación de lenguaje es aprender estructuras semánticas y sintácticas (representaciones), las cuales el preentrenamiento de AR ya ha dominado. La tarea restante es meramente adaptar estas características a un mecanismo de decodificación iterativo y de cualquier orden.

El método involucra dos modelos con arquitectura idéntica:

1. Profesor: Un transformador autoregresivo (AR) preentrenado y congelado con atención causal.
2. Estudiante: Un modelo de lenguaje de difusión enmascarado (DLM) con atención bidireccional, inicializado a partir de los mismos pesos de AR.

Procedimiento de Entrenamiento

Durante el entrenamiento, el modelo estudiante consume una secuencia corrupta (con tokens enmascarados) e intenta predecir los tokens enmascarados utilizando un objetivo estándar de desruido enmascarado. Simultáneamente, el modelo profesor congelado procesa la secuencia limpia bajo atención causal.

Para preservar la geometría interna del modelo AR, los autores introducen una pérdida de alineación de representaciones por capa. Para cada capa $\ell$ en la pila del transformador, los estados ocultos del estudiante ($h^{(\ell)}_D$) se alinean con los estados ocultos del profesor ($h^{(\ell)}_{AR}$) utilizando una pérdida de similitud del coseno.

El objetivo total de entrenamiento es:
$$\mathcal{L}(\theta) = \mathcal{L}_{diff}(\theta) + \lambda \mathcal{L}_{align}(\theta)$$
Donde:

• $\mathcal{L}_{diff}$ es la pérdida estándar de entropía cruzada de desruido enmascarado.
• $\mathcal{L}_{align}$ es la distancia coseno media entre los estados ocultos del estudiante y del profesor a través de todas las capas y posiciones enmascaradas.
• $\lambda$ es un hiperparámetro de ponderación (valor predeterminado 10).

Las decisiones clave de diseño incluyen:

• Sin Adaptadores: El método no requiere parámetros adicionales ni cambios arquitectónicos más allá de cambiar la máscara de atención de causal a bidireccional.
• Profesor Congelado: El profesor AR permanece congelado, proporcionando un sistema de coordenadas representacional estable.
• Convención de Desplazamiento: El método se adhiere a la convención estándar de desplazamiento utilizada en la adaptación de AR a DLM para reconciliar la indexación del siguiente token con el desruido por posición.

Contribuciones Clave

1. Identificación de la Preservación de Representaciones: El artículo identifica que la preservación explícita de la geometría de representación interna del modelo AR es un ingrediente faltante en las recetas actuales de conversión de AR→DLM.
2. Una Receta de Conversión Minimalista: REPR-ALIGN introduce un método simple, sin adaptadores, que ancla al estudiante DLM al profesor AR congelado mediante alineación coseno por capa. Este enfoque mejora la calidad de la conversión y la eficiencia de la muestra a través de las escalas de modelos.
3. Eficiencia en Datos y Cómputo: Los autores demuestran que la conversión de AR→DLM no es inherentemente hambrienta de datos. Con la alineación de representaciones, el entrenamiento en un subconjunto diminuto de datos (0.8B de tokens) puede superar al entrenamiento en el conjunto de datos completo (50B de tokens) bajo el mismo presupuesto de pasos. Además, congelar grandes bloques (incrustaciones y MLPs) mejora el rendimiento de entrenamiento hasta en 2$\times$ sin degradar la calidad.
4. Validación de Escalado: Los autores entrenaron un "oDLM" de 4B de parámetros utilizando esta receta. En comparación con Dream-7B (un DLM público fuerte también construido sobre inicialización de AR), oDLM logró mejores puntuaciones de HumanEval y HumanEval+ pass@10 a pesar de utilizar menos parámetros y un presupuesto de datos y cómputo significativamente más ligero.

Resultados Experimentales

El método fue evaluado en benchmarks de generación de código (HumanEval, HumanEval+, MBPP, MBPP+) utilizando checkpoints de Qwen3 en escalas de 0.6B, 1.7B y 4B.

• Velocidad de Adaptación: REPR-ALIGN aceleró consistentemente la adaptación. En HumanEval pass@10, el modelo alineado superó tanto al ajuste fino de AR como al entrenamiento desde cero durante la conversión temprana. A 200k pasos, el modelo de 0.6B vio una ganancia de 6.1 puntos, y el modelo de 1.7B vio una ganancia de 9.4 puntos sobre la línea base.
• Escalado: Los beneficios de la alineación crecieron con el tamaño del modelo. El modelo de 1.7B mostró mejoras absolutas mayores que el modelo de 0.6B, lo que sugiere que las geometrías de preentrenamiento de AR más grandes son más valiosas de preservar.
• Eficiencia de Datos: El entrenamiento en una muestra aleatoria de 0.8B de tokens con alineación produjo puntuaciones pass@10 más altas que el entrenamiento en el flujo completo de 50B de tokens con el mismo conteo de pasos.
• Congelamiento: Congelar las incrustaciones de tokens y los bloques MLP aumentó el rendimiento en $\sim$2$\times$ mientras mantenía o mejoraba ligeramente el rendimiento, apoyando la visión de que la señal de aprendizaje principal está en adaptar el mecanismo de desruido, no en reaprender representaciones.
• Ablaciones:
  • Métrica: La alineación del coseno superó a la distancia L2, lo que sugiere que la señal útil es geométrica (dirección) en lugar de métrica (magnitud).
  • Peso: $\lambda=10$ proporcionó el equilibrio óptimo entre anclaje y adaptación.
  • Capas: Alinear todas las capas proporcionó el mejor pass@1, mientras que las capas medias y superiores fueron particularmente importantes para pass@10.

Significado y Afirmaciones

El artículo afirma que las representaciones lingüísticas son universales independientemente del orden de generación. Al reformular el entrenamiento de DLM como un problema de alineación, los autores argumentan que el alto costo del preentrenamiento de difusión se debe en gran medida a la reaprendizaje innecesario de representaciones que el preentrenamiento de AR ya ha proporcionado.

El significado de este trabajo radica en ofrecer una "técnica de referencia" para entrenar modelos de lenguaje de difusión que:

• Reduce el Cómputo: Logra un rendimiento competitivo con una fracción del costo estándar de preentrenamiento.
• Simplifica el Entrenamiento: Elimina la necesidad de adaptadores complejos o ajuste fino pesado, confiando en cambio en un objetivo de alineación simple.
• Democratiza el Acceso: Hace viable el entrenamiento de DLM de alta calidad en regímenes de bajos datos y con recursos de cómputo limitados.

Los autores posicionan su trabajo como un cambio de "reentrenar" a "alinear", sugiriendo que los avances futuros en la generación discreta pueden provenir de una mejor alineación de mecanismos en lugar de paradigmas de entrenamiento desde cero.