Not all tokens are needed(NAT): token efficient reinforcement learning

El artículo presenta NAT, un marco unificado que optimiza el aprendizaje por refuerzo en modelos de lenguaje al actualizar la política utilizando solo un subconjunto de tokens mediante un estimador de gradiente no sesgado, logrando un rendimiento comparable al método completo con una reducción significativa en costos computacionales y uso de memoria.

Lo esencial

Imagina que estás entrenando a un estudiante muy inteligente (una Inteligencia Artificial) para que resuelva problemas de matemáticas muy difíciles. Para aprender, el estudiante debe "pensar en voz alta", escribiendo paso a paso todo su razonamiento. A esto le llamamos Cadena de Pensamiento (o Chain-of-Thought).

El problema es que, para que el estudiante mejore, el profesor (el algoritmo de aprendizaje) tiene que revisar cada palabra que escribió el estudiante, desde la primera hasta la última, para corregir sus errores. Si el estudiante escribe una respuesta de 1.000 palabras, el profesor tiene que leer y corregir las 1.000. Esto consume muchísima energía, tiempo y memoria de la computadora, como si tuvieras que revisar cada letra de un libro entero solo para encontrar un par de errores.

Los autores de este paper se preguntaron: ¿Realmente necesitamos revisar cada palabra para que el estudiante aprenda?

La respuesta es: No.

Presentan una nueva técnica llamada NAT (Not All Tokens are Needed o "No todos los tokens son necesarios"). Aquí te explico cómo funciona con analogías sencillas:

1. El Problema: La "Tasa Oculta" de las Palabras

Imagina que el estudiante escribe un ensayo de 10 páginas.

• Páginas 1 a 3: Son la introducción y los pasos lógicos importantes.
• Páginas 4 a 8: Son solo el estudiante repitiendo lo que ya dijo, escribiendo "y luego...", "además...", o rellenando espacio.
• Página 9 y 10: La conclusión y la respuesta final.

El método tradicional (GRPO) revisa las 10 páginas por igual. Esto es como si un entrenador de fútbol revisara cada paso que da un jugador, incluso cuando el jugador solo está caminando por el campo sin el balón. Es un desperdicio de energía.

2. La Solución: NAT (El Profesor Inteligente)

NAT propone que el profesor no revise todo el texto, sino que seleccione solo una parte importante para corregir, pero de una manera muy especial para no engañar al estudiante.

Usan dos trucos principales:

A. Muestreo Aleatorio (URS): "El Sorteo"

Imagina que el profesor cierra los ojos y marca al azar 50% de las palabras del ensayo para corregirlas.

• Ventaja: Ahorra tiempo de corrección.
• Desventaja: Como el profesor sigue teniendo que "leer" todo el texto antes de marcar las palabras (porque la computadora necesita procesar la historia completa para saber qué sigue), no ahorra mucha energía al principio. Es como si el profesor leyera todo el libro, pero solo corrigiera la mitad de las páginas.

B. El Corte Aleatorio del Prefijo (RPC): "La Historia Corta"

Este es el truco estrella. En lugar de marcar palabras sueltas al azar, el profesor decide: "Hoy solo voy a leer y corregir los primeros 500 caracteres de tu respuesta".

• La magia: Como la computadora solo necesita procesar esos primeros 500 caracteres para hacer la corrección, ahorra mucha memoria y energía (no tiene que "pensar" en las palabras que no va a corregir).
• El riesgo: ¿Y si la respuesta importante estaba al final? ¿No aprenderá el estudiante?
• La solución mágica (Horvitz-Thompson): Aquí entra la parte matemática brillante. El profesor usa una fórmula de "reajuste". Si decide corregir solo la mitad del texto, duplica la importancia de las correcciones que hace en esa mitad.
  • Analogía: Imagina que tienes que pagar 100 dólares. Si solo revisas 50 dólares de tu cuenta, pero decides que esos 50 valen el doble en importancia para el cálculo final, el resultado matemático es el mismo que si hubieras revisado los 100.
  • Gracias a esto, el estudiante aprende exactamente lo mismo que si el profesor hubiera revisado todo el texto, pero usando la mitad de recursos.

3. ¿Por qué es mejor que simplemente "cortar" el texto?

