Online Causal Kalman Filtering for Stable and Effective Policy Optimization

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

Imaginez que vous essayez d'enseigner à un très grand et très intelligent robot (un modèle de langage) comment résoudre des problèmes de mathématiques complexes. Pour l'entraîner, vous lui donnez des réponses, vous vérifiez si elles sont bonnes, et vous lui dites : « Bravo, fais plus de ça ! » ou « Non, évite ça ». C'est ce qu'on appelle l'apprentissage par renforcement.

Le problème, c'est que parfois, le robot change d'avis trop vite et de manière chaotique. Il oscille entre « Je vais faire comme ça » et « Non, je vais faire comme ça » à chaque mot qu'il écrit. C'est comme si un chef cuisinier, en préparant un plat, décidait soudainement d'ajouter du sel, puis de l'enlever, puis d'en remettre, à chaque cuillère de sauce. Le résultat est une soupe immangeable et l'entraînement devient instable.

Voici comment les auteurs de cette paper (KPO) ont résolu ce problème, expliqué simplement :

1. Le Problème : Le « Bruit » des Mots

Dans les méthodes actuelles, le robot regarde chaque mot individuellement pour décider s'il doit le garder ou le modifier.

L'analogie : Imaginez que vous essayez d'écouter une conversation dans une pièce très bruyante. Si vous essayez de comprendre chaque syllabe séparément, vous entendez juste du bruit (des « tch » et des « ch ») et vous ne comprenez pas le sens de la phrase.
La réalité : Les chercheurs ont découvert que les « ratios d'importance » (un outil mathématique qui dit au robot à quel point il a changé d'avis) sont très bruyants. Ils sautent partout, de haut en bas, sans logique. Cela rend l'apprentissage fou et instable.

2. La Solution : Le Filtre Kalman (Le « Filtre à Bruit »)

Pour arranger ça, les auteurs ont utilisé une technique appelée Filtre de Kalman. C'est un outil mathématique utilisé depuis des décennies pour guider les fusées et les satellites.

L'analogie du GPS : Imaginez que vous conduisez une voiture avec un GPS.
- Parfois, le GPS fait une erreur de calcul et vous dit soudainement : « Tournez à gauche ! » alors que vous êtes sur une autoroute. C'est une erreur isolée (du bruit).
- Si vous suivez aveuglément ce GPS, vous allez vous écraser.
- Le Filtre de Kalman, lui, dit : « Attends, je sais que tu es sur une autoroute. Ce signal de « tourner à gauche » est probablement une erreur. Je vais lisser ce signal en me basant sur ta position précédente et ta trajectoire actuelle. »
- Il ne supprime pas le changement de direction si c'est réel, mais il ignore les petits sauts bizarres.

3. Comment ça marche pour le robot ?

Au lieu de regarder chaque mot isolément, le nouveau système (KPO) regarde le mot actuel en gardant à l'esprit les mots précédents.

Avant : Le robot voyait un mot, pensait « C'est super ! », puis le mot suivant, pensait « C'est nul ! », puis le suivant « Super ! ». C'était le chaos.
Avec KPO : Le robot dit : « Tiens, j'ai dit que ce mot était super. Le mot d'avant était aussi super. Donc, il est très probable que ce mot soit aussi super. Je vais ignorer le petit signal bizarre qui me dit le contraire. »
Le résultat : Au lieu de sauter partout, le robot suit une trajectoire fluide et cohérente. Il garde la structure logique de sa phrase tout en éliminant le bruit.

4. Les Résultats : Un Robot Plus Calme et Plus Intelligent

Les chercheurs ont testé cette méthode sur des problèmes de mathématiques très difficiles (comme les Olympiades de mathématiques).

Sans le filtre : Le robot apprenait vite, puis s'effondrait (il oubliait tout ou devenait bête).
Avec le filtre KPO : Le robot est devenu beaucoup plus stable. Il a appris à résoudre des problèmes complexes avec beaucoup plus de succès que les méthodes précédentes.

