Resumen Técnico: Nemotron-TwoTower

Planteamiento del Problema

Los modelos de lenguaje autorregresivos (AR) dominan la generación de texto, pero sufren un cuello de botella en el rendimiento debido a su proceso de decodificación secuencial, un token a la vez. Los modelos de lenguaje de difusión discreta ofrecen una alternativa prometedora al generar tokens en paralelo y refinarlos iterativamente, lo que potencialmente ofrece un mayor rendimiento y un control más fino. Sin embargo, los enfoques de difusión existentes suelen depender de una única red para desempeñar dos roles distintos simultáneamente: representar los tokens de contexto limpios y denotar iterativamente los tokens corruptos. Este entrelazamiento obliga a un único conjunto de pesos a servir a objetivos conflictivos, lo que limita la capacidad del modelo para sobresalir en cualquiera de los dos roles. Además, no está claro si es posible desacoplar completamente estos roles a gran escala, particularmente en arquitecturas híbridas modernas que combinan capas Mamba, de atención y de Mezcla de Expertos (MoE).

Metodología

Los autores proponen TwoTower, un modelo de difusión autorregresivo por bloques que desacopla los roles de representación de contexto y de denotación en dos torres separadas, construidas sobre la arquitectura híbrida Nemotron-3-Nano-30B-A3B (que intercala capas Mamba-2, de autoatención y de MoE).

Arquitectura

El modelo utiliza dos copias de la red preentrenada con roles complementarios:

1. Torre de Contexto AR Congelada: Esta torre permanece congelada y actúa como un modelo autorregresivo causal sobre el prompt y los tokens previamente comprometidos. Procesa tokens limpios para producir cachés de Clave-Valor (KV) alineadas por capa y estados de Mamba. Preserva las capacidades autorregresivas originales del backbone y mantiene estados reutilizables a través de los bloques.
2. Torre de Denotación de Difusión Entrenable: Esta torre se entrena para refinar bloques ruidosos. Emplea atención de bloque bidireccional dentro del bloque ruidoso actual (permitiendo que los tokens se atiendan entre sí) mientras mantiene una atención cruzada causal a las capas correspondientes de la torre de contexto congelada. Esta atención cruzada alineada por capa proporciona al denotador acceso multiescala a las representaciones del backbone, contrastando con métodos previos que solo transmitían el estado oculto final.

Proceso de Generación

La generación procede bloque por bloque (Algoritmo 1):

• La torre de contexto procesa el prompt y los tokens comprometidos de forma causal.
• Se inicializa un nuevo bloque con $S$ tokens enmascarados.
• La torre de denotación refina iterativamente este bloque durante $T$ pasos. En cada capa, el denotador atiende al caché KV de la torre de contexto para los bloques pasados y a su propio caché KV bidireccional para el bloque actual.
• Una vez que un bloque es suficientemente refinado (basado en umbrales de confianza), se compromete. La torre de contexto actualiza entonces sus estados con los nuevos tokens limpios, y el proceso se repite para el siguiente bloque.

Entrenamiento y Condicionamiento

• Difusión con Máscaras: El denotador se entrena utilizando un objetivo de difusión con máscaras. Los bloques limpios se corrompen reemplazando tokens con [MASK] de acuerdo con un esquema lineal. El modelo minimiza la log-verosimilitud negativa de los tokens originales en las posiciones enmascaradas.
• Condicionamiento Temporal: El denotador está condicionado al paso de tiempo de difusión $t$ mediante la normalización de capa adaptativa (adaLN), una técnica común en la difusión de imágenes pero menos utilizada en modelos de lenguaje con máscaras. Esto añade parámetros mínimos (~1.5M) pero mejora la calidad del muestreo.
• Currículo: El entrenamiento sigue un currículo de dos etapas. La Fase 1 utiliza una mezcla de datos diversa para una cobertura amplia, mientras que la Fase 2 aumenta el peso de fuentes de alta calidad y centradas en STEM. El tamaño del bloque de entrenamiento se reduce gradualmente (de $S=32$ a $S=16$) para equilibrar el paralelismo con la calidad del condicionamiento.
• Muestreo: Los autores emplean una estrategia de desenmascarado por confianza (Algoritmo 2). En cada paso, el denotador predice todos los tokens enmascarados en paralelo; los tokens con una confianza superior a un umbral $\gamma$ se comprometen, mientras que los tokens inciertos permanecen enmascarados para un mayor refinamiento. Esto hace que el número de tokens comprometidos sea adaptativo.

Resultados Clave

Evaluado sobre el backbone Nemotron-3-Nano-30B-A3B (30B de parámetros, MoE híbrido), Nemotron-TwoTower demuestra lo siguiente:

• Compromiso Calidad-Rendimiento: El modelo retiene el 98.7% de la calidad agregada de los benchmarks de la línea base autorregresiva, logrando al mismo tiempo una mejora de 2.42× en el rendimiento de generación en tiempo real (wall-clock).
• Desempeño en Benchmarks:
  • Conocimiento General: El desempeño se mantiene dentro de ~1 punto de la línea base AR.
  • Código y Matemáticas: Muestra una degradación modesta en comparación con la línea base.
  • Sentido Común y Multilingüe: El desempeño se recupera o mejora ligeramente.
• Ablaciones de Diseño:
  • Mamba Bidireccional: Probar una variante de Mamba bidireccional mostró ganancias marginales de calidad pero duplicó el cómputo de SSM; el diseño final mantiene Mamba causal.
  • Condicionamiento Temporal: El condicionamiento temporal explícito mejoró significativamente las puntuaciones de generación, código y matemáticas.
  • Desacoplamiento de Torres: Mantener la torre de contexto congelada y entrenar una torre de denotación separada produjo los mejores resultados. Vincular los pesos entre las torres o continuar el entrenamiento AR en la torre de contexto condujo a caídas significativas de precisión.
• Sensibilidad del Tamaño de Bloque: Los tamaños de bloque de entrenamiento más pequeños ($S=16$) mejoraron la calidad sobre los tamaños más grandes ($S=128$), pero redujeron el rendimiento. Los autores seleccionaron $S=16$ como el predeterminado para equilibrar la retención de calidad con la velocidad. El muestreo con bloques más grandes que el bloque de entrenamiento resultó en una degradación severa del desempeño.

Significado y Reivindicaciones

El artículo afirma que Nemotron-TwoTower demuestra con éxito que las representaciones autorregresivas preentrenadas pueden soportar un denotador de difusión entrenable sin necesidad de re-preentrenar todo el modelo. Al desacoplar los roles de contexto y denotación, este enfoque preserva las capacidades amplias del modelo híbrido MoE de 30B original, al tiempo que desbloquea los beneficios de la generación en paralelo de la difusión.

Los autores enfatizan que esta recuperación de calidad se logra mediante la adaptación en una fracción de los datos utilizados para el preentrenamiento original (~2.1T de tokens frente a 25T de tokens). El trabajo establece una nueva frontera de Pareto para los modelos de lenguaje, ofreciendo un camino viable hacia una generación de mayor rendimiento sin sacrificar la calidad fundamental de los modelos autorregresivos a gran escala. El código y los pesos del modelo se liberan para facilitar la investigación adicional en esta dirección.