Optimización de razonadores de cadena de pensamiento mediante la minimización de la varianza del gradiente en muestreo por rechazo y aprendizaje por refuerzo

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

Resumen

El razonamiento de cadena de pensamiento (CoT, por sus siglas en inglés) en modelos de lenguaje de gran escala (LLMs) puede formalizarse como un problema de variables latentes, donde el modelo necesita generar pasos intermedios de razonamiento. Aunque enfoques previos, como el ajuste fino iterativo basado en recompensas clasificadas (RAFT), han dependido de tales formulaciones, generalmente aplican presupuestos de inferencia uniformes en todas las indicaciones (prompts), lo que no tiene en cuenta la variabilidad en la dificultad y el comportamiento de convergencia. Este trabajo identifica el principal cuello de botella en el entrenamiento de CoT como la estimación ineficiente del gradiente estocástico debido a estrategias de muestreo estáticas. Proponemos GVM-RAFT, una Estrategia de Asignación Dinámica de Muestras específica para cada indicación, diseñada para minimizar la varianza del gradiente estocástico bajo una restricción de presupuesto computacional. El método asigna dinámicamente recursos computacionales monitoreando las tasas de aceptación de indicaciones y las normas del gradiente estocástico, asegurando que la varianza resultante del gradiente se minimice. Nuestro análisis teórico muestra que la estrategia de muestreo dinámico propuesta conduce a garantías de convergencia acelerada bajo condiciones adecuadas. Los experimentos en razonamiento matemático demuestran que GVM-RAFT logra una aceleración de 2 a 4 veces y mejoras considerables en precisión en comparación con RAFT estándar. La estrategia de muestreo dinámico propuesta es general y puede incorporarse en otros algoritmos de aprendizaje por refuerzo, como GRPO, lo que conduce a mejoras similares en convergencia y precisión en pruebas. Nuestro código está disponible en 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.