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

Each language version is independently generated for its own context, not a direct translation.

Imaginez que vous apprenez à un élève très intelligent (une Intelligence Artificielle) à résoudre des problèmes de mathématiques complexes. Pour cela, vous lui donnez un énoncé et vous attendez qu'il écrive toute sa démarche, étape par étape, sur un tableau noir. C'est ce qu'on appelle une "chaîne de pensée".

Le problème, c'est que pour apprendre, l'ordinateur doit relire chaque mot de cette longue explication, du premier au dernier, pour corriger ses erreurs. C'est comme si un professeur devait relire et annoter chaque virgule d'un roman entier pour donner une note à un élève. Cela prend énormément de temps, d'énergie et de mémoire, surtout quand les explications deviennent très longues.

C'est là qu'intervient le papier que vous avez partagé, qui propose une méthode géniale appelée NAT ("Pas tous les jetons sont nécessaires").

Voici l'explication simple, avec quelques analogies :

1. Le Problème : La "Taxe Cachée" des Mots

Actuellement, pour entraîner ces IA, on utilise une méthode où l'ordinateur doit faire un aller-retour mental sur tous les mots générés.

L'analogie : Imaginez que vous devez nettoyer une très longue tapisserie. La méthode actuelle vous oblige à frotter chaque centimètre carré, même les zones où il n'y a aucune tache, juste parce que c'est la règle. C'est épuisant et ça consomme beaucoup d'eau (de mémoire et de puissance de calcul).

2. La Solution : NAT (Le Tri Intelligent)

Les auteurs de NAT se sont dit : "Est-ce qu'on a vraiment besoin de relire chaque mot pour apprendre ?"
La réponse est non. Souvent, les derniers mots d'une phrase mathématique sont juste des répétitions ou des étapes mécaniques. Ce qui compte vraiment, ce sont les moments clés où l'élève a pris une décision importante.

NAT propose de ne relire et de ne corriger qu'un sous-ensemble de mots choisis au hasard, tout en gardant la note finale basée sur la réponse complète.

3. Les Deux Méthodes de NAT

Le papier teste deux façons de faire ce tri :

A. Le "Tirage au Sort" (URS)

C'est comme si le professeur prenait un stylo rouge et barrait au hasard 50 % des mots de la copie de l'élève pour ne les corriger que sur les autres.

Le problème : Même si le professeur ne corrige pas les mots barrés, il a dû lire tout le texte pour savoir où ils étaient. L'effort de lecture (le "calcul vers l'avant") n'est pas réduit. On économise un peu de temps de correction, mais pas beaucoup.

B. La "Coupe de Prefixe Aléatoire" (RPC) - La Star de l'histoire

C'est la méthode la plus ingénieuse. Au lieu de barrer des mots au hasard dans le texte, le professeur dit : "Je vais lire et corriger seulement les 50 % premiers mots de ta copie, et je m'arrête là."

Pourquoi c'est génial ?
1. Économie de lecture : Le professeur n'a même pas besoin de voir la fin de la copie. Il arrête de lire à mi-chemin. Cela économise énormément de temps et d'énergie.
2. Économie de mémoire : Il n'a pas besoin de garder en tête toute l'histoire, juste le début.
3. Le tour de magie mathématique (Repondération) : Pour s'assurer que l'élève n'apprend pas des choses fausses parce qu'on ne lui a pas montré la fin, le professeur utilise une astuce mathématique (appelée repondération de Horvitz-Thompson). C'est comme si le professeur disait : "J'ai corrigé seulement la moitié de ta copie, donc je vais doubler l'importance de chaque erreur que j'ai trouvée pour que cela compte autant que si j'avais tout lu."

4. Les Résultats Concrets

Les chercheurs ont testé cette méthode sur des modèles d'IA (comme Qwen) avec des problèmes de maths difficiles.

Performance : L'IA apprend aussi bien avec la méthode "Coupe de Prefixe" (RPC) qu'avec la méthode classique qui lit tout. Elle ne perd pas en intelligence.
Vitesse et Mémoire :
- L'entraînement est plus rapide (jusqu'à 29 % plus rapide).
- L'IA a besoin de moins de mémoire (la carte graphique chauffe moins et ne plante pas aussi facilement).

