Source-linked AI summary
Cog-DRIFT: Exploration on Adaptively Reformulated Instances Enables Learning from Hard Reasoning Problems
Justin Chih-Yao Chen, Archiki Prasad, Zaid Khan, Joykirat Singh, Runchu Tian, Elias Stengel-Eskin, Mohit Bansal
TL;DR
RLVR struggles to learn from problems that the current policy cannot solve, because these problems provide little or no reward signal. Cog-DRIFT reformulates them into answer-preserving variants and trains across an adaptive difficulty curriculum. Across two models and six benchmarks, it improves performance over strong baselines, including +4.72% for Qwen and +3.23% for Llama versus the second-best baseline.
Problem
RLVR cannot reliably learn from problems with pass@64=0 because rollout samples usually receive zero reward, producing no meaningful advantage updates.
Method
Cog-DRIFT converts hard open-ended problems into answer-preserving multiple-choice and cloze variants organized into an adaptive curriculum by difficulty.
Results
Across two models and six challenging benchmarks, Cog-DRIFT consistently outperforms standard and guided-exploration baselines, with +4.72% average improvement for Qwen and +3.23% for Llama over the second-best baseline.
Takeaways & Limitations
Cog-DRIFT improves pass@1 and higher-pass@k performance while its instance-level curriculum remains more sample-efficient than static data mixtures.
Takeaways & Limitations
Using pass@k to estimate difficulty introduces selection bias because the hard set may include problems that are intrinsically unsolvable, not merely beyond the model’s capability.
Abstract
from arXiv · showhide
Reinforcement learning from verifiable rewards (RLVR) has improved the reasoning abilities of LLMs, yet a fundamental limitation remains: models cannot learn from problems that are too difficult to solve under their current policy, as these yield no meaningful reward signal. We propose a simple yet effective solution based on task reformulation. We transform challenging open-ended problems into cognitively simpler variants -- such as multiple-choice and cloze formats -- that preserve the original answer while reducing the effective search space and providing denser learning signals. These reformulations span a spectrum from discriminative to generative tasks, which we exploit to bootstrap learning: models first learn from structured, easier formats, and this knowledge transfers back to improve performance on the original open-ended problems. Building on this insight, we introduce Cog-DRIFT, a framework that constructs reformulated variants and organizes them into an adaptive curriculum based on difficulty. Training progresses from easier to harder formats, enabling the model to learn from problems that previously yielded zero signal under standard RL post-training. Cog-DRIFT not only improves on the originally unsolvable hard problems (absolute +10.11% for Qwen and +8.64% for Llama) but also generalizes well to other held-out datasets. Across 2 models and 6 reasoning benchmarks, our method consistently outperforms standard GRPO and strong guided-exploration baselines. On average, Cog-DRIFT shows +4.72% (Qwen) and +3.23% (Llama) improvements over the second-best baseline. We further show that Cog-DRIFT improves pass@k at test time, and the curriculum improves sample efficiency. Overall, our results highlight task reformulation and curriculum learning as an effective paradigm for overcoming the exploration barrier in LLM post-training.
1 Introduction
Cog-DRIFT addresses the RLVR exploration barrier by reformulating overly difficult open-ended problems into simpler structured tasks, then transferring learning back to the originals and held-out datasets.
- Motivation: RLVR can produce no meaningful updates when open-ended problems are too difficult for the current policy to solve.Under such conditions, on-policy sampling yields zero advantages and no effective learning signal.
- Task Reformulation: Task reformulation reduces cognitive load and output-space difficulty while preserving the original problem’s answer.The proposed variants include multiple-choice and cloze formats spanning discriminative to generative tasks.
- Adaptive Curriculum: Cog-DRIFT organizes reformulated variants into an adaptive difficulty curriculum analogous to cognitive scaffolding.The framework generates families of variants with different difficulty levels and orders training progression accordingly.
- Results: Across six benchmarks and two models, Cog-DRIFT improves over the second-best baseline by +4.72% for Qwen and +3.23% for Llama.The method also improves pass@k and generalizes to held-out tasks, whereas direct GRPO on hard problems can be unstable and harm generalization.
2 Related Work
Prior work frames hard-problem learning as an exploration challenge in RLVR and addresses it through verifiers, expert guidance, or policy-generated assistance.
- RLVR: RLVR trains reasoning models with objective, rule-based feedback and is especially effective for mathematics and code generation.Deterministic verifiers can provide binary rewards based on final-answer correctness.
- Learning from Hard Problems: Human-learning theory motivates matching task difficulty to a learner’s capabilities through structured progression rather than relying on human-like cognition in models.The paper explicitly uses “cognitive” to mean progression in task difficulty.
- Learning from Hard Problems: Existing approaches reduce difficulty by using stronger models for expert trajectories or the same policy for hints, critiques, and partial solutions.These strategies aim to help the model discover correct solutions when direct on-policy sampling fails.
3 Problem Statement
The paper defines hard problems by zero pass@64 performance and proposes answer-preserving reformulations that reduce difficulty while enabling transfer to the original open-ended tasks.
- Problem Definition: The dataset consists of open-ended problem-answer pairs, with each prompt paired to an exact answer and sampled through independent generations.The commonly used pass@k metric checks whether at least one of k generations is correct.
- Hard Problems: Hard problems are examples with pass@64=0, for which rollout-based RL generally receives only zero rewards and cannot form meaningful advantage updates.Direct GRPO on this subset yields little to no improvement because correct trajectories are rarely explored.
- Reformulation Criteria: An ideal reformulation preserves the answer, reduces empirical difficulty, and transfers learning gains back to the original problem.The transformation maps an open-ended pair (x, y) to a simplified variant (x′, y).
- Reformulation Formats: The considered variants include four-choice, ten-choice, and cloze formats, while requiring the model to output the full answer value rather than an option letter.These formats constrain the output space while retaining the original gold answer.
- Evaluating Difficulty Reduction: 4-choice reformulation reaches 33.8% accuracy versus a 25% random baseline, and empirical difficulty follows MCQ4 ≤ MCQ10 ≤ Cloze ≤ Open-ended.The original hard open-ended baseline is 0.0% because the selected problems have pass@64=0.
- Transferability: Training on reformulated problems transfers to the original hard questions, with accuracy rising from 11.1% for 4-choice to 18.9% for cloze.The format-specific models are also evaluated on unseen MATH500 to assess transfer beyond memorization.
4 Cognitive Demand Reduction via Incremental Format Transitions
Cog-DRIFT constructs verified reformulations ordered by increasing empirical difficulty and trains with an adaptive, instance-level curriculum. GRPO uses correctness and format-compliance rewards to advance examples when their estimated accuracy reaches a threshold.
- Reformulated task sequence: Cog-DRIFT orders reformulated variants from 4-choice MCQ through 10-choice MCQ and cloze to open-ended questions.All variants share the same ground-truth answer and are generated through the paper’s reformulation procedure.
- Instance-level Curriculum: Each training example starts at the easiest difficulty level and is promoted when its empirical accuracy reaches τ = 0.5.Otherwise, the example remains at its current level, creating a per-sample curriculum.
- Instance-level Curriculum: The curriculum adapts training to each problem’s evolving competence by focusing on the most informative difficulty level.Promotion depends on the model’s estimated accuracy for the current format.
- GRPO Training Objective: GRPO samples response groups, computes rule-based rewards, normalizes them within each group, and increases probability for relatively better samples.The reward combines answer correctness with format compliance, including whether the final answer appears in \boxed{}.
- GRPO Training Objective: The reward combines correctness in [0, 1] with format compliance in [0, 0.2].Format compliance checks whether the required final-answer format is followed.
5 Experimental Setup
The experiments evaluate Cog-DRIFT on two instruction-tuned models using a hard BigMath training subset and compare it with prompting, few-shot, and rejection-sampling baselines. Evaluation uses pass@1 with repeated runs across datasets.
- Models: Cog-DRIFT is evaluated on Llama3.2-3B-Instruct and Qwen3-4B-Instruct-2507.The models were selected because they exhibit different post-training behaviors.
- Dataset: The training data come from BigMath, which contains 251k samples with difficulty-related pass-rate information.The authors use this information to select harder training examples.
- Dataset: The hard subset begins with the hardest 20% of BigMath samples, totaling 37.5k examples, and retains problems without a successful sampled reasoning chain.Qwen samples 64 reasoning chains when filtering the hard subset.
- Baselines: Baselines include zero-shot prompting, few-shot prompting with eight in-context examples, and rejection-sampling fine-tuning using cloze reformulations.These comparisons avoid privileged information and stronger teacher models.
- Implementation Details: Evaluation uses pass@1, averaging 16 runs for AIME and three runs for other datasets.Math-Verify checks equivalence between predictions and references.
6 Results and Analysis
Cog-DRIFT unlocks learning from hard problems through reformulation, outperforming baselines across benchmarks and improving transfer, pass@k, and curriculum efficiency. Results also identify reformulation diversity, adaptive progression, and data-quality sensitivity as important factors.
- Baseline comparison: GRPO yields a +6.37% BMH gain with Qwen but no improvement with Llama, and it harms held-out generalization.The Qwen gain may reflect spurious rewards or occasional successful rollouts rather than transferable reasoning.
- Overall performance: Cog-DRIFT improves average accuracy by +4.72% for Qwen and +3.23% for Llama, while raising BMH from 0 to 10.11% and 8.64%, respectively.These gains are reported across six benchmarks, including problems collected with pass@64=0.
- Pass@k: At k = 128, Cog-DRIFT improves over the Qwen base model by +3.33% on AIME2024, +3.33% on AIME2025, and +2.02% on GPQA.Cog-DRIFT consistently achieves higher pass@k as k increases, outperforming both the base and GRPO-trained models.
- Reformulation diversity: Multiple reformulations outperform single-format training, with 4 Choice + Cloze achieving gains comparable to mixtures that also include open-ended problems.The findings identify diversity in problem representation, especially discriminative formats, as critical for transfer.
- Curriculum: The instance-level curriculum advances each instance after average accuracy exceeds 50%, producing continued test improvements while non-curriculum performance plateaus.Easier 4-choice MCQ samples decrease over training, while harder formats, including open-ended questions, become more prominent.
- Limitations and data quality: Pass@k-based difficulty estimation can introduce selection bias by including problems that are truly unsolvable for reasons beyond model capability.The hard dataset also contains incomplete, unanswerable, and incorrectly labeled samples, motivating data-quality filtering.
7 Conclusion
Cog-DRIFT reformulates hard open-ended questions into cognitively simpler variants that preserve the correct answer and tested knowledge. Across two models and six benchmarks, it improves reasoning performance, generalizes beyond training tasks, and benefits from instance-level curriculum learning.
- Cog-DRIFT reformulates hard open-ended questions into cognitively simplified variants that preserve the same correct answer and knowledge tested.
- Problem reformulation offers a practical strategy for reducing task difficulty and unlocking learning from previously intractable problems.
- Cog-DRIFT improves pass@1 performance across two models and six challenging benchmarks while generalizing effectively.
- Cog-DRIFT’s gains extend to higher pass@k values, suggesting learning from problems that naive approaches fail to exploit.
- Instance-level curriculum training continues improving performance, whereas static data mixtures quickly plateau.
A Implementation Details
The implementation uses fixed training and inference settings across GRPO-based methods and baselines. Detailed configurations and hyperparameters are reported in Table 4.
- Training uses a temperature of 1.0, clip-high parameter ϵ of 0.28, and eight rollouts for every GRPO-based method.
- Training and testing output length is capped at 8192 tokens, while inference temperature is fixed at 0.7.
- All methods run on 4 A6000 GPUs with aligned token limits and inference temperatures across evaluation baselines.
- Table 4 reports the detailed configurations and hyperparameters used for training.
B Dataset Statistics and Licenses
The study evaluates reasoning tasks across mathematics, science, and date arithmetic, using structured reformulation prompts that preserve the original question and correct answer. The appendix provides examples of four-choice, ten-choice, and cloze variants.
- The evaluation spans BigMathHard, OmniMATH-Hard, AIME 24/25, GPQA Diamond, and Date Understanding across STEM and domain-specific reasoning tasks.
- Dataset records: Table 5 reports the statistics and licenses of datasets used in the study.
- Reformulation prompts: Four-choice prompts retain the original question, add four options, include one correct answer, and generate plausible distractors.
- Reformulation prompts: Ten-choice prompts retain the original question, add ten options, include one correct answer, and generate nine plausible distractors.
- Reformulation prompts: Cloze prompts mask approximately 50–80% of answer digits while preserving at least one visible digit and LaTeX formatting.
- Reformulation examples: Examples show the same answer, 21, preserved across open-ended, four-choice, ten-choice, and cloze formulations.
E Example of Truly Unsolvable Problems
The hard-problem collection includes examples that are unsolvable because of incorrect gold answers, incomplete or ambiguous wording, or missing figures. These defects can prevent reliable reasoning from the supplied task.
- Some collected problems are described as truly unsolvable, motivating inspection of their data quality.
- Lack of Figures: A question that refers to an unavailable figure can omit information needed to compute the answer, despite providing a gold answer of 2150.
F Qualitative Examples of Baselines and Cog-DRIFT Outputs
The qualitative examples contrast a GRPO-generated solution with a Cog-DRIFT-trained solution, showing both the baseline’s detailed mathematical attempt and the trained model’s concise correct conclusion.
- GRPO baseline: GRPO’s example works through polynomial identities and root constraints before arriving at the correct minimum, 16.The solution rewrites the product as |P(i)|², constructs roots satisfying b − d = 5, and verifies the value.
- GRPO baseline: The baseline polynomial solution ultimately constructs the all-one root case, verifies b − d = 5, and obtains the final answer 16.It checks the coefficients of (x − 1)^4 and computes (1² + 1)^4 = 16.
- Cog-DRIFT: Cog-DRIFT’s example concludes that the only real matrix satisfying the stated condition is the zero matrix.Its reasoning derives that A^k is skew-symmetric, then uses the real-eigenvalue condition to conclude A^k = 0 and A = 0.