NVIDIA lanza Nemotron-Labs-TwoTower: un modelo de lenguaje de difusión de peso abierto construido sobre una columna vertebral autorregresiva congelada Nemotron-3-Nano-30B-A3B

NVIDIA ha lanzado Nemotron-Labs-TwoTower, un modelo de lenguaje de difusión construido sobre una columna vertebral autorregresiva previamente entrenada. Se envía como pesas abiertas bajo la licencia de modelo abierto NVIDIA Nemotron. La versión apunta a un cuello de botella en el rendimiento en la generación de texto.

Los modelos autorregresivos (AR) decodifican un token a la vez. Ese proceso en serie limita el rendimiento de generación. Los modelos de lenguaje de difusión discreta toman otra ruta. Generan tokens en paralelo y los refinan de forma iterativa.

La mayoría de los modelos de lenguaje de difusión utilizan una red para dos trabajos. Representa tokens limpios y elimina el ruido de los corruptos en cada paso. TwoTower separa estos trabajos en dos torres. Mantiene el 98,7% de la calidad de referencia agregada de la línea base AR. También informa un rendimiento de generación de relojes de pared 2,42 veces mayor.

TL;DR

TwoTower divide la difusión en una torre de contexto AR congelada y una torre de eliminación de ruido entrenada. Conserva el 98,7% de la calidad AR con un rendimiento de 2,42× (γ=0,8, S=16, 2×H100). El eliminador de ruido se entrenó en ~2,1T de tokens; la columna vertebral usaba 25T. Un punto de control ejecuta los modos de difusión, AR simulado y decodificación AR.

Nemotron-Labs-TwoTower

TwoTower es un modelo de difusión autorregresivo por bloques. Se crea una instancia en Nemotron-3-Nano-30B-A3B, una columna vertebral híbrida de peso abierto. Esa columna vertebral entrelaza capas de Mamba-2, autoatención y mezcla de expertos (MoE).

Cada torre tiene 52 capas: 23 Mamba-2, 6 de autoatención y 23 MoE. El puesto de control liberado envía ambas torres, aproximadamente 60 mil millones de parámetros en total. Los parámetros activos por token son aproximadamente 3 mil millones por torre. El MoE utiliza 128 expertos enrutables, de los cuales 6 se activan, más 2 expertos compartidos.

Ambas torres comienzan como copias del mismo punto de control principal. Sólo se entrena la torre de eliminación de ruido. La torre de contexto AR permanece congelada. El eliminador de ruido se entrenó con ~2,1T de tokens, una fracción del preentrenamiento de 25T de tokens de la columna vertebral.

Cómo funcionan las dos torres

La torre de contexto AR se ejecuta causalmente sobre los tokens solicitados y comprometidos. Produce caché KV por capa y estados finales de Mamba-2. Preserva la capacidad autorregresiva de la columna vertebral.

La torre de difusión de ruido refina los bloques ruidosos. Dentro de un bloque, utiliza atención bidireccional dentro del bloque. Sigue siendo causal con respecto a bloques limpios pasados.

Las torres se conectan capa por capa. La capa i de eliminación de ruido atiende cruzadamente a la capa i de la torre de contexto. Esta atención cruzada alineada en capas brinda acceso a múltiples escalas a las representaciones de la columna vertebral. Los enfoques anteriores sólo difunden el último estado oculto.

Dos modificaciones más para eliminar el ruido son importantes. Las capas de Mamba-2 generan su estado inicial a partir del estado Mamba de la torre de contexto. El paso de tiempo de difusión modula cada capa a través del acondicionamiento de tiempo único adaLN. Ese módulo adaLN agrega solo ~1,5 millones de parámetros.

La generación se ejecuta bloque a bloque. Cada bloque comienza como S [MASK] fichas. El eliminador de ruido lo refina en T pasos y luego lo confirma. Luego, la torre de contexto procesa los tokens comprometidos para actualizar sus cachés.

Esto explica por qué múltiples pasos de eliminación de ruido aún pueden superar la decodificación de un token. La decodificación autorregresiva confirma exactamente un token por paso. TwoTower confirma varios tokens por paso al principio del refinamiento.

Puntos de referencia

Las evaluaciones utilizan BF16 en 2 GPU H100. El punto operativo predeterminado es el desenmascaramiento de confianza, umbral γ=0,8, tamaño de bloque S=16. La tabla compara la línea base de AR con la decodificación de difusión TwoTower.

