Stanford CS329A Self‑Improving AI Agents ｜ Part 6 ｜ Train Time Scaling⧸Scaling RL — Annotated Lecture Notes & Key Takeaways

Course Information

Institution: Stanford

Original Course: Stanford CS329A Self‑Improving AI Agents ｜ Part 6 ｜ Train Time Scaling⧸Scaling RL

Instructor Bio: This session is co-instructed by **Aakanksha Chowdhery** and **Azalia Mirhoseini**, faculty and researchers at Stanford University and Google DeepMind. Aakanksha Chowdhery earned her PhD from Stanford University and is a Senior Researcher at Google DeepMind. She led end-to-end training of the 540B PaLM model and drove pre-training and scaling efforts for multiple generations of Gemini MoE models. She has also contributed core components to PaLM-E, Med-PaLM, and the Pathways infrastructure that underpins Google’s large language model ecosystem, with deep expertise in dense and Mixture-of-Experts architectures and large-scale model training. Azalia Mirhoseini is an Assistant Professor of Computer Science at Stanford University and founder of the Scaling Intelligence Lab. She also serves as a Senior Staff Scientist at Google DeepMind and co-founded Ricursive Intelligence. Previously she held research roles at Google Brain and Anthropic, contributing to the development of the Claude and Gemini model families. Her pioneering work spans MoE architectures, test-time compute scaling, and deep reinforcement learning for chip design (AlphaChip), with a core research focus on recursive self-improving AI systems.

Course Description: This lecture focuses on training-time scaling and reinforcement learning as engines for agent self-improvement. It covers preference optimization methods including RLHF and DPO, scaling laws for reinforcement learning in agentic settings, and recursive self-improvement training techniques such as SWiRL. It also examines how agents can generate their own training data through synthetic reasoning traces, creating a closed loop of capability advancement.

This lecture focuses on train‑time scaling. It walks through three research papers, draws clear comparisons between train‑time and test‑time scaling, and explains how reinforcement learning can boost a model’s reasoning skills, most notably for mathematics problems. One key takeaway from the session: bigger model parameter counts are not the only path to better benchmark results. By dedicating compute resources during training and building on outputs the model itself generates, smaller‑size models can reach surprisingly strong performance numbers. The recording includes plenty of student questions that surface untested assumptions, confusing edge cases, and real‑world headaches when trying to put these algorithms into practice.
 

The big picture: Train‑time vs test‑time scaling

 

Before diving into individual papers, the lecturer lays out several established workflows for improving model capability. Pre‑training is the familiar foundation step, where models learn patterns from huge collections of internet text. Fine‑tuning covers chat‑focused development work, including reinforcement learning built on human preference feedback. Test‑time scaling techniques operate entirely at inference time.
 

Test‑time scaling covers things like sampling many different answers from a single model run, majority voting, or running verification checks against generated outputs. All the extra compute happens when you run the finished model, not while you train it. Course homework one even had students experiment with these exact inference‑side approaches.
 

Train‑time scaling works by taking data produced through that test‑time‑style inference workflow, filtering it carefully, and feeding high‑quality examples back into further fine‑tuning. This creates a closed self‑improvement loop. The lecturer keeps returning to this simple core idea: generate outputs from your model, validate those outputs, and retrain your model using the good examples you collect.
 

Students raised questions early on. One attendee wanted to know how to split limited compute budget between train‑time scaling and test‑time scaling, and whether performance gains eventually flatten out. The lecturer said they would circle back to this topic after covering the third paper. Another student asked whether fine‑tuning models to handle harder problems would make them worse at simpler tasks. This question was also tabled for later discussion, though the lecturer briefly mentioned “overthinking” as a known failure mode, where models produce long, broken reasoning chains.
 

One important detail came up repeatedly: these self‑improve loops deliver far better results in domains where outputs can be objectively verified. Math and coding fit this description perfectly; you can check whether a final answer is right or wrong. Tasks like creative writing or text editing see much smaller improvements, since there is no clear, objective way to label what counts as a good output.
 

The lecturer also talked about observable behaviors from modern reasoning‑focused models. These models break down problems, analyze what a question is actually asking, check their own intermediate work, backtrack when they spot mistakes, and test multiple possible solution paths. Chain‑of‑thought prompting draws these behaviors out, and the papers covered in class try to teach models to reliably produce these reasoning traces on their own.
 

Paper 1: STAR — bootstrapping reasoning from few starting examples

 

STAR is the first paper covered. Its main goal is improving model reasoning without building large datasets of human‑written reasoning steps.
 

The lecturer lays out the practical barriers that motivated this work. Raw internet text almost never contains complete step‑by‑step reasoning. Manually writing out reasoning traces costs significant time and labor. Automatically generating reasoning paths from known correct answers only works for narrow, limited use cases. Few‑shot prompting can coax models to show their working, but it still falls short of results you get when fine‑tuning on larger datasets.
 

STAR starts with a small set of examples that already include complete reasoning steps. Those examples serve as few‑shot prompts for the base model, which then generates solutions for a much larger pool of problems. The system keeps only outputs that end with a correct final answer, and uses those to build new fine‑tuning training data.
 

