Vive l'équilibre : optimisation de politiques basée sur les arbres pilotée par l'Information Bottleneck

Résumé

Les avancées récentes en apprentissage par renforcement en ligne (RL) pour les grands modèles de langage (LLMs) ont démontré des performances prometteuses dans des tâches de raisonnement complexes. Cependant, elles présentent souvent un compromis déséquilibré entre exploration et exploitation, entraînant une optimisation instable et des performances sous-optimales. Nous introduisons IB-Score, une nouvelle métrique fondée sur la théorie du goulot d'information (Information Bottleneck) qui évalue l'équilibre exploration-exploitation d'une politique en quantifiant le compromis entre la diversité de raisonnement au niveau des étapes et l'information mutuelle partagée avec la réponse correcte. L'analyse basée sur IB-Score montre que les approches populaires de RL en ligne (par exemple, GRPO) avec des régularisateurs courants ne parviennent pas à maintenir un équilibre constant durant l'entraînement, ce qui conduit à des résultats sous-optimaux. Pour remédier à cela, nous proposons IB-TPO (Information Bottleneck-driven Tree-based Policy Optimization), un cadre théorique qui formule IB-Score comme un objectif d'optimisation à grain fin et utilise une nouvelle stratégie d'échantillonnage par arbre guidée par IB. Cette stratégie améliore non seulement l'efficacité de l'échantillonnage en ligne avec 50 % de trajectoires supplémentaires sous le même budget de tokens, mais réutilise également la structure de l'arbre pour une estimation efficace de IB-Score par Monte Carlo. Des expériences approfondies sur des benchmarks standards montrent que notre méthode surpasse significativement la baseline GRPO de 2,9 % à 3,6 % et dépasse également d'autres approches de RL en ligne de pointe. Notre code est disponible à l'adresse https://github.com/alibaba/EfficientRL.

English

Recent advances in online reinforcement learning (RL) for large language models (LLMs) have demonstrated promising performance in complex reasoning tasks. However, they often exhibit an imbalanced exploration-exploitation trade-off, resulting in unstable optimization and sub-optimal performance. We introduce IB-Score, a novel metric grounded in Information Bottleneck theory that evaluates policy's exploration-exploitation balance by quantifying the trade-off between step-level reasoning diversity and mutual information shared with the correct answer. Analysis based on IB-Score shows that popular online RL approaches (e.g., GRPO) with common regularizers fail to consistently maintain balance during training with suboptimal results. To address this, we propose Information Bottleneck-driven Tree-based Policy Optimization (IB-TPO), a principled framework that formulates IB-Score as a fine-grained optimization objective and utilizes a novel IB-guided tree sampling strategy that not only improves the efficiency of online sampling with 50% more trajectories under the same token budget, but also reuses the tree structure for effective IB-Score Monte Carlo estimation. Extensive experiments across standard benchmarks show that our method significantly outperforms GRPO baseline by 2.9% to 3.6% and also outperforms other state-of-the-art online RL approaches. Our code is available at https://github.com/alibaba/EfficientRL.