Placing Puzzle Pieces Where They Matter: A Question Augmentation Framework for Reinforcement Learning

2026-04-17

AI Summary

This work addresses the challenges in applying reinforcement learning to large language model (LLM) reasoning—namely, overfitting on easy problems and sparse rewards on hard ones—exacerbated by existing prompting methods that uniformly inject hints, often introducing redundancy or overlooking critical bottlenecks. To overcome this, the authors propose PieceHint, a framework that dynamically assesses problem difficulty and step-wise reasoning importance to inject task-critical hints on demand. Coupled with a progressive hint-removal strategy, PieceHint enables a smooth transition from assisted learning to autonomous reasoning. Integrating reinforcement learning, adaptive hint injection, and training annealing, the method achieves performance on par with 32B-parameter baselines across six mathematical reasoning benchmarks using only a 1.5B-parameter model, while consistently maintaining high pass@k diversity across all k values.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Knowledge Representation and Reasoning: Computational Complexity of Reasoning

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Reinforcement learning has become a powerful approach for enhancing large language model reasoning, but faces a fundamental dilemma: training on easy problems can cause overfitting and pass@k degradation, while training on hard problems often results in sparse rewards. Recent question augmentation methods address this by prepending partial solutions as hints. However, uniform hint provision may introduce redundant information while missing critical reasoning bottlenecks, and excessive hints can reduce reasoning diversity, causing pass@k degradation. We propose \textbf{PieceHint}, a hint injection framework that strategically identifies and provides critical reasoning steps during training. By scoring the importance of different reasoning steps, selectively allocating hints based on problem difficulty, and progressively withdrawing scaffolding, PieceHint enables models to transition from guided learning to independent reasoning. Experiments on six mathematical reasoning benchmarks show that our 1.5B model achieves comparable average performance to 32B baselines while preserving pass@k diversity across all $k$ values.

Problem

• reinforcement learning
• question augmentation
• reasoning bottlenecks
• pass@k degradation
• sparse rewards

Innovation

• question augmentation
• hint injection
• reasoning bottlenecks
• scaffolding withdrawal
• reinforcement learning