Lots of problems still stump the base model and produce wrong answers. For these failed attempts, STAR uses rationalization: feed the model the known correct answer as a hint, then ask it to write the reasoning path that arrives at that answer. These question‑reason‑answer triplets created with hints also get added to the fine‑tuning dataset. People repeat this whole process multiple times. With each iteration, the model gets slightly better and can solve more problems.
 

The lecturer points out three major assumptions baked into STAR’s design. First, a correct final answer works as a reasonable stand‑in for high‑quality intermediate reasoning. In practice, this means filtering out wrong answers under the assumption that incorrect final outputs usually come from flawed reasoning. Second, when given the correct answer as a hint, the model can produce valid reasoning steps. Third, your starting base model must already hold at least partial competence for the problem domain. If every single problem sits far beyond what the base model can handle, the bootstrap loop cannot deliver meaningful improvements.
 

Student questions poked at these assumptions during Q&A. One student brought up a real risk: you can end up with a correct final answer built entirely on flawed intermediate reasoning. The original STAR implementation does not filter for this problem. Follow‑up research can add process‑level reward models to inspect each reasoning step, though that sits outside the scope of the original paper. Another student asked what happens when the model cannot produce usable reasoning even after getting the correct answer as a hint. The lecturer responded that benchmark problems need to sit partially within the model’s existing skill range. Iterative improvement can slowly expand what the model can solve, but hard boundaries still exist.
 

In practical testing, STAR worked reasonably well on commonsense QA benchmarks. Human evaluators rated its generated reasoning as fairly solid, and it hit competitive accuracy numbers while using less training data than standard supervised fine‑tuning. Performance gains were much more muted on the GSM8K grade‑school math dataset. The lecturer’s take: when problems already sit mostly within the model’s existing ability, this bootstrap loop does not move performance numbers very much.
 

STAR has clear limits. Fundamentally it uses supervised fine‑tuning, not true reinforcement learning. Performance plateaus after a certain number of loop iterations. Evaluating reasoning quality remains difficult; you either need human reviewers or separate process‑oriented reward models. It cannot make huge logical jumps into completely new domains. It can only build on capabilities the base model already partially possesses.
 

I found one part a little tricky to follow. The difference between rationalizing failed samples versus learning only from successful model outputs matters a great deal, yet the lecture moved fairly quickly between these two workflows. A simple concrete example of a rationalized sample would have helped clarify things.
 

Paper 2: DeepSeek‑Math and the GRPO algorithm

 

Next the lecture moves to DeepSeek‑Math, which brings actual reinforcement learning into mathematical reasoning work. Discussion opens with results from the AME math benchmark. A model thought to match GPT‑3.5 scale only reached roughly 5% accuracy. Meanwhile, a 7‑billion‑parameter DeepSeek‑Math model using train‑time scaling hit 51.7% accuracy, and could approach 60% with extra adjustments. These numbers reinforce the lecture’s central idea: compute invested during training can offset sheer model parameter size.
 

The DeepSeek‑Math project split its work into two major phases. First, researchers built a strong starting base model. Older STEM‑focused research mined arXiv papers for training text. Instead, the DeepSeek team built Open Web Math, curating math‑related content pulled from Common Crawl web pages. They also started from a base model pre‑trained on code. Code‑based pre‑training delivered bigger gains for mathematical reasoning and tool‑use than arXiv‑sourced text material. This pre‑fine‑tune step gets the model ready before reinforcement learning runs.
 

Then they introduced GRPO, or Group Relative Policy Optimization. Traditional PPO, widely used in RLHF, requires maintaining multiple model copies: old policy, new policy, critic network, reward model. That creates heavy memory overhead, which gets worse as you scale up model size.
 

GRPO removes the critic value network entirely. For each question, you sample multiple different outputs from the model. Every sample receives a reward score. You compute advantage by normalizing each sample’s reward value: subtract the group’s average reward, then divide by the group’s reward standard deviation. This normalized value becomes your advantage estimate used for gradient updates.
 

This group‑baseline approach cuts down memory usage significantly. It works especially well for single‑step objective tasks like generating math answers. In benchmark testing, 7‑billion‑parameter models running GRPO crossed the 50% mark on math benchmarks, an important milestone for open‑weight models.
 

One student raised a practical concern: updating model weights for hard math problems might hurt performance on simpler questions. The lecturer explained a core requirement for GRPO to learn effectively. Your training batch needs diversity in reward values. If every sample gets maximum reward, or every sample gets zero reward, normalization gives you no usable learning signal. If you feed the model only problems it cannot solve at all, there is nothing for it to improve toward.
 

You can also keep the KL divergence penalties inherited from classic PPO. These penalties stop the updated policy from drifting too far away from the original base model, preserving capabilities the model already mastered.
 

One notable observation from DeepSeek‑Math results: the method improved majority‑at‑K performance. When drawing multiple samples for each question, most of those samples would become correct. Pass‑at‑K performance, which measures whether at least one sample out of many is correct, saw far less improvement. The lecturer frames this observation: the model grew more consistent at producing correct answers, but it did not unlock completely new fundamental problem‑solving ability.
 

