Epistemic Uncertainty for Test-Time Discovery

Riaz, Kainat

Epistemic Uncertainty for Test-Time Discovery

Kainat Riaz^*, Muhammad Ahmed Mohsin^*, Ahsan Bilal, Muhammad Umer, Ayesha Mohsin, Aqib Riaz, Ali Subhan, John M. Cioffi

Stanford University · National University of Sciences and Technology · University of Oklahoma

^*Equal contribution (core contributors).

Paper Code

UG-TTT preserves the diversity that drives discovery. The training objective (top) augments the baseline clipped importance-sampling surrogate with a mutual-information (MI) shaped advantage and a nuclear-norm regulariser (purple terms). Left: per-epoch Shannon entropy H of the algorithmic family distribution, averaged over the four benchmarks (thin lines show individual tasks); UG-TTT (purple) sustains a cross-task mean H ≈ 1.42 bits while the baseline (gray) collapses to 0.65 bits by epoch 5, a 2.20× gap. Right: final-epoch entropy gains of +1.01 (AC1), +0.62 (AC2), +0.53 (CP26), and +0.94 (Erdős) bits, which directly enable the R_max improvements reported in the experiments.

Abstract

Automated scientific discovery using large language models relies on identifying genuinely novel solutions. Standard reinforcement learning penalises high-variance mutations, which leads the policy to prioritise familiar patterns. As a result, the maximum reward plateaus even as the average reward increases. Overcoming this limitation requires a signal that distinguishes unexplored regions from intrinsically difficult problems, which necessitates measuring disagreement across independently adapted weight hypotheses rather than relying on a single network's confidence.

UG-TTT addresses this challenge by maintaining a small ensemble of low-rank adapters over a frozen base model. The per-token disagreement, quantified as the mutual information between ensemble predictions and weight hypotheses, isolates epistemic uncertainty and identifies positions where insufficient coverage leads to adapter divergence rather than intrinsic problem difficulty. This measure is incorporated as an exploration bonus into the policy gradient, directing the policy toward positions where persistent adapter disagreement signals low training coverage, the same frontier where genuine discovery is possible. A nuclear norm regularizer ensures the adapters remain distinct from one another, thereby preserving the exploration signal throughout training.

Across four scientific discovery benchmarks, UG-TTT increases the maximum reward on three tasks, maintains substantially higher solution diversity, and an ablation study confirms that the regularizer is essential for sustaining exploration.

Method

UG-TTT augments test-time RL with three ingredients:

LoRA ensemble. A small set of low-rank adapters over a single frozen base model produces multiple weight hypotheses at near-zero memory overhead.
Mutual-information exploration bonus. Per-token disagreement between adapters, measured as the mutual information between ensemble predictions and weight hypotheses, isolates epistemic from aleatoric uncertainty and is added to the policy-gradient reward.
Nuclear norm regularizer. An explicit diversity term on the stacked LoRA weights prevents the adapters from collapsing into the same hypothesis, keeping the exploration signal alive across training.

Results

UG-TTT improves the maximum reward on three of four scientific discovery benchmarks and maintains substantially higher solution diversity throughout training. Below, per-epoch Shannon entropy over solution families shows that the baseline collapses to a single family while UG-TTT preserves a multi-modal mixture.

AC1 — **Figure 2. Training dynamics across all four benchmarks.** Per-epoch Shannon entropy H of the family distribution over correct rollouts (orange, left axis) and cumulative *R_max* (teal, right axis); dashed is baseline, solid is UG-TTT. The baseline contracts H monotonically on every problem while UG-TTT sustains higher H and matches or exceeds the baseline reward ceiling, so diversity preservation does not trade off against *R_max*.

AC2 — **Figure 2. Training dynamics across all four benchmarks.** Per-epoch Shannon entropy H of the family distribution over correct rollouts (orange, left axis) and cumulative *R_max* (teal, right axis); dashed is baseline, solid is UG-TTT. The baseline contracts H monotonically on every problem while UG-TTT sustains higher H and matches or exceeds the baseline reward ceiling, so diversity preservation does not trade off against *R_max*.

AC1 family composition — **Figure 3. Diversity collapse is universal and UG-TTT averts it.** Per-epoch family composition of correct rollouts on each benchmark; within each panel, the top row is the baseline (1 LoRA) and the bottom is UG-TTT (5 LoRA + MI + NNM), with *H_start → H_end* reported above. Family vocabularies are problem-specific and identical across conditions, so the contrast lies in which families remain populated. Terminal H = 1.71 / 1.12 / 1.23 / 1.62 bits under UG-TTT versus 0.70 / 0.50 / 0.70 / 0.67 under the baseline.

AC2 family composition — **Figure 3. Diversity collapse is universal and UG-TTT averts it.** Per-epoch family composition of correct rollouts on each benchmark; within each panel, the top row is the baseline (1 LoRA) and the bottom is UG-TTT (5 LoRA + MI + NNM), with *H_start → H_end* reported above. Family vocabularies are problem-specific and identical across conditions, so the contrast lies in which families remain populated. Terminal H = 1.71 / 1.12 / 1.23 / 1.62 bits under UG-TTT versus 0.70 / 0.50 / 0.70 / 0.67 under the baseline.

BibTeX

@inproceedings{riaz2026ugttt,
  title     = {Epistemic Uncertainty for Test-Time Discovery},
  author    = {Riaz, Kainat and Mohsin, Muhammad Ahmed and Bilal, Ahsan and Umer, Muhammad and Mohsin, Ayesha and Riaz, Aqib and Subhan, Ali and Cioffi, John M.},
  booktitle = {Submitted to the 40th Conference on Neural Information Processing Systems (NeurIPS 2026)},
  year      = {2026},
  url       = {https://kainatriaz98.github.io/uncertainty-guided-ttt/}
}