🤖 machine learning

Training nGPT

Este artículo presenta una receta de entrenamiento práctica para el Transformer normalizado (nGPT) que incorpora técnicas como el Preacondicionamiento de Gradiente de Logit y GatedAdamW, permitiendo que modelos híbridos de Mamba-Transformer MoE de 14B de parámetros alcancen la misma pérdida de validación que sus contrapartes no normalizados utilizando aproximadamente la mitad de los tokens de entrenamiento.

Autores originales: Ilya Loshchilov, Boris Ginsburg

Publicado 2026-08-04

📖 7 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Ilya Loshchilov, Boris Ginsburg

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina el mundo de la inteligencia artificial como una biblioteca gigante y bulliciosa donde las computadoras están tratando de aprender a escribir historias, resolver problemas matemáticos y charlar como humanos. Para hacer esto, utilizan estructuras matemáticas especiales llamadas redes neuronales, que son como redes masivas e interconectadas de neuronas. Dentro de estas redes, la información viaja como números, y la red aprende ajustando la fuerza de las conexiones entre estas neuronas. Durante mucho tiempo, los científicos han notado que estos números pueden volverse desordenados: a veces crecen demasiado, otras veces se encogen hasta la nada, y todo el sistema puede confundirse sobre en qué dirección moverse para volverse más inteligente. Esto es un poco como intentar subir una montaña empinada y con niebla mientras cargas una mochila que cambia de peso constantemente; podrกว่า dar un paso adelante, solo para deslizarte de nuevo hacia abajo. El objetivo de esta investigación es encontrar una mejor manera de cargar esa mochila para que la computadora pueda escalar la montaña del aprendizaje de manera mucho más rápida y eficiente.

El artículo que vas a leer proviene del equipo de NVIDIA, liderado por Ilya Loshchilov y Boris Ginsburg. Ellos están abordando un problema específico: cómo hacer que sus modelos de IA aprendan mejor obligando a todos sus números internos a permanecer en una esfera perfecta e invisible. Piensa en esto como una regla que dice: "No importa qué tan lejos camines, siempre debes mantenerte exactamente a una milla de distancia del centro de la ciudad". Esta regla, llamada "normalización", evita que los números sean demasiado grandes o pequeños, lo que ayuda a la computadora a mantenerse enfocada. Los autores se basaron en una idea previa llamada "nGPT" (GPT normalizado) y lo probaron en un tipo de IA muy moderno y potente que mezcla dos tecnologías diferentes: "Mamba-2" (que es excelente para recordar secuencias largas) y "Transformers" (que son excelentes para entender el contexto). Querían ver si sus nuevas reglas de entrenamiento podían hacer que estos modelos híbridos aprendieran más rápido que el método estándar.

La Gran Idea: Una Nueva Receta de Entrenamiento

Los autores no solo ajustaron una cosa; cocinaron toda una nueva "receta de entrenamiento" para estos modelos de IA. Imagina que estás enseñando a un robot a caminar. La forma antigua (usando un método estándar llamado AdamW) es como dejar que el robot tropiece por doquier, a veces dando pasos gigantes y torpes y otras veces pasos diminutos y vacilantes, mientras espera eventualmente encontrar el camino. La nueva receta, que llaman la receta de entrenamiento nGPT, es como poner al robot en una cinta de correr con un conjunto muy específico de reglas para mantener sus pasos perfectos.

Aquí te explico cómo lo hicieron, desglosado en conceptos sencillos y cotidianos:

1. El "Preacondicionamiento del Gradiente de Logit" (El control de volumen)
Cuando la IA intenta adivinar la siguiente palabra en una oración, produce una lista de puntuaciones llamadas "logits". En el sistema antiguo, el volumen de estas puntuaciones podía volverse extrañamente fuerte o silencioso a medida que avanzaba el entrenamiento, confundiendo al robot sobre qué tan fuerte debía presionar. Los autores añadieron un control de volumen especial (un vector de escala aprendible) que podían ajustar. Pero aquí está la parte ingeniosa: se dieron cuenta de que si solo giraban el control, arruinarían la memoria del robot sobre cómo caminar. Así que inventaron un truco llamado Preacondicionamiento del Gradiente de Logit. Es como tener un control de volumen que cambia el sonido que escuchas, pero cuando el robot intenta aprender de un error, el sistema automáticamente baja el volumen a la normalidad para que el robot no se confunda. Esto mantiene el proceso de aprendizaje estable, incluso si la "intensidad" de las predicciones cambia.

