Actor critic method improves large language model chain of thought reasoning
S3AC: Stabilized, Selective, and Scalable Actor-Critic for LLM Reasoning
Machine LearningComputation and Language
Summary
Long chains of reasoning in language models are hard to train because it's difficult to figure out how each step affects the final answer. The authors developed a new training method called S3AC, which uses a helper model (critic) to give better guidance for each word during training. They also created a way to pick easier or harder prompts based on past difficulty to save training effort. Their approach works well on very large models and beats other recent methods.
What this means in practice
- •For ai model trainers: Train large language models more efficiently for complex reasoning tasks using selective prompt sampling and token-level feedback.
- •For natural language processing engineers: Develop language model applications that need reliable and scalable chain-of-thought reasoning by integrating critic-guided reinforcement learning.
Authors
Shuo Shen, Guangran Cheng, Chengqi Lyu, Xiangsen Wang, Yanan Sun, Wenwei Zhang
Abstract
Critic-free methods are widely used in reinforcement learning with verifiable rewards (RLVR) for long chain-of-thought (CoT) reasoning, avoiding difficult and costly critic training. Common alternatives lack token-level credit assignment and typically rely on grouped rollouts to estimate prompt-specific baselines. A critic can address both limitations, but long-CoT critic training is hindered by insufficient coverage of trajectories with diverse lengths and early gradients dominated by the newly initialized value head. We introduce S3AC (Stabilized, Selective, and Scalable Actor-Critic) for LLM reasoning, built around a practical critic-training recipe. The trained critic provides token-specific advantages and prompt-specific baselines. Using the prompt-level baseline to estimate difficulty, we propose Difficulty-Aware Prompt Sampling (DAPS). DAPS combines this estimate with outcome history to select prompts before generation, reducing rollout waste without offline difficulty initialization. S3AC improves group-rollout performance across dense and mixture-of-experts models from 1.5B to 397B parameters, outperforming a modern GRPO variant (GRPO-X). At 35B, it sustains single-rollout training while tested critic-free baselines deteriorate after initial gains.