Otra idea sería simplemente decir: "Siempre voy a cortar la respuesta a la mitad".

• El problema: Si siempre cortas al 50%, el estudiante nunca aprende a escribir la conclusión o la parte final. Se vuelve "tonto" al final de la frase.
• La ventaja de NAT (RPC): NAT corta el texto en un punto aleatorio cada vez. A veces corta al 30%, a veces al 70%, a veces al 50%.
  • Esto asegura que, a lo largo de muchas sesiones de entrenamiento, el estudiante tenga la oportunidad de aprender de todas las partes de la respuesta (principio, medio y final), pero sin que la computadora tenga que procesar todo cada vez.

Los Resultados en la Vida Real

Los autores probaron esto con modelos de inteligencia artificial reales (como Qwen3) resolviendo problemas de matemáticas:

• Ahorro de memoria: Redujeron el uso de la memoria de la tarjeta gráfica en un 18%.
• Ahorro de tiempo: El entrenamiento fue un 29% más rápido.
• Calidad: ¡El estudiante aprendió igual de bien! Obtuvo las mismas calificaciones que si se hubiera revisado todo el texto.

En Resumen

Este paper nos dice que, al entrenar inteligencias artificiales para pensar mucho, no necesitamos revisar cada palabra. Podemos ser más inteligentes: revisar solo una parte aleatoria, pero "pesar" esa parte correctamente para que el aprendizaje sea justo.

Es como si, en lugar de leer todo un libro para aprender una lección, el profesor te dijera: "Lee solo el primer capítulo, pero piensa en él con el doble de intensidad, y mañana leeremos otro capítulo diferente". Al final, habrás aprendido todo el libro, pero habrás gastado la mitad de la energía.

Resumen técnico

1. El Problema: La "Tasa Oculta" de los Tokens en RLVR

El Aprendizaje por Refuerzo a partir de Recompensas Verificables (RLVR) se ha convertido en un motor clave para el avance de los Modelos de Lenguaje Grandes (LLM) en tareas de razonamiento complejo (como matemáticas y código). Sin embargo, escalar RLVR a trayectorias largas de "Cadena de Pensamiento" (Chain-of-Thought, CoT) enfrenta un cuello de botella crítico:

• Costo Computacional Excesivo: Los pipelines estándar de RL (como GRPO) realizan actualizaciones de política (backpropagation) sobre cada token generado en la trayectoria.
• Desajuste de Eficiencia: Aunque la generación de muestras (rollout) se ha optimizado, la fase de aprendizaje (forward/backward pass) sigue siendo intensiva en memoria y cómputo.
• Limitación de Escala: A medida que las trayectorias de razonamiento se vuelven más largas, el costo de retropropagar a través de todos los tokens se convierte en un "impuesto oculto", limitando la capacidad de entrenar modelos con contextos más largos o aumentando el riesgo de desbordamiento de memoria (OOM).

La pregunta central del trabajo es: ¿Realmente necesitamos todos los tokens para entrenar razonadores de RL efectivos? La intuición sugiere que no todos los tokens contribuyen equitativamente al señal de aprendizaje; muchos son continuaciones mecánicas o boilerplate, mientras que solo un subconjunto representa puntos de decisión de alto impacto.

2. Metodología: El Marco NAT (Not All Tokens are Needed)

Los autores proponen NAT, un marco unificado que trata el presupuesto de tokens como un primitivo de optimización de primer nivel. NAT actualiza la política utilizando solo un subconjunto seleccionado de tokens generados, pero mantiene la señal de recompensa calculada sobre la respuesta completa.

Componentes Clave:

1. Enmascaramiento de Tokens y Muestreo:
   NAT introduce una máscara binaria $m_{i,t}$ que determina si un token $t$ de la trayectoria $i$ participa en la actualización de la política. Se proponen dos esquemas de selección:

   • Muestreo Aleatorio Uniforme (URS): Selecciona tokens individuales de forma independiente con una probabilidad $p$.
   • Corte Aleatorio de Prefijo (RPC): Selecciona un prefijo contiguo de la respuesta (de longitud $L_i$) y descarta el sufijo. Esto es crucial porque permite reducir el cómputo del forward pass en transformadores causales.