TaskNemotron-3-Nano-30B-A3B (AR)Nemotron-Labs-TwoTower (difusión)MMLU (5 disparos, acc)78.5678.24MMLU-Pro (5 disparos, CoT EM)62.5960.93ARC-Challenge (25 disparos, acc_norm)91.7292.66WinoGrande (5 disparos, acc)76.0976.09RACE (0 disparos, acc)88.9088.90HumanEval (0 disparos)79.2775.58MBPP-Sanitized (3 disparos)74.7174.28GSM8K (8 disparos, acc)92.4990.14MATH-500 (4 disparos)84.4080.60MMLU Global Lite (5 disparos) 73,9773,94 MGSM (8 disparos, acc promedio) 80,8080,40 Calidad retenida 100 % 98,7 % Rendimiento de generación (× AR) 1,0 × 2,42 ×

El conocimiento general se mantiene dentro de aproximadamente un punto de la línea base AR. El código y las matemáticas muestran una degradación modesta. Se recuperan o mejoran ligeramente las puntuaciones de sentido común y multilingües. Reducir γ compromete más tokens por paso y aumenta el rendimiento, con una calidad reducida.

Ejecutándolo: modos de tres generaciones

El punto de control expone tres caminos de inferencia. La difusión completa de dos torres utiliza 2 GPU, aproximadamente 59 GB por GPU en BF16. El modo solo AR se ejecuta en una única GPU de 80 GB.

importar antorcha desde transformadores importar AutoTokenizer, AutoModelForCausalLM model_name = “nvidia/Nemotron-Labs-TwoTower-30B-A3B-Base-BF16” tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=torch.bfloat16, trust_remote_code=True, ) # torre de contexto -> GPU 0, torre de eliminación de ruido -> GPU 1 model.place_towers_on_devices(“cuda:0”, “cuda:1”) model.eval() mensaje = “Francia es un país ” entradas = tokenizer(prompt, return_tensors=”pt”).to(“cuda:0”) salidas = model.generate_mask_diffusion( entradas[“input_ids”]max_new_tokens=128, tamaño_bloque=16, pasos_por_bloque=16, mask_token_id=3, temperatura=0.1, umbral_confianza=0.8, eos_token_id=tokenizer.eos_token_id, ) print(tokenizer.decode(salidas)[0][inputs[“input_ids”].forma[1]:], skip_special_tokens=Verdadero))

Los tres modos son generate_mask_diffusion(), generate_mock_ar() y generate_ar(). La difusión de máscara compromete hasta tokens de tamaño de bloque por paso. Mock-AR y AR comprometen un token por paso.

Dónde encaja: casos de uso

El caso de uso más directo es la generación de lotes más rápida. Un equipo de datos que produce texto sintético puede cambiar una pequeña caída de calidad por rendimiento. Con γ=0,8, esa operación tiene una calidad del 1,3% para una velocidad de 2,42×.

Un segundo caso de uso es ajustar la relación calidad-rendimiento. Aumentar γ preserva más calidad, según el artículo de NVIDIA. Reducir γ compromete más tokens por paso para aumentar la velocidad.

Un tercer caso de uso es la adaptación directa. La torre de contexto mantiene su cabeza LM para decodificación especulativa, verificación o puntuación AR. Los equipos pueden ejecutar AR y difusión desde un punto de control.

Fortalezas y debilidades

Fortalezas:

Pesos abiertos bajo la licencia de modelo abierto NVIDIA Nemotron; listo para uso comercial 98,7 % de la calidad de AR retenida con un rendimiento de 2,42 × en el punto operativo predeterminado Un punto de control admite difusión, AR simulado y decodificación de AR Denoiser entrenado en ~2,1 T de tokens, no un re-entrenamiento completo La memoria caché de longitud de secuencia se escala como la línea base de AR

Debilidades:

La difusión completa de dos torres necesita 2 GPU y ~59 GB por GPU en BF16 El código y las matemáticas se degradan más que el conocimiento general (HumanEval 79.27 → 75.58) Mantener ambas torres residentes aumenta la huella de memoria de peso fijo del modelo El punto de control liberado es un modelo base, antes del ajuste o la alineación de las instrucciones. El rendimiento superior a 3 veces conlleva una mayor pérdida de calidad.

Explicador interactivo

¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros