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

AI summary

Overview

Research area: Machine learning — specifically test-time reinforcement learning (RL) for large language models (LLMs) applied to open-ended scientific and algorithmic discovery, with an emphasis on exploration.

Technical level: Intermediate. The paper is readable with a working knowledge of policy-gradient RL, but it is notation-heavy: it defines token-level states, actions, prediction residuals, epoch normalization, and a masked bonus, and readers need to track the distinction between task coefficients and intrinsic bonuses.

Scope: A single paper introducing CURIO, a framework that adds an Intrinsic Curiosity World Model to a test-time RL discovery loop, evaluated on six mathematical discovery tasks and single-cell denoising across Qwen3 backbones at 8B, 14B, and 235B scale.

What This Paper Is About

Open-ended discovery systems use an LLM to propose programs, evaluate them, and refine them, but search-based methods keep the model frozen and cannot learn from their own successes and failures on the problem at hand, while RL methods that do update the model tend to concentrate on currently high-reward attempts and abandon directions that look weak now but might pay off later. CURIO's goal is to let a policy keep learning from evaluator feedback while retaining an incentive to revisit underexplored directions, by adding an exploration signal learned from the policy's own hidden-state dynamics rather than from task reward alone.

Key Contributions

1. An Intrinsic Curiosity World Model (ICWM) for policy representations. The ICWM pairs a feature encoder with a forward predictor that learns token-level transitions in the policy's hidden-state representation, using hidden states and sampled-token embeddings. Its prediction residual becomes a token-local intrinsic signal; no ICWM gradient is propagated into the actor.

2. A selective, normalized, annealed curiosity bonus. The bonus is standardized using epoch-level mean and variance of masked residuals, retained only at sampled tokens outside the policy's top-k choices (k = 1 in every run), and added to a token-local task coefficient under a cosine-annealed weight with a floor of α = 0.1. Setting the initial curiosity weight η₀ = 0 recovers the matched task-only control.

3. A controlled empirical comparison against both frozen-search and trainable-RL baselines. CURIO is compared against AlphaEvolve (via OpenEvolve), ShinkaEvolve, Best-of-25600, SimpleTES, ThetaEvolve, and TTT-Discover across six mathematical objectives and two scRNA-seq denoising corpora, at three backbone scales, with three-run means.

4. A code-level diversity diagnostic. Using OpenEvolve's island code diversity metric, the authors track structural variation in generated programs across training epochs and report that CURIO generally maintains higher diversity than the task-only RL control, while explicitly noting this is not a semantic-diversity measurement.

Main Findings

• Improvement over the matched task-only RL control on five mathematical objectives. Relative to TTT-Discover (the same procedure with η₀ = 0), CURIO improves Erdős, AC1, AC2, AC3, and HadamardMatrix at each tested scale. Examples: Erdős at 8B is 0.380932 versus 0.380995; at 235B, 0.380880 versus 0.380918. HadamardMatrix at 235B is 0.9284 versus 0.7846.

• Circle Packing is matched, not beaten. On the n = 26 Circle Packing task (best-known ≈ 2.635983), CURIO and TTT-Discover both report 2.63598 at all three scales; the paper describes this as matching the best reported performance at the displayed precision.

• Gains are not monotonic in model scale. The Hadamard relative improvement over TTT-Discover is 4.11%, 3.36%, and 18.33% at 8B, 14B, and 235B respectively. The largest advantage appears at 235B, but the trend is not monotonic, and the authors state that attributing this to richer hidden representations would require additional controls.

• Denoising Score and MSE improve on both held-out corpora at every tested scale. Examples: 8B Tabula Score 0.7165 versus 0.6468; 235B Tabula Score 0.7408 versus 0.7254 and MSE 0.1334 versus 0.1401; 235B PBMC Score 0.7156 versus 0.7011. The abstract places the denoising Score relative gain at 10.8% on 8B Tabula.

• Not a uniform win over every baseline. On 8B Tabula, ThetaEvolve reports a higher Score (0.7207 versus 0.7165) and lower MSE (0.1441 versus 0.1485) than CURIO.

• Higher code-level structural diversity. CURIO shows higher code diversity D than TTT-Discover at every displayed stage on Erdős and Circle Packing — across the early (epochs 0–12), middle (13–25), late (26–38), and final (39–49) stages — and generally maintains higher D across the six mathematical tasks, although curves cross or approach one another at some stages.

• Curiosity strength is task-dependent. On Qwen3-14B, η₀ = 0.3 performs best on the autocorrelation tasks, while η₀ = 0.5 is preferable on Erdős, Hadamard, and aggregated single-cell; the largest tested weight, η₀ = 0.8, underperforms the task-only control on some tasks.

• Architecture and representation ablations are mixed. The default MLP forward predictor is competitive with GRU and Transformer alternatives, but the best architecture depends on the task; the policy-coupled representation leads on most plotted tasks, with exceptions on individual tasks for frozen-actor and frozen-external-LLM representations.

Methodology in Plain English

The setup is a loop that runs for 50 epochs. In each epoch, the system picks a small number of parent programs from an archive, generates 64 code completions per parent from the current policy at temperature 1.0, executes and scores them, updates the archive, and then takes one training step on the policy. The configured budget is therefore 25,600 generation slots. Only LoRA adapters (rank 32) are trained, and constant-reward groups are dropped because they carry no learning signal.

The task reward is turned into a per-rollout coefficient using exponential reward weighting, with a group-specific β calibrated to target a KL divergence of log 2 from uniform when attainable, plus a token-local correction toward a reference policy with coefficient λ = 0.1. That part is shared with the baseline, so any difference must come from the curiosity term.

The curiosity term works as follows. During generation, the system records the policy's last-layer hidden state before and after each sampled token, along with the token's embedding. A small encoder-predictor network is trained to predict the encoded next hidden state from the current encoded state plus the token embedding, and its prediction error measures how surprising that token was to the model's own learned model of its representation dynamics. Errors are cached before each update, standardized within the epoch, and then used as a bonus only at tokens the policy did not already rank in its top-k choices — that is, at points where generation departed from the most obvious continuation. The bonus is added to the task coefficient, weighted by an annealed coefficient that starts at η₀ and decays toward a floor, and one importance-weighted policy-gradient step follows. Setting η₀ = 0 turns the whole thing off and yields the control.

Evaluation uses three-run arithmetic means. Mathematical tasks report the best feasible objective found per run; denoising selects the epoch by held-in MSE and then reports Score, MSE, and Poisson NLL on both held-out corpora, with neither held-out corpus used for selection. SimpleTES is evaluated only at 8B and 14B, and the 235B comparison includes only TTT-Discover and CURIO.

Why This Matters

Impact on research. The paper targets a specific tension in test-time RL for discovery: reward-weighted updates are what make learning work, but they are also what collapse exploration. CURIO's proposal — derive the exploration signal from the policy's own representation dynamics rather than from reward, token entropy, or sampling temperature — is a general recipe that could be attached to any policy-gradient discovery loop. Just as important is the paper's framing of its own evidence: prediction error and code-level diversity are labeled as proxies that do not establish semantic novelty, and the auxiliary loss is explicitly noted to admit constant-feature solutions with no anti-collapse guarantee.

Real-world applications:

• Automated algorithm and construction search, such as packing problems, combinatorial matrices, and inequality discovery, where the best-known answer is a record rather than a fixed target.
• Scientific program synthesis where an evaluator exists but the space of viable strategies is wide, such as computational biology pipelines; the paper's own denoising task has the model evolve a denoising program for scRNA-seq data.
• AI R&D automation, where the same loop of propose-execute-evaluate applies to reward design, kernels, or agentic system construction (areas the related work cites as prior applications).
• Any setting where a model must adapt to a single hard problem at inference time rather than to a broad training distribution.

Industry relevance. The method trains only LoRA adapters, uses DeepSpeed for training and vLLM via Ray Data for generation, and checkpoints each of its five epoch stages to disk so runs are resumable at epoch granularity after preemption — a practical shape for serving-time adaptation. Because the curiosity head is small (a three-layer MLP in fp32) and the actor stays frozen between epochs, the added machinery is bounded relative to the cost of generation.

Future Directions

• Isolating the token-selection rule. The authors state directly that the comparisons do not independently validate the non-top-k masking heuristic, so an ablation separating the masking from the normalization and from the ICWM itself is the obvious missing experiment. The paper notes that high prediction error "does not establish semantic novelty," so connecting residuals to actual semantic branching is open.

• Guarding against encoder collapse. Joint encoder-predictor training under the stated loss admits constant-feature solutions, and the paper says it "does not provide an anti-collapse guarantee." Adding a mechanism that provably or empirically prevents trivial representations is a clear next step, alongside stability checks across random seeds, which the ablations explicitly do not establish.

• Explaining the scale trend. Why the Hadamard advantage is 4.11%, 3.36%, and 18.33% at 8B, 14B, and 235B rather than increasing smoothly is unexplained; the authors say attribution to richer hidden representations would require additional controls. Related open questions include whether the bonus scale should be matched to task-coefficient scale, since the paper notes normalization "does not ensure that their scale matches the task coefficients."

• Removing the per-task tuning burden. The initial curiosity weight is the only task-dependent hyperparameter, set to 0.5 for Erdős, Hadamard, and denoising and 0.3 for AC1–AC3 and Circle Packing, following the sweep in which η₀ = 0.8 underperformed the control on some tasks. Automating that choice, or making the method robust to it, would broaden applicability.

Target Audience

Researchers and practitioners working on LLM-based automated discovery, evolutionary program search, and test-time or continual RL for language models. It is most useful to readers who already understand policy-gradient methods and reward weighting and want to know how an intrinsic, prediction-error-based exploration signal can be grafted onto an existing discovery loop. Readers focused on statistical rigor should note that the paper reports three-run means but omits variability for space and states that mean differences alone do not establish statistical significance; readers interested in computational cost will find that wall-clock time, token counts, and training cost are not reported, and the paper explicitly warns that equal generation counts do not imply equal token counts, wall-clock time, or training cost.

Authors’ 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.