Viva o Equilíbrio: Otimização de Política Baseada em Árvores Orientada por Gargalo de Informação

Resumo

Avanços recentes em aprendizado por reforço online (RL) para grandes modelos de linguagem (LLMs) demonstraram desempenho promissor em tarefas complexas de raciocínio. No entanto, eles frequentemente exibem uma compensação desequilibrada entre exploração e explotação, resultando em otimização instável e desempenho subótimo. Apresentamos o IB-Score, uma nova métrica fundamentada na teoria do Gargalo da Informação que avalia o equilíbrio entre exploração e explotação da política, quantificando a compensação entre a diversidade de raciocínio em nível de passo e a informação mútua compartilhada com a resposta correta. Análises baseadas no IB-Score mostram que abordagens populares de RL online (por exemplo, GRPO) com regularizadores comuns não conseguem manter consistentemente o equilíbrio durante o treinamento, resultando em resultados subótimos. Para resolver isso, propomos o Information Bottleneck-driven Tree-based Policy Optimization (IB-TPO), uma estrutura fundamentada que formula o IB-Score como um objetivo de otimização refinado e utiliza uma nova estratégia de amostragem em árvore guiada por IB que não apenas melhora a eficiência da amostragem online com 50% mais trajetórias sob o mesmo orçamento de tokens, mas também reutiliza a estrutura da árvore para uma estimação Monte Carlo eficaz do IB-Score. Extensos experimentos em benchmarks padrão mostram que nosso método supera significativamente a linha de base GRPO em 2,9% a 3,6% e também supera outras abordagens de RL online de última geração. Nosso código está disponível em 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.