Cog-DRIFT: La exploración de instancias reformuladas de forma adaptativa permite el aprendizaje a partir de problemas de razonamiento complejos

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

Resumen

El aprendizaje por refuerzo a partir de recompensas verificables (RLVR) ha mejorado las capacidades de razonamiento de los LLM, pero persiste una limitación fundamental: los modelos no pueden aprender de problemas que son demasiado difíciles de resolver bajo su política actual, ya que estos no generan una señal de recompensa significativa. Proponemos una solución simple pero efectiva basada en la reformulación de tareas. Transformamos problemas abiertos desafiantes en variantes cognitivamente más simples —como formatos de opción múltiple y cloze— que preservan la respuesta original mientras reducen el espacio de búsqueda efectivo y proporcionan señales de aprendizaje más densas. Estas reformulaciones abarcan un espectro que va desde tareas discriminativas hasta generativas, lo cual aprovechamos para impulsar el aprendizaje: los modelos primero aprenden de formatos estructurados y más fáciles, y este conocimiento se transfiere para mejorar el rendimiento en los problemas abiertos originales. Basándonos en esta idea, presentamos Cog-DRIFT, un marco que construye variantes reformuladas y las organiza en un currículo adaptativo basado en la dificultad. El entrenamiento progresa de formatos más fáciles a más difíciles, permitiendo que el modelo aprenda de problemas que previamente generaban señal cero bajo el post-entrenamiento RL estándar. Cog-DRIFT no solo mejora en los problemas duros originalmente irresolubles (+10.11% absoluto para Qwen y +8.64% para Llama), sino que también generaliza bien a otros conjuntos de datos retenidos. En 2 modelos y 6 benchmarks de razonamiento, nuestro método supera consistentemente al GRPO estándar y a líneas base fuertes de exploración guiada. En promedio, Cog-DRIFT muestra mejoras de +4.72% (Qwen) y +3.23% (Llama) sobre la segunda mejor línea base. Además, demostramos que Cog-DRIFT mejora pass@k en el momento de la prueba, y que el currículo mejora la eficiencia muestral. En general, nuestros resultados destacan la reformulación de tareas y el aprendizaje curricular como un paradigma efectivo para superar la barrera de exploración en el post-entrenamiento de LLM.

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.