En Résumé

Imaginez que vous voulez apprendre à conduire.

L'ancienne méthode : Vous devez conduire 100 km, puis l'instructeur vous fait refaire 100 km en arrière, en vous expliquant chaque virage, chaque freinage, même ceux qui étaient parfaits.
La méthode NAT (RPC) : L'instructeur vous laisse conduire 50 km, s'arrête, et vous dit : "Très bien, on va analyser ces 50 km. On va doubler l'importance de tes bons réflexes pour que tu apprennes aussi vite que si on avait fait 100 km."

Le résultat ? Vous apprenez aussi bien, mais vous y passez la moitié du temps et vous usez la moitié de l'essence (de la mémoire de l'ordinateur). C'est une façon plus intelligente et économe de faire grandir les intelligences artificielles.

Each language version is independently generated for its own context, not a direct translation.

1. Problématique

L'apprentissage par renforcement (RL) est devenu un moteur clé pour l'amélioration des modèles de langage (LLM), en particulier pour les tâches de raisonnement complexe nécessitant des chaînes de pensée (Chain-of-Thought ou CoT) longues. Cependant, l'extension du RL à ces trajectoires longues se heurte à des contraintes majeures de coût et d'évolutivité :

Coût de la rétropropagation (Backpropagation) : Les pipelines standards de RL (comme GRPO - Group Relative Policy Optimization) traitent tous les tokens générés comme également importants. Ils calculent les gradients et rétropropagent à travers chaque token de chaque trajectoire.
Goulot d'étranglement mémoire et calcul : À mesure que les trajectoires s'allongent pour permettre un raisonnement plus profond, la mémoire d'activation et le nombre d'opérations (FLOPs) augmentent considérablement. Cela peut entraîner des échecs de mémoire (OOM) ou une instabilité de l'optimisation.
Inefficacité des optimisations actuelles : Bien que les moteurs d'inférence aient été optimisés pour accélérer la génération (rollout), la phase d'apprentissage (forward/backward pass) reste limitée par la mémoire et le calcul, créant un plafond d'efficacité.
Hypothèse sous-jacente : Tous les tokens d'une longue chaîne de pensée ne contribuent pas également au signal d'apprentissage. Beaucoup sont des continuations mécaniques ou du "boilerplate" à faible entropie, tandis que quelques-uns représentent des points de décision critiques.

Question centrale : Faut-il vraiment tous les tokens pour entraîner de puissants raisonnateurs par RL ?

2. Méthodologie : Le cadre NAT

Les auteurs proposent NAT (Not All Tokens Are Needed), un cadre unifié qui rend le budget de tokens une primitive d'optimisation de premier ordre. L'idée est de mettre à jour la politique en utilisant uniquement un sous-ensemble sélectionné de tokens, tout en conservant le signal d'apprentissage complet de la séquence.

A. Principe Fondamental : Estimation de Horvitz-Thompson (HT)

Pour garantir que la mise à jour basée sur un sous-ensemble de tokens soit statistiquement correcte (non biaisée), NAT utilise l'estimation de Horvitz-Thompson :

Chaque token $t$ est associé à une probabilité d'inclusion $p_{i,t}$ .
Lors du calcul du gradient, les tokens sélectionnés sont repondérés par l'inverse de leur probabilité d'inclusion ($1/p_{i,t}$).
Théorème : Cette repondération garantit que l'estimateur du gradient est non biaisé par rapport à l'objectif RL complet, même si seuls certains tokens sont utilisés pour la rétropropagation.

B. Deux schémas de sélection de tokens

Le cadre NAT est implémenté via deux méthodes simples et interchangeables :

