Optimalisatie van Chain-of-Thought Redeneerders via Minimalisatie van Gradiëntvariantie in Afwijzingssteekproeven en Reinforcement Learning

Optimizing Chain-of-Thought Reasoners via Gradient Variance Minimization in Rejection Sampling and RL

May 5, 2025

Auteurs: Jiarui Yao, Yifan Hao, Hanning Zhang, Hanze Dong, Wei Xiong, Nan Jiang, Tong Zhang

cs.AI

Samenvatting

Chain-of-thought (CoT) redeneren in grote taalmodellen (LLMs) kan worden geformaliseerd als een latent variabelenprobleem, waarbij het model tussenliggende redeneerstappen moet genereren. Hoewel eerdere benaderingen zoals iteratieve beloning-gestuurde fine-tuning (RAFT) op dergelijke formuleringen hebben vertrouwd, passen ze meestal uniforme inferentiebudgetten toe voor alle prompts, wat geen rekening houdt met variabiliteit in moeilijkheidsgraad en convergentiegedrag. Dit werk identificeert het belangrijkste knelpunt in CoT-training als inefficiënte schatting van stochastische gradienten door statische steekproefstrategieën. Wij stellen GVM-RAFT voor, een promptspecifieke Dynamische Steekproefallocatiestrategie die is ontworpen om de variantie van stochastische gradienten te minimaliseren onder een rekeningkundig budgetbeperking. De methode wijst rekenkundige middelen dynamisch toe door acceptatiepercentages van prompts en normen van stochastische gradienten te monitoren, waardoor wordt gegarandeerd dat de resulterende gradientvariantie wordt geminimaliseerd. Onze theoretische analyse toont aan dat de voorgestelde dynamische steekproefstrategie onder geschikte voorwaarden leidt tot versnelde convergentiegaranties. Experimenten op het gebied van wiskundig redeneren laten zien dat GVM-RAFT een 2-4x versnelling en aanzienlijke nauwkeurigheidsverbeteringen bereikt ten opzichte van standaard RAFT. De voorgestelde dynamische steekproefstrategie is algemeen en kan worden geïntegreerd in andere reinforcement learning-algoritmen, zoals GRPO, wat leidt tot vergelijkbare verbeteringen in convergentie en testnauwkeurigheid. Onze code is beschikbaar op https://github.com/RLHFlow/GVM.

English

Chain-of-thought (CoT) reasoning in large language models (LLMs) can be formalized as a latent variable problem, where the model needs to generate intermediate reasoning steps. While prior approaches such as iterative reward-ranked fine-tuning (RAFT) have relied on such formulations, they typically apply uniform inference budgets across prompts, which fails to account for variability in difficulty and convergence behavior. This work identifies the main bottleneck in CoT training as inefficient stochastic gradient estimation due to static sampling strategies. We propose GVM-RAFT, a prompt-specific Dynamic Sample Allocation Strategy designed to minimize stochastic gradient variance under a computational budget constraint. The method dynamically allocates computational resources by monitoring prompt acceptance rates and stochastic gradient norms, ensuring that the resulting gradient variance is minimized. Our theoretical analysis shows that the proposed dynamic sampling strategy leads to accelerated convergence guarantees under suitable conditions. Experiments on mathematical reasoning show that GVM-RAFT achieves a 2-4x speedup and considerable accuracy improvements over vanilla RAFT. The proposed dynamic sampling strategy is general and can be incorporated into other reinforcement learning algorithms, such as GRPO, leading to similar improvements in convergence and test accuracy. Our code is available at https://github.com/RLHFlow/GVM.