2. El "Decaimiento de la Tasa de Aprendizaje Logarítmico" (El corredor de velocidad)
Normalmente, al entrenar una IA, comienzas con una velocidad de aprendizaje rápida y la desaceleras lentamente, como un conductor de carreras que suelta el pedal del acelerador gradualmente. Los autores descubrieron que la forma estándar de desacelerar (como una curva suave) no era la mejor para su nuevo sistema de "esfera". En su lugar, utilizaron un "Decaimiento de la Tasa de Aprendizaje Logarítmico". Imagina un coche de carreras que arranca con fuerza, reduce la velocidad muy rápido al principio para ubicarse, y luego mantiene un crucero constante y largo durante el resto de la carrera. Este programa "centrado al principio" permite que el modelo aprenda lo básico rápidamente y luego refine sus habilidades durante un largo periodo, lo cual resultó ser mucho más eficiente.

3. "GatedAdamW" (El portero inteligente)
La forma estándar de actualizar el cerebro del robot es AdamW. Los autores actualizaron esto a GatedAdamW. Piensa en el método estándar como un portero que deja pasar cada actualización, incluso si la actualización es diminuta y apenas perceptible. A veces, estas actualizaciones diminutas son solo ruido. El nuevo G gatedAdamW tiene un "portero inteligente" que observa cada actualización. Si una actualización es demasiado pequeña (como un susurro), el portero dice gentilmente: "Hoy no", y la bloquea. Si la actualización es un grito (un cambio significativo), el portero abre la puerta de par en par. Esto evita que el robot desperdicie energía en ajustes diminutos y de utilidad nula, y lo ayuda a concentrarse en los cambios grandes e importantes.

4. El "Límite de Paso Angular" (La correa de seguridad)
Dado que el robot está caminando sobre una esfera, dar un paso demasiado grande podría hacer que pierda el control o salte al lado equivocado del mundo. Los autores añadieron un "Límite de Paso Angular", que es como una correa de seguridad. Si el robot intenta dar un paso demasiado grande, la correa lo tira suavemente de vuelta a un ángulo más pequeño y seguro. Esto asegura que el robot nunca dé un salto salvaje y peligroso, manteniendo su viaje suave y estable.

Los Resultados: Aprendiendo de Manera Más Eficiente

El equipo probó esta nueva receta en una serie de modelos de IA, que van desde 1 mil millones hasta 14 mil millones de parámetros (el "tamaño del cerebro" del robot). Compararon sus nuevos modelos nGPT contra los modelos GPT estándar entrenados con los métodos antiguos.

Los resultados fueron impresionantes. Los nuevos modelos nGPT alcanzaron consistentemente tasas de error más bajas (medidas por la "pérdida de validación") que los modelos estándar. Específicamente, el modelo nGPT de 14 mil millones de parámetros alcanzó el mismo nivel de habilidad que el modelo estándar de 14 mil millones, pero requirió aproximadamente la mitad de los tokens de entrenamiento (la cantidad de texto que lee el modelo) para lograrlo. En otras palabras, para llegar al mismo destino, el modelo nGPT solo necesitó leer aproximadamente la mitad de la cantidad de texto comparado con el modelo estándar.

El artículo sugiere que esta mejora proviene de la combinación de todas estas nuevas reglas trabajando juntas, no solo de un único truco. También probaron una configuración específica para su "portero inteligente" (llamada el parámetro de nitidez $a$ ) y descubrieron que un portero más suave ( $a=0.5$ ) funcionaba ligeramente mejor que uno estricto, pero la mayor victoria vino de todo el nuevo sistema de entrenamiento, no solo del portero.

