Stanford CS329A Self-Improving AI Agents | Part 6 | Train Time Scaling/Scaling RL

The lecture explored train time scaling and reinforcement learning techniques from three key papers—STaR, DeepSeekMath, and DAPO—that enhance AI agents’ mathematical reasoning by iteratively fine-tuning models using their own verified outputs, addressing challenges like training stability and output control. While these methods improve performance and consistency on known problem types, they do not fundamentally extend models’ abilities to solve novel problems, highlighting ongoing challenges in verification, reward design, and integrating diverse training approaches.

The sixth lecture of Stanford’s CS329A course focused on train time scaling and scaling reinforcement learning (RL) to improve AI agents’ reasoning capabilities. The lecture covered three key papers: STaR, DeepSeekMath, and DAPO, each addressing different aspects of enhancing model reasoning through train time scaling. The motivation stemmed from the need to improve performance on complex mathematical reasoning benchmarks like AIME, where smaller models trained with these techniques achieved surprisingly high accuracy compared to much larger models. The core idea behind train time scaling is to use a model’s own outputs, filtered for correctness, to fine-tune and improve the model iteratively, creating a self-improving feedback loop that complements pre-training and test-time scaling.

The first paper, STaR, introduced a simple yet effective approach to boost reasoning by generating rationales for problems and fine-tuning the model on only those outputs with correct answers. When the model failed to solve a problem, it was given the correct answer as a hint and asked to generate a rationale backward, which was then used for further training. This iterative process allowed the model to bootstrap its reasoning capabilities without requiring extensive human annotation of reasoning steps. While STaR showed promising results in improving reasoning efficiency and sample efficiency, it had limitations such as plateauing performance after several iterations and challenges in verifying the quality of generated rationales.

DeepSeekMath built upon these ideas by focusing on mathematical reasoning and introducing a novel RL algorithm called GRPO, which simplified the training process by reducing memory requirements and improving reward normalization. Starting from a code-pretrained model and curating high-quality math data from web sources, DeepSeekMath demonstrated significant improvements in math benchmarks by effectively scaling RL. The approach emphasized the importance of having a good base model and well-curated data before applying RL. It also addressed concerns about potential regression on simpler problems by incorporating mechanisms like KL divergence to prevent the model from deviating too far from previously learned capabilities.

The third paper, DAPO, tackled challenges encountered when scaling RL to larger models and harder problems, such as entropy collapse, unstable training, and uncontrollable output lengths. DAPO introduced several technical innovations, including asymmetric clipping in PPO to maintain exploration, dynamic sampling to ensure a meaningful distribution of rewards, and token-level loss penalties to control response length and stabilize training. These improvements allowed the Qwen-32B model to achieve state-of-the-art performance on the AIME benchmark, surpassing previous methods. The paper highlighted the complexity of optimizing RL loops for reasoning tasks and the necessity of carefully balancing exploration, reward shaping, and output control.

Throughout the lecture, the instructor emphasized that while train time scaling and RL techniques can significantly improve model consistency and reasoning on known problem types, they do not fundamentally expand the model’s ability to solve entirely new or out-of-domain problems. Verification signals and reward models are crucial for effective training, but they also present challenges such as potential reward hacking and noise. Open problems remain in learning from failures, evaluating rationale quality, and combining different training paradigms. The lecture concluded with a discussion on the practical implications of these methods for frontier models and the ongoing research directions aimed at making AI agents more self-improving and robust in complex reasoning tasks.