2. Corrección de Horvitz-Thompson (HT):
   Para garantizar que el entrenamiento no se sesgue al ignorar tokens, NAT utiliza el estimador de Horvitz-Thompson.

   • Cada token seleccionado se repondera por el inverso de su probabilidad de inclusión ($1/p_{i,t}$).
   • Proposición de Sesgo Cero: Se demuestra teóricamente que este estimador ponderado es insesgado. Es decir, el gradiente esperado de la actualización parcial es idéntico al gradiente de la secuencia completa, preservando la señal de aprendizaje original en expectativa.

3. Integración con GRPO:
   El método se implementa sobre el algoritmo GRPO (Group Relative Policy Optimization). La recompensa y la ventaja se calculan sobre la respuesta completa, pero la pérdida se calcula y retropropaga solo sobre los tokens seleccionados y reponderados.

3. Contribuciones Principales

1. Marco Unificado Token-Eficiente: NAT permite actualizaciones de RLVR utilizando subconjuntos de tokens sin modificar el cálculo de recompensas ni el pipeline de generación.
2. Fundamento Teórico Insesgado: La aplicación de la corrección de Horvitz-Thompson establece una conexión rigurosa entre el enmascaramiento de tokens y la optimización fiel del objetivo original de RLVR, evitando el sesgo sistemático que sufren métodos de truncamiento determinista.
3. Estrategias Prácticas (RPC): La introducción de Random Prefix Cutting (RPC) como una solución superior a URS. RPC permite ahorros reales en el forward pass (al no procesar el sufijo truncado) y en el backward pass, manteniendo la estructura causal del prefijo.

4. Resultados Experimentales

Los experimentos se realizaron en benchmarks de razonamiento matemático (MATH, AIME24, AIME25) utilizando modelos Qwen2.5-Math-7B y Qwen3-8B.

• Rendimiento (Precisión):

  • NAT (RPC y URS): Logra un rendimiento comparable al GRPO completo (full-token), con intervalos de confianza que se superponen en las métricas de precisión (Acc@16 y Pass@16).
  • Truncamiento Determinista: El método que simplemente corta el 50% final de los tokens (Det. Trunc.) muestra un rendimiento significativamente peor y mayor varianza, confirmando que ignorar sistemáticamente los tokens finales (donde a menudo se encuentra la respuesta o verificación) es perjudicial.

• Eficiencia de Recursos:

  • Memoria GPU: RPC reduce la memoria pico de GPU en aproximadamente un 18-20% (ej. de 47.72 GB a 39.23 GB en Qwen3-8B). URS no reduce significativamente la memoria pico porque el forward pass sigue procesando la secuencia completa.
  • Tiempo de Entrenamiento: RPC reduce el tiempo de entrenamiento (sin inferencia) en un 29% y el tiempo total por paso en un 36% para Qwen3-8B. Esto se debe a que RPC acorta la longitud efectiva de la secuencia procesada en ambos pasos (forward y backward).

• Estabilidad: Las curvas de entropía muestran que NAT converge de manera similar al GRPO completo, mientras que el truncamiento determinista exhibe una entropía más alta e inestable.

5. Significado e Impacto

Este trabajo ofrece una vía ortogonal a las optimizaciones de sistemas existentes (como motores de inferencia más rápidos o decodificación especulativa). En lugar de hacer que la generación sea más rápida, NAT hace que el consumo de esas trayectorias sea más eficiente.

• Escalabilidad: Permite entrenar RL en trayectorias de contexto más largo sin aumentar los costos de memoria o cómputo de manera prohibitiva.
• Eficiencia Teórica y Práctica: Demuestra que se puede reducir el presupuesto de tokens a la mitad (50%) sin sacrificar la calidad del razonamiento, siempre que se utilice una corrección estadística adecuada (HT) y una estrategia de muestreo que respete la causalidad (RPC).
• Futuro: Abre la puerta a futuras investigaciones sobre selección de tokens "consciente de la información" (basada en incertidumbre o magnitud del gradiente) dentro del mismo marco de Horvitz-Thompson, y a la co-diseño de kernels de sistema para aprovechar estas estructuras de prefijo truncado.

En resumen, NAT demuestra que "no todos los tokens son necesarios" para el aprendizaje, proporcionando una solución teóricamente fundamentada y empíricamente validada para escalar el RL en modelos de razonamiento avanzado.