Curiosity improves test-time learning for open-ended discovery tasks

CURIO: Curiosity-Driven Test-Time Learning for Open-Ended Discovery

Machine LearningComputation and Language

Summary

Solving open-ended problems means trying many solutions and learning from mistakes while still exploring new ideas, even if they don’t seem immediately promising. The authors created a method called CURIO that helps AI models stay curious by rewarding exploration of less obvious choices during problem solving. This curiosity encourages more diverse solutions and helps the model learn better on test problems without forgetting past successes. They tested CURIO on math problems and biological data tasks, showing improvements over usual methods that only focus on immediate rewards. This approach lets AI explore more broadly and adapt more effectively in complex tasks.

What this means in practice

  • For machine learning engineers: Enhance AI models’ adaptability on novel tasks by integrating curiosity signals to encourage exploration beyond immediate high-reward outcomes.
  • For biomedical data scientists: Improve denoising and restoring single-cell biological data by applying curiosity-driven learning to refine model predictions at test time.

Authors

Tao Feng, Fangxu Yu, Zijie Lei, Jiaru Zou, Changjiang Jiang, Yi Yan, Jiaxuan You, Pan Lu

Abstract

Open-ended discovery requires learning from repeated attempts while continuing to explore directions whose value is not yet apparent. Search with a frozen large language model (LLM) can reuse previous solutions in context, but cannot update the model from its successes and failures on the test problem. Reinforcement learning (RL) enables such adaptation; however, strongly favoring high-reward trajectories may suppress low-reward yet potentially promising directions too early. We introduce CURIO, a curiosity-driven test-time learning framework that complements task feedback with an Intrinsic Curiosity World Model (ICWM). The ICWM learns transitions in the policy's hidden-state representation and supplies prediction-error bonuses at sampled tokens outside the policy's top-k choices. Epoch normalization and an annealed weight regulate their contribution to the policy update. On six mathematical discovery tasks and single-cell denoising with Qwen3 backbones from 8B to 235B, three-run means improve over a matched task-only RL control on five mathematical objectives, match the best reported performance on Circle Packing, and improve denoising Score and mean squared error (MSE) on both held-out corpora at every tested scale. Relative gains reach 18.3% on Hadamard and 10.8% on denoising Score. Code-diversity measurements show greater structural variation among generated programs, supporting curiosity as a complementary exploration signal for learning in open-ended discovery.