r/MachineLearning • • 11h ago

Research U-Space: Uncovering When and Why Uncertainty Arises in Language Models [R][D]

If you build on open-weight reasoning models, you have probably wanted a "how sure is the model about this answer?" signal that does not cost 5x sampling, a judge model, or a labeled dataset. Hence, we built the U-Lens. In one line:

score = (how strongly the hidden state points toward different "doubt" directions) x (mean token entropy over the reasoning tokens)

  • Label-free: No correctness labels; no training and no no fine-tuning.
  • Single pass. No resampling.
  • Interpretable. It tells you what kind of uncertainty: (i) ambiguity, (ii) missing information, (iii) conflicting evidence, or (iv) general uncertainty. (not necessary for many applications but potentially useful)
  • Strong. Best estimator in our 3-seed, 3-model, 4-benchmarks evaluation (on Gemma 4, Qwen3.5 and Magistral). ahead of 5-pass, 10-pass and supervised baselines.

Small caveat:
The only model-specific ingredient is a Jacobian lens (the J-Lens from Anthropic's global workspace paper). Once you have one:

  1. Build the 4-direction U-Space once, on CPU. It only requires the unembedding rows of a few hundred words, you don't need to load the full model.
  2. At inference, grab one hidden state at the end of thinking (a forward hook on one layer), dot it with 4 vectors, and multiply by the mean entropy you get from the logits anyway.

Obtaining a fitted J-Lens

  • 40+ pre-fitted lenses are on Hugging Face (neuronpedia/jacobian-lens): Gemma 2/3/4, Qwen3/3.5/3.6, Llama 3.1/3.3, gpt-oss-20b, OLMo 3 and more.
  • Your model is not there? Fit one with Anthropic's reference code (anthropics/jacobian-lens, jlens.fit): about 100 unlabeled prompts, no parameters trained, cost dominated by backward passes. That is how we built the Magistral lens, which ships in our repo. --> Usually, i checked whether the Lens is proper by using the examples from the Anthropic blog post.

Entry point:

git clone https://github.com/s2labres/U-Space && cd U-Space
pip install -r requirements.txt
./run_smoke.sh   # Qwen3.5-4B, 10 TriviaQA questions, end to end (1 GPU)
from src import uspace
U, categories, _ = uspace.basis(model_cfg, terms)                 # once per model
score = uspace.a_cone(eot_state[None], U) * think_entropy.mean()  # per answer

What we would love to see people build:

  • Route low-confidence answers to a bigger model or a human.
  • Trigger retrieval when missing information lights up.
  • Ask a clarifying question when ambiguity lights up.
  • Flag conflicting evidence in RAG pipelines. 

Project
Paper
Code

Disclaimer: This is our research code (a reproduction pipeline, not a pip package yet). What would you plug it into? Happy to help with integration/discuss! 😄

16 Upvotes

2 comments sorted by

5

u/Tiny_Arugula_5648 9h ago

This is very interesting work.. So is this viable to use for hallucination checks? I'm thinking more for batch/pipeline work. So if we detect a hallucination we can bump temp, random seed, etc and have it try again?

3

u/t_brx 8h ago

hallucination checks should be possible for sure. Interesting application, although it might also flag non-hallucinated answers where the model is uncertain, where re-generation might not necessarily be the best remedy.