Lo Que Esto Significa

Los autores señalan con cautela que no intentaron todas las variaciones posibles de estas reglas (no hicieron un "estudio de ablación completo"), por lo que podría haber configuraciones aún mejores que no han encontrado. Sin embargo, los resultados que sí encontraron son sólidos: al obligar a la IA a caminar en una esfera perfecta y usar una forma más inteligente y controlada de dar pasos, podemos enseñar a estos modelos masivos a ser más inteligentes usando significativamente menos datos y tiempo. Es una receta práctica que sugiere que podemos construir una mejor IA sin necesidad de alimentarla con todo el internet, solo con una porción mucho más eficiente de este.

Resumen Técnico: Entrenamiento de nGPT en Modelos Híbridos MoE Modernos

Planteamiento del Problema

El marco de trabajo del Transformer normalizado (nGPT) propone representar los parámetros y las activaciones del modelo en una hiperesfera unitaria para facilitar el aprendizaje de representaciones hiperesféricas. Si bien la formulación original de nGPT (Loshchilov et al., 2024) demostró la viabilidad de este enfoque en entornos de Transformers densos, la aplicación a arquitecturas híbridas modernas y a gran escala —específicamente aquellas que combinan modelos de espacio de estados Mamba-2 con Mixture-of-Experts (MoE) de Transformer— requiere una receta de entrenamiento robusta. Los optimizadores estándar como AdamW, diseñados para espacios euclidianos no restringidos, pueden no interactuar de manera óptima con las restricciones geométricas de nGPT, particularmente en lo que respecta al escalado de gradientes, los cronogramas de tasa de aprendizaje y el manejo de estadísticas de segundo momento pequeñas.

Metodología

Los autores proponen una "receta de entrenamiento nGPT" integral adaptada para modelos híbridos Mamba-2–Transformer MoE (específicamente la arquitectura Nemotron-3). Esta receta integra restricciones de normalización con modificaciones de optimización específicas para estabilizar el entrenamiento y mejorar la convergencia.

1. Normalización y Escalado

Restricciones de Parámetros: Todos los vectores de parámetros que forman matrices están restringidos a una norma unitaria a lo largo de la dimensión de incrustación (embedding).
Normalización de Activaciones: Los estados ocultos se renormalizan tras la recombinación con las salidas de atención/SSM y MLP/MoE.
Vectores de Escala Aprendibles: Para evitar una sobre-restricción en las capas no lineales, el modelo emplea vectores de escala aprendibles:
- $s_{qk}$ para reescalar queries y keys normalizadas.
- $s_u$ para reescalar las activaciones intermedias de MLP/MoE.
- $s_z$ para reescalar los logits de salida para controlar la nitidez de la distribución.
Especificidades de Mamba-2 y MoE: Se introducen escalares entrenables ( $s_{mamba}$ , $s_{moe}$ ) para reescalar las entradas a las funciones SiLU y sigmoid, respectivamente, reemplazando las capas de normalización estándar como RMSNorm.

2. Ingredientes de Optimización

La receta introduce componentes novedosos en el proceso de optimización:

Precondicionamiento de Gradiente de Logit (LGP): En nGPT, el vector de escala de logit aprendido $s_z$ crece logarítmicamente durante el entrenamiento, actuando como un multiplicador oculto y variante en el tiempo sobre los gradientes. LGP desacopla este efecto normalizando la escala del gradiente hacia atrás: $\partial L / \partial u \leftarrow (s_z / \text{mean}(s_z))^q \odot \partial L / \partial z$ . Con $q=1$ , esto elimina la escala media global mientras preserva las variaciones relativas por vocabulario, estabilizando el flujo de gradientes sin alterar el paso hacia adelante (forward pass).
Decaimiento Logarítmico de la Tasa de Aprendizaje: En lugar del decaimiento de coseno estándar, los autores emplean un cronograma de decaimiento logarítmico. Esto concentra el presupuesto de la tasa de aprendizaje al inicio, permitiendo un descenso inicial más pronunciado seguido de una cola larga, lo cual se alinea mejor con la dinámica de parámetros restringidos de nGPT.
GatedAdamW: Este optimizador extiende AdamW separando la estabilidad numérica de la supresión de actualizaciones.
- Utiliza una constante numérica $\epsilon_{num}$ estrictamente para la estabilidad.
- Introduce una compuerta suave $\gamma_t = \sigma(a \log(d_t / \epsilon_{gate}))$ basada en la escala del segundo momento $d_t$ .
- Esto permite un control explícito sobre qué coordenadas reciben actualizaciones cuando las estadísticas del segundo momento son pequeñas, evitando la supresión de actualizaciones que podría ocurrir con un único valor $\epsilon$ en el AdamW estándar.
Controles Geométricos:
- Tope de Paso Angular (ASC): Limita el desplazamiento angular de los vectores de parámetros no de incrustación para asegurar la estabilidad geométrica.
- Recorte de Crecimiento de Segundo Momento (SMGC): Una salvaguarda contra picos extremos de gradiente mediante el recorte del crecimiento de los segundos momentos.
- Ruido de Exploración Post-Momento (PMEN): Añade ruido para fomentar la exploración en regiones de baja varianza.

Resultados Clave

Los autores evaluaron la receta en modelos híbridos Mamba-2–Transformer MoE que van desde 1B hasta 14B de parámetros totales, utilizando una mezcla de datos de grado industrial.

Eficiencia de Datos: El hallazgo principal es una mejora significativa en la eficiencia de datos. Específicamente, el modelo nGPT de 14B de parámetros alcanza la misma pérdida de validación que un baseline de GPT no normalizado (entrenado con AdamW) utilizando aproximadamente la mitad del número de tokens de entrenamiento. Este resultado se deriva de la interpolación de las curvas de pérdida–token en la Figura 7, en lugar de una comparación directa en tablas entre todos los tamaños de modelo.
Escalado Consistente: Las ganancias de rendimiento son consistentes en todos los tamaños de modelo (1B, 2B, 4B, 7B y 14B). Las variantes nGPT logran consistentemente pérdidas de validación aproximadamente un 2.5% menores que sus contrapartes no normalizadas entrenadas con AdamW.
Dinámicas de Convergencia: nGPT exhibe una disminución inicial de la pérdida más lenta seguida de una convergencia más rápida en etapas avanzadas del entrenamiento, un comportamiento atribuido a que el cronograma de la tasa de aprendizaje controla más directamente el tamaño de paso efectivo ante la ausencia de interacciones de decaimiento de peso (weight decay).
Análisis de Componentes: Los estudios de ablación sobre el parámetro de nitidez de la compuerta $a$ en GatedAdamW sugieren que, si bien $a=0.5$ ofrece una ventaja modesta sobre $a=1$ (que recupera el comportamiento estándar de Adam), la mayor parte de la ganancia de rendimiento proviene de la parametrización normalizada y de la receta de entrenamiento completa, más que del mecanismo de compuerta por sí solo.

Significado y Reivindicaciones

El artículo afirma proporcionar una receta de entrenamiento práctica que logra extender nGPT de Transformers densos a arquitecturas híbridas MoE modernas y a gran escala. Los autores enfatizan que la mejora central en la eficiencia de datos no es específica de los modelos MoE, sino que es el resultado de la sinergia entre la normalización y la optimización.

El trabajo se presenta con modestia respecto a su alcance:

Los experimentos se realizaron con un presupuesto de cómputo limitado, centrándose en una receta que "funcionó en la práctica" en lugar de un estudio de ablación exhaustivo de cada hiperparámetro.
Los autores declaran explícitamente que las ablaciones completas por componentes y las leyes de escala de hiperparámetros quedan para trabajos futuros.
La importancia radica en demostrar que el aprendizaje de representación hiperesférica puede escalarse a modelos de 14B de parámetros con ganancias sustanciales en la eficiencia del entrenamiento, ofreciendo una alternativa viable a los pipelines de entrenamiento estándar de Transformer.