En résumé

Cette paper propose une astuce intelligente pour apprendre aux IA à ne pas paniquer à chaque petit changement. Au lieu de réagir à chaque mot comme s'il était isolé, elles utilisent un « filtre de mémoire » (le Filtre de Kalman) pour comprendre la tendance globale de la phrase.

C'est comme passer d'un conducteur qui tourne le volant à chaque mètre de route (et finit par sortir de la route) à un conducteur qui regarde la route devant lui, anticipe les virages et garde une trajectoire fluide. Résultat : le robot apprend mieux, plus vite, et ne s'effondre jamais.

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

1. Problématique : Instabilité de l'Optimisation par Renforcement (RL) pour les LLM

L'apprentissage par renforcement (RL) est devenu une méthode clé pour dépasser les limites du pré-entraînement des grands modèles de langage (LLM), notamment via des algorithmes comme GRPO (Group Relative Policy Optimization). Cependant, à grande échelle, ces méthodes souffrent d'une instabilité critique lors des mises à jour off-policy (hors politique).

Le cœur du problème : Les ratios d'échantillonnage d'importance (IS - Importance Sampling) au niveau des tokens présentent une variance élevée.
L'observation empirique clé : Les auteurs ont identifié que la déviation off-policy au niveau des tokens individuels est structurellement incohérente.
- Contrairement à l'intuition d'une homogénéité locale, les tokens adjacents dans une séquence changent fréquemment d'état (de on-policy à off-policy et vice-versa).
- Ces changements rapides et aléatoires (bruit haute fréquence) déforment les mises à jour du gradient de politique entre tokens voisins, entraînant souvent un effondrement de l'entraînement (training collapse) ou une variance excessive du gradient.
Limites des solutions actuelles :
- Les méthodes récentes (GSPO, GMPO) utilisent un ratio IS de niveau séquence (fixe pour tous les tokens d'une réponse). Cela lisse le bruit mais efface toute structure locale et hétérogénéité au sein de la séquence.
- D'autres méthodes ajustent les ratios token par token sans tenir compte de la structure temporelle, laissant subsister le bruit.

2. Méthodologie : KPO (Kalman Policy Optimization)

Pour résoudre ce problème, les auteurs proposent KPO, une méthode qui applique un filtre de Kalman causal aux ratios IS token par token. L'objectif est de lisser le bruit tout en préservant la cohérence structurelle locale.

A. Modélisation comme Série Temporelle

Les ratios IS ( $r_t$ ) sont traités comme une série temporelle observée ( $z_t = \log r_t$ ) évoluant au fur et à mesure de la génération des tokens. Le modèle suppose l'existence d'un état latent ( $\rho_t$ ) représentant le "vrai" ratio IS lissé, qui évolue de manière cohérente sur les tokens adjacents.

Le modèle d'espace d'état est défini par :

Équation d'état (Processus) : $\rho_t = \rho_{t-1} + \eta_t$ (marche aléatoire avec bruit de processus $Q$ ).
Équation d'observation : $z_t = \rho_t + \epsilon_t$ (observation bruitée avec bruit d'observation $V$ ).

B. Le Filtre de Kalman Causal

Contrairement aux filtres classiques qui peuvent utiliser le futur, KPO utilise une approche causale et autoregressive : l'estimation du ratio pour le token $t$ ne dépend que des tokens passés et du token courant. Le filtre alterne trois étapes :

Prédiction : Estimation de l'état latent et de sa variance basée sur l'historique ( $\hat{\rho}_{t|t-1}$ ).
Gain de Kalman ( $K_t$ ) : Calcul adaptatif qui pondère la confiance entre la prédiction temporelle et la nouvelle observation.
- Si le bruit d'observation est élevé ( $V$ grand), le filtre ignore l'observation et suit la prédiction (lissage fort).
- Si la prédiction est incertaine ( $P$ grand), le filtre s'adapte rapidement à l'observation.
Mise à jour : Correction de l'estimation latente en fonction de l'innovation (différence entre observation et prédiction).

