Cog-DRIFT: Exploração de Instâncias Reformuladas Adaptativamente Permite o Aprendizado a Partir de Problemas Complexos de Raciocínio

Cog-DRIFT: Exploration on Adaptively Reformulated Instances Enables Learning from Hard Reasoning Problems

April 6, 2026

Autores: Justin Chih-Yao Chen, Archiki Prasad, Zaid Khan, Joykirat Singh, Runchu Tian, Elias Stengel-Eskin, Mohit Bansal

cs.AI

Resumo

O aprendizado por reforço com recompensas verificáveis (RLVR) melhorou as capacidades de raciocínio de LLMs, mas uma limitação fundamental persiste: os modelos não conseguem aprender com problemas muito difíceis de resolver sob sua política atual, pois estes não produzem nenhum sinal de recompensa significativo. Propomos uma solução simples, porém eficaz, baseada na reformulação de tarefas. Transformamos problemas desafiadores de resposta aberta em variantes cognitivamente mais simples – como formatos de múltipla escolha e cloze – que preservam a resposta original, ao mesmo tempo que reduzem o espaço de busca efetivo e fornecem sinais de aprendizagem mais densos. Essas reformulações abrangem um espectro que vai de tarefas discriminativas a generativas, o qual exploramos para impulsionar o aprendizado: os modelos primeiro aprendem com formatos estruturados e mais fáceis, e esse conhecimento é transferido de volta para melhorar o desempenho nos problemas originais de resposta aberta. Com base nessa ideia, introduzimos o Cog-DRIFT, um framework que constrói variantes reformuladas e as organiza em um currículo adaptativo baseado na dificuldade. O treinamento progride de formatos mais fáceis para os mais difíceis, permitindo que o modelo aprenda com problemas que anteriormente produziam sinal zero sob o pós-treinamento padrão por RL. O Cog-DRIFT não só melhora os problemas difíceis originalmente insolúveis (ganhos absolutos de +10,11% para Qwen e +8,64% para Llama), mas também generaliza bem para outros conjuntos de dados retidos. Em 2 modelos e 6 benchmarks de raciocínio, nosso método supera consistentemente o GRPO padrão e baselines fortes de exploração guiada. Em média, o Cog-DRIFT mostra melhoras de +4,72% (Qwen) e +3,23% (Llama) sobre o segundo melhor baseline. Mostramos ainda que o Cog-DRIFT melhora a métrica pass@k no momento do teste, e que o currículo melhora a eficiência amostral. No geral, nossos resultados destacam a reformulação de tarefas e o aprendizado curricular como um paradigma eficaz para superar a barreira de exploração no pós-treinamento de LLMs.

English

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.