Échantillonnage Aléatoire Uniforme (URS - Uniform Random Sampling) :
- Chaque token a une probabilité constante $p$ d'être inclus.
- Avantage : Réduit le coût de rétropropagation (backward pass).
- Limite : Dans les Transformers causaux, le calcul avant (forward pass) doit tout de même traiter tous les tokens précédents pour calculer les probabilités, même si le token est masqué dans la perte. URS ne réduit donc pas significativement la mémoire d'activation ni le calcul avant.
Découpage Aléatoire de Préfixe (RPC - Random Prefix Cutting) :
- Au lieu de sélectionner des tokens individuellement, on sélectionne un préfixe contigu de la réponse (de la longueur $L_i$ ).
- La longueur du préfixe est échantillonnée aléatoirement selon une distribution (par exemple, uniforme).
- Avantage clé : Puisque le modèle ne traite que le préfixe, cela permet une véritable réduction du calcul avant (forward pass) et de la mémoire d'activation, en plus de réduire le calcul arrière.
- Différence avec la troncature déterministe : Contrairement à une troncature fixe (qui supprime toujours la fin de la séquence et introduit un biais systématique), RPC assure que chaque position a une probabilité non nulle d'être incluse, préservant l'absence de biais grâce à la correction HT.

3. Contributions Clés

Cadre unifié pour le RLVR efficace en tokens : Introduction de NAT, qui permet des mises à jour de politique avec des sous-ensembles masqués de tokens tout en évaluant les récompenses sur la réponse complète.
Preuve d'absence de biais (Unbiasedness) : Démonstration théorique que la repondération de Horvitz-Thompson fournit un estimateur non biaisé du gradient GRPO complet pour n'importe quelle probabilité d'inclusion positive.
Accélérations pratiques sans perte de qualité : Validation empirique montrant que NAT (surtout RPC) atteint des performances équivalentes au GRPO complet tout en réduisant le volume de rétropropagation de jusqu'à 50%.

4. Résultats Expérimentaux

Les expériences ont été menées sur des modèles Qwen2.5-Math-7B et Qwen3-8B avec l'algorithme GRPO, sur des benchmarks de raisonnement mathématique (MATH, AIME24, AIME25).

Performance (Précision)

RPC et URS : Les méthodes basées sur l'échantillonnage aléatoire (NAT) obtiennent des résultats équivalents au GRPO complet (les intervalles de confiance à 95 % se chevauchent) sur tous les benchmarks.
Troncature Déterministe (Det. Trunc.) : La suppression systématique de la moitié finale des tokens entraîne une baisse significative de la précision et une variance élevée, confirmant que la troncature fixe introduit un biais nuisible à l'apprentissage.

Efficacité Système (Mémoire et Temps)

Mémoire GPU (Peak Memory) :
- RPC réduit la mémoire GPU de pointe d'environ 18 % à 20 % (ex: de 47.72 Go à 39.23 Go pour Qwen3-8B).
- URS n'offre presque aucune réduction de mémoire (car le calcul avant reste complet).
- La troncature déterministe économise le plus de mémoire mais au détriment de la performance.
Temps d'entraînement :
- RPC réduit le temps de calcul (forward/backward, hors inférence) d'environ 29 % pour Qwen3-8B.
- Le temps total par étape (incluant l'inférence) est réduit d'environ 36 %.
- URS offre des gains minimes sur le temps total car il ne réduit pas le coût de l'inférence avant.

5. Signification et Impact

Changement de paradigme : NAT remet en cause l'hypothèse selon laquelle la rétropropagation complète sur toute la séquence est nécessaire pour un apprentissage efficace. Elle propose une approche "orthogonale" aux optimisations d'inférence (comme vLLM ou le speculative decoding).
Évolutivité du RL : En réduisant la pression mémoire et le coût calculatoire, NAT permet d'entraîner des modèles sur des trajectoires de raisonnement plus longues sans exploser les coûts matériels, ouvrant la voie à l'évolutivité des systèmes d'IA de pointe.
Régularisation structurelle : Le découpage aléatoire de préfixe (RPC) agit comme une forme de "dropout" au niveau de la séquence, préservant la structure causale tout en introduisant de la stochasticité bénéfique pour l'optimisation.
Applicabilité : La méthode est "plug-and-play" et peut être combinée avec n'importe quel pipeline d'inférence ou d'optimisation existant, offrant une voie directe pour réduire les coûts de formation des LLMs de raisonnement.

En conclusion, NAT démontre que l'on peut obtenir des gains de performance significatifs en termes de coût et de mémoire (jusqu'à 50 % de tokens utilisés pour la mise à jour) sans sacrifier la qualité du raisonnement, grâce à une correction statistique rigoureuse (Horvitz-Thompson).