C. Intégration dans l'Objectif de Politique

Le ratio filtré $\hat{r}_t = \exp(\hat{\rho}_{t|t})$ remplace le ratio brut dans la fonction objectif de GRPO/PPO. Cela permet de :

Lisser les pics de bruit (réduction de la variance).
Préserver la structure locale (les segments cohérents de déviation off-policy sont maintenus, contrairement à l'approche séquence unique).

3. Contributions Principales

Révélation Empirique : Identification et démonstration que la déviation off-policy au niveau des tokens est structurellement incohérente (changement fréquent d'état), ce qui est une cause majeure d'instabilité dans GRPO.
Nouvelle Méthode (KPO) : Introduction d'un filtrage de Kalman causal pour lisser les ratios IS token par token. C'est une approche "consciente de la structure" qui équilibre lissage du bruit et fidélité aux dynamiques réelles.
Résultats Empiriques Robustes : Démonstration que KPO surpasse les méthodes de l'état de l'art (GRPO, GSPO, GMPO) sur des tâches de raisonnement mathématique complexes, tout en assurant une stabilité d'entraînement supérieure.

4. Résultats Expérimentaux

Les expériences ont été menées sur le modèle Qwen3-4B avec six benchmarks de raisonnement mathématique (AIME'24/25, AMC'23, MATH500, OlympiadBench, Minerva).

Performance Globale : KPO (surtout la version avec clipping, KPO-clipped) obtient les meilleurs résultats sur la majorité des métriques (avg@16 et pass@16).
- Exemple : Sur AIME'24, KPO-clipped atteint 37.91 en avg@16 contre 32.70 pour GSPO et 27.29 pour GRPO.
- Sur les tâches les plus difficiles (AIME, Olympiades), l'amélioration est particulièrement marquée, soulignant l'importance d'une optimisation stable pour le raisonnement à long horizon.
Stabilité de l'Entraînement :
- GRPO diverge souvent après ~200 étapes et subit un effondrement de l'entropie (comportement déterministe prématuré).
- GSPO/GMPO sont stables mais saturent tôt (plateau de performance).
- KPO maintient une courbe de récompense ascendante stable, une entropie préservée (meilleure exploration) et une variance de perte de gradient faible.
Analyse des Ratios Filtrés :
- Après filtrage, la fréquence de changement d'état (switch frequency) chute de 0.43 à 0.01.
- La longueur des séquences continues (run-length) passe de 1.6 tokens à **120 tokens** pour les états directionnels.
- La variance locale et globale des ratios est réduite de plusieurs ordres de grandeur, confirmant la suppression du bruit haute fréquence.

5. Signification et Impact

Changement de Paradigme Temporel : L'article propose une nouvelle perspective temporelle sur les ratios IS, passant d'une vue statique (token par token ou séquence entière) à une vue dynamique et causale.
Efficacité pour le Raisonnement : En stabilisant les gradients locaux, KPO permet aux modèles d'explorer des chemins de raisonnement complexes sans s'effondrer, ce qui est crucial pour les tâches de type "Chain-of-Thought".
Légèreté et Compatibilité : La méthode est légère (ajout minimal de calcul) et compatible avec les pipelines RL existants (comme GRPO), ne nécessitant pas de changement d'architecture majeur.
Limites et Perspectives : Le filtrage de Kalman étant intrinsèquement séquentiel (causal), il est difficile à paralléliser comme la génération de tokens. Les auteurs suggèrent que le développement d'algorithmes de filtrage parallélisé serait une direction de recherche intéressante pour l'avenir.

En résumé, KPO résout un problème fondamental de stabilité dans l'entraînement RL des LLMs en introduisant une structure temporelle intelligente dans le traitement des ratios d'importance, permettant ainsi des gains significatifs sur des tâches de raisonnement de haut niveau.