Оптимизация цепочек рассуждений через минимизацию дисперсии градиента в отбраковочной выборке и обучении с подкреплением

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

May 5, 2025

Авторы: Jiarui Yao, Yifan Hao, Hanning Zhang, Hanze Dong, Wei Xiong, Nan Jiang, Tong Zhang

cs.AI

Аннотация

Цепочка рассуждений (Chain-of-Thought, CoT) в больших языковых моделях (LLMs) может быть формализована как задача с латентными переменными, где модель должна генерировать промежуточные шаги рассуждений. Хотя предыдущие подходы, такие как итеративная тонкая настройка с ранжированием по вознаграждению (RAFT), опирались на такие формулировки, они обычно применяли равномерные вычислительные бюджеты для всех запросов, что не учитывало изменчивость сложности и поведения сходимости. В данной работе основным узким местом в обучении CoT идентифицируется неэффективная оценка стохастического градиента из-за статических стратегий выборки. Мы предлагаем GVM-RAFT — стратегию динамического распределения выборок, специфичную для каждого запроса, которая минимизирует дисперсию стохастического градиента при ограничении на вычислительный бюджет. Метод динамически распределяет вычислительные ресурсы, отслеживая показатели принятия запросов и нормы стохастических градиентов, что обеспечивает минимизацию дисперсии градиента. Наш теоретический анализ показывает, что предложенная стратегия динамической выборки приводит к ускоренным гарантиям сходимости при подходящих условиях. Эксперименты на задачах математического рассуждения демонстрируют, что GVM-RAFT достигает ускорения в 2-4 раза и значительного улучшения точности по сравнению с базовым RAFT. Предложенная стратегия динамической выборки является универсальной и может быть интегрирована в другие алгоритмы обучения с подкреплением, такие как GRPO, приводя к аналогичным улучшениям в сходимости и точности на тестах. Наш код доступен по адресу 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.