Paper 3: DAPO — fixing instability for long‑chain reasoning

 

DAPO takes GRPO and addresses failure modes that pop up when you run reinforcement learning on harder math benchmarks like AME, where reasoning chains grow much longer. If you run vanilla GRPO on the Quen‑32B model, training becomes unstable. Common issues include collapsing generation entropy (the model loses exploration and grows over‑confident), output response lengths spiraling out of control, and noisy gradients created by reasoning sequences that get cut off when hitting context‑window limits.
 

DAPO introduces several targeted fixes.
 

First comes asymmetric clipping for PPO‑style updates. Standard symmetric clipping applies equal restrictions to both large increases and large decreases for token probabilities. DAPO loosens the upper bound, allowing larger positive updates while keeping stricter limits on probability drops. This preserves model exploration and keeps entropy from collapsing.
 

Second is dynamic sampling. You over‑generate many responses for each question. Then filter out groups where every sample is fully correct or every sample is completely wrong. Only batches containing a mix of high‑reward and low‑reward samples move forward to gradient calculation. This keeps useful reward variance intact and avoids wasting compute on groups that offer no learning signal.
 

Third, researchers switch from sample‑level loss to token‑level loss. Sample‑level loss assigns identical weight to an extremely long, garbage‑filled answer and a short, correct solution. Token‑level loss weights each individual token position, and includes soft penalties to manage runaway output length.
 

Fourth, they address truncated reasoning sequences. When a reasoning chain gets chopped off by the context window, it injects noise into model learning. DAPO applies gradual token‑wise penalties to reduce this noise.
 

Stacking these adjustments one after another on Quen‑32B pushed benchmark scores steadily higher. Starting at 30% accuracy with vanilla GRPO, each new component lifted performance, until the full DAPO setup hit 50% accuracy on AME. This outperformed results from a distilled DeepSeek‑R1 variant running on that same base model.
 

The lecturer stresses that monitoring reinforcement‑learning training means looking well beyond basic loss values. People working with these systems need to track three practical metrics: generation entropy, average response length, and the share of fully‑correct samples inside each batch. If response length blows up, entropy drops sharply, or nearly all samples are either all‑right or all‑wrong, your training loop runs into trouble. Reward‑model saturation is another possible failure state, when your reward signal stops producing meaningful differences between different outputs.
 

Open questions and unresolved gaps from the lecture

 

Below are research gaps the lecturer called out, plus points I found confusing listening to the talk.
 

All three papers covered deliver gains in majority‑at‑K performance. Yet pass‑at‑K often barely improves. Researchers still do not fully understand why this gap appears. One possibility: these methods only strengthen reasoning patterns the model already partially knows, rather than unlocking genuinely new reasoning capacity.
 

Model outputs show reasoning‑related behaviors like self‑correction and backtracking. Even so, people still debate whether these are truly emergent new skills, or just statistical shifts that make rare existing behaviors show up more frequently.
 

Learning directly from failed examples remains largely unsolved. Most existing workflows simply throw away bad samples. A small number of early papers explore learning from failure cases, but no mature standard approach exists.
 

Reward signals themselves create vulnerabilities. Reward hacking stays a real risk; models can find ways to game reward functions without improving actual reasoning quality. Working with small benchmark datasets such as AME or competition‑level IMO math also creates risk of data leakage. The lecturer notes reinforcement learning is data‑efficient, but you still need enough reward variance to enable hill‑climbing. Ensemble verifiers are one potential path forward, though the lecture did not dive into implementation details.
 

There are real practical trade‑offs between supervised fine‑tuning and reinforcement learning. SFT is simpler and faster when you have lots of high‑quality training data. RL requires careful tuning across many hyperparameters, but can deliver improvements with less raw training data — assuming you have solid reward or verification signals to work with.
 

Closing thoughts

 

Looking across STAR, DeepSeek‑Math (GRPO), and DAPO, each train‑time scaling approach fits different starting conditions. If you only have a small handful of reasoning examples and lack full reinforcement‑learning infrastructure, STAR‑style iterative supervised fine‑tuning makes a reasonable starting point. GRPO works well for standard math tasks on mid‑size models, with relatively low memory overhead. DAPO becomes relevant for harder problems requiring long reasoning chains, where you need active stabilization for your reinforcement‑learning loop.
 

Even so, none of these techniques produce magical jumps in out‑of‑domain generalization. They refine and amplify reasoning skills that the base model already partially holds. The lecturer mentions that in real‑world production systems, pre‑training on raw internet text still consumes the vast majority of total training compute. Self‑improvement work built on reinforcement learning takes up only a small fraction of total resources, though that fraction keeps growing over time.
 

The field of train‑time scaling is still very much active research territory. Many core questions have no settled answers, opening lots of possible directions for student projects.

Content Disclaimer: This article is for general reference only and does not constitute professional R&D guidance, production process advice or quality certification. All material performance data has specific test premises; readers should verify parameters against actual equipment and working conditions.