r/MachineLearning · · 2 min read

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

Mirrored from r/MachineLearning for archival readability. Support the source by reading on the original site.

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! 😄

submitted by /u/AdRevolutionary3242
[link] [comments]

Discussion (0)

Sign in to join the discussion. Free account, 30 seconds — email code or GitHub.

Sign in →

No comments yet. Sign in and be the first to say something.

More from r/MachineLearning