AutoScaffold boosts reinforcement learning by guiding language models
AutoScaffold: AI-Guided Scaffolding for Reinforcement Learning of Language Models
Machine LearningComputation and Language
Summary
Training language models with reinforcement learning typically reduces success to a single number, missing lots of helpful details. The authors created AutoScaffold, a system where a fixed teacher program watches the student model learn and adds hints or strategies during training to help it get better. These aids automatically adjust over time and are removed when testing the model. The approach helped models perform better on various tasks like question answering, math problems, and generating fast GPU kernels.
What this means in practice
- •For ai research engineers: Improve training efficiency and final performance of language models on tasks with verifiable rewards by automating training guidance.
- •For software developers: Generate optimized GPU kernels more reliably and faster than standard examples by training models with adaptive hints.$Commercial implications: Enables tools that produce high-performance code kernels for GPU computing, appealing to software vendors in high-performance computing.
Authors
Zijian Zhang, Rizhen Hu, Jiangweizhi Peng, Shiyang Li, Sujay Bhatt, Sihan Zeng, Mingyi Hong
Abstract
Reinforcement learning with verifiable rewards (RLVR) produces rich training experience, yet outcome rewards reduce each attempt to a single number and leave most of this information unused. Training-time scaffolds, such as strategies or hints added to the prompt during training, offer a natural way to turn this information into guidance, but their design and adaptation are largely left to humans and tailored to specific domains. We introduce AutoScaffold, a unified AI-for-AI framework that turns a student's training experience into adaptive scaffolds, with minimal domain-specific information and no human intervention during training. A frozen teacher periodically examines this experience, diagnoses bottlenecks, and designs scaffolds in two forms: (i) strategies shared across a class of tasks and (ii) hints for individual hard tasks, adjusting the content and amount of support as the student learns. The scaffolds are removed at evaluation. We apply the same framework to ALFWorld, search-based QA, mathematical reasoning, and GPU kernel generation. When evaluated without scaffolds, students trained with AutoScaffold outperform those trained with the corresponding RL baselines in all four domains, with gains of 7.2 points in ALFWorld success and 4.0 points in average QA exact match. On GPU kernel generation, AutoScaffold increases the fraction of tasks with a correct kernel that runs faster than the PyTorch reference by up to 5.0 points. On math, our student surpasses QuestA's reported accuracy using approximately one-third as many student training steps. Analyses show that the teacher identifies a different failure in each task class and shifts its scaffolds from step-by-step workflows to targeted fixes as the student improves, and that its revised scaffolds help the student learn more than keeping the previous scaffold or using no scaffold.