Ottimizzazione dei Ragionatori a Catena di Pensiero tramite Minimizzazione della Varianza del Gradiente nel Campionamento di Rifiuto e nell'Apprendimento per Rinforzo

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

May 5, 2025

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

cs.AI

Abstract

Il ragionamento a catena di pensiero (Chain-of-Thought, CoT) nei grandi modelli linguistici (Large Language Models, LLMs) può essere formalizzato come un problema di variabile latente, in cui il modello deve generare passaggi intermedi di ragionamento. Sebbene approcci precedenti come il fine-tuning iterativo con classificazione basata su ricompense (iterative reward-ranked fine-tuning, RAFT) si siano basati su tali formulazioni, essi applicano tipicamente budget di inferenza uniformi tra i prompt, senza tenere conto della variabilità nella difficoltà e nel comportamento di convergenza. Questo lavoro identifica il principale collo di bottiglia nell'addestramento CoT come una stima inefficiente del gradiente stocastico dovuta a strategie di campionamento statiche. Proponiamo GVM-RAFT, una strategia dinamica di allocazione dei campioni specifica per prompt, progettata per minimizzare la varianza del gradiente stocastico sotto un vincolo di budget computazionale. Il metodo assegna dinamicamente le risorse computazionali monitorando i tassi di accettazione dei prompt e le norme del gradiente stocastico, garantendo che la varianza del gradiente risultante sia minimizzata. La nostra analisi teorica mostra che la strategia di campionamento dinamico proposta porta a garanzie di convergenza accelerate in condizioni appropriate. Esperimenti sul ragionamento matematico dimostrano che GVM-RAFT ottiene un'accelerazione di 2-4 volte e miglioramenti significativi in termini di accuratezza rispetto al RAFT standard. La strategia di campionamento dinamico proposta è generale e può essere integrata in altri algoritmi di apprendimento per rinforzo, come GRPO, portando a miglioramenti simili nella convergenza e nell'accuratezza nei test. Il nostro codice è disponibile all'indirizzo 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.