Reproduce “Benefits and Pitfalls of Reinforcement Learning for Language Model Planning: A Theoretical Perspective” on OpenResearch

About this paper

Siwei Wang, Yifei Shen, Haoran Sun, Shi Feng, Shang-Hua Teng, Li Dong, Yaru Hao, Wei Chen

Researchers from Microsoft Research Asia, Peking University, Harvard University, and the University of Southern California present a theoretical framework examining reinforcement learning (RL) for Large Language Model (LLM) planning, comparing policy gradient (PG) and Q-learning against supervised fine-tuning (SFT). The study identifies SFT's memorization limitations, reveals a diversity collapse in PG which KL regularization can address, and demonstrates Q-learning's capacity to recover graph structures and preserve output diversity when utilizing appropriate process rewards.

Key insights

• Supervised Fine-Tuning (SFT) primarily memorizes co-occurrence patterns in training data, failing to infer transitive relationships crucial for robust planning and resulting in limited generalization.
• Policy Gradient (PG) benefits from on-policy data generation for exploration but inherently suffers from 'diversity collapse,' converging to a single correct path unless KL regularization is applied, which then trades off some accuracy gains for diversity.
• Q-learning, particularly with well-designed process rewards, effectively recovers underlying graph structures and preserves output diversity by converging to multiple feasible paths, and naturally supports efficient off-policy learning, unlike outcome rewards which lead to Q-value bias and model collapse.

Summary by alphaXiv.