Otimização de Raciocinadores de Cadeia de Pensamento por meio da Minimização da Variância do Gradiente em Amostragem por Rejeição e Aprendizado por Reforço

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

May 5, 2025

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

cs.AI

Resumo

O raciocínio em cadeia de pensamento (Chain-of-Thought, CoT) em modelos de linguagem de grande escala (Large Language Models, LLMs) pode ser formalizado como um problema de variável latente, onde o modelo precisa gerar etapas intermediárias de raciocínio. Embora abordagens anteriores, como o ajuste fino iterativo baseado em recompensas (iterative reward-ranked fine-tuning, RAFT), tenham se baseado em tais formulações, elas geralmente aplicam orçamentos de inferência uniformes em todos os prompts, o que não leva em consideração a variabilidade na dificuldade e no comportamento de convergência. Este trabalho identifica o principal gargalo no treinamento de CoT como a estimativa ineficiente do gradiente estocástico devido a estratégias de amostragem estáticas. Propomos o GVM-RAFT, uma Estratégia Dinâmica de Alocação de Amostras específica para prompts, projetada para minimizar a variância do gradiente estocástico sob uma restrição de orçamento computacional. O método aloca dinamicamente recursos computacionais monitorando as taxas de aceitação de prompts e as normas do gradiente estocástico, garantindo que a variância do gradiente resultante seja minimizada. Nossa análise teórica mostra que a estratégia de amostragem dinâmica proposta leva a garantias de convergência acelerada sob condições adequadas. Experimentos em raciocínio matemático mostram que o GVM-RAFT alcança uma aceleração de 2 a 4 vezes e melhorias consideráveis na precisão em relação ao RAFT padrão. A estratégia de amostragem dinâmica proposta é geral e pode ser incorporada a outros algoritmos de aprendizado por reforço, como o GRPO, levando a melhorias semelhantes na convergência e na precisão dos testes. Nosso código está disponível em 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.