TreePO: Overbrugging van de kloof tussen beleidsoptimalisatie en effectiviteit, en inferentie-efficiëntie met heuristisch boomgebaseerd modelleren

Samenvatting

Recente vooruitgang in het afstemmen van grote taalmodellen via reinforcement learning heeft opmerkelijke verbeteringen geboekt in het oplossen van complexe redeneerproblemen, maar tegen de kosten van dure on-policy rollouts en beperkte verkenning van diverse redeneerpaden. In dit werk introduceren we TreePO, waarbij een zelfgeleid rollout-algoritme wordt gebruikt dat sequentiegeneratie ziet als een boomgestructureerd zoekproces. Bestaande uit een dynamisch boomsteekproefbeleid en decodering van vaste-lengte segmenten, benut TreePO lokale onzekerheid om extra vertakkingen te garanderen. Door berekeningen te amortiseren over gemeenschappelijke prefixen en vroegtijdig laagwaardige paden te snoeien, vermindert TreePO in essentie de rekenlast per update terwijl de verkenning diversiteit behouden of verbeterd blijft. Belangrijke bijdragen omvatten: (1) een segmentgewijs steekproefalgoritme dat de KV-cache last verlicht door middel van aaneengesloten segmenten en nieuwe vertakkingen genereert samen met een vroegtijdig-stopmechanisme; (2) een boomgebaseerde segmentniveau voordeelschatting die zowel globale als lokale proximale beleidsoptimalisatie in overweging neemt; en (3) analyse van de effectiviteit van dynamische divergentie en terugvalstrategie gedreven door waarschijnlijkheid en kwaliteit. We valideren empirisch de prestatieverbetering van TreePO op een set redeneerbenchmarks en de efficiëntiewinst van GPU-uren van 22\% tot 43\% van het steekproefontwerp voor de getrainde modellen, terwijl we een reductie van tot 40\% op trajectniveau en 35\% op tokenniveau steekproefberekening laten zien voor bestaande modellen. Terwijl TreePO een gratis efficiëntieverbetering biedt voor inferentie, toont het een praktisch pad naar het schalen van RL-gebaseerde na-training met minder steekproeven en minder rekenkracht. De homepage is te vinden op https://m-a-p.ai/TreePO.

English

Recent advancements in aligning large language models via reinforcement learning have achieved remarkable gains in solving complex reasoning problems, but at the cost of expensive on-policy rollouts and limited exploration of diverse reasoning paths. In this work, we introduce TreePO, involving a self-guided rollout algorithm that views sequence generation as a tree-structured searching process. Composed of dynamic tree sampling policy and fixed-length segment decoding, TreePO leverages local uncertainty to warrant additional branches. By amortizing computation across common prefixes and pruning low-value paths early, TreePO essentially reduces the per-update compute burden while preserving or enhancing exploration diversity. Key contributions include: (1) a segment-wise sampling algorithm that alleviates the KV cache burden through contiguous segments and spawns new branches along with an early-stop mechanism; (2) a tree-based segment-level advantage estimation that considers both global and local proximal policy optimization. and (3) analysis on the effectiveness of probability and quality-driven dynamic divergence and fallback strategy. We empirically validate the performance gain of TreePO on a set reasoning benchmarks and the efficiency saving of GPU hours from 22\% up to 43\% of the sampling design for the trained models, meanwhile showing up to 40\% reduction at trajectory-level and 35\% at token-level sampling compute for the existing models. While offering a free lunch of inference efficiency, TreePO reveals a practical path toward scaling RL-based post-training with fewer samples and less compute. Home page locates at https://m-a-p.ai/TreePO.

TreePO: Overbrugging van de kloof tussen beleidsoptimalisatie en effectiviteit, en inferentie-efficiëntie met heuristisch boomgebaseerd modelleren

TreePO: Bridging the Gap of Policy Optimization and Efficacy and Inference Efficiency with Heuristic Tree-based Modeling

Samenvatting

Support