atlas

Research map / Large language models / Language, safety and interpretability

I Don't Know: Explicit Modeling of Uncertainty with an [IDK] Token

NeurIPS 2024 · Large language models › Language, safety and interpretability · approach: Uncertainty quantification for LLMs · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Adds a special [IDK] token and a continued-pretraining objective that shifts probability mass from wrong next-token predictions to [IDK], improving factual precision at a small recall cost.

Problem

LLMs hallucinate factually incorrect text and current calibration methods do not let them express uncertainty reliably in closed-book factual completion.

Approach

Introduce a new [IDK] vocabulary token and modify cross-entropy during continued pretraining so that, when the gold token is not the argmax, part of the target mass is moved to [IDK], scaled by an adaptive uncertainty factor (1 - p(gold)/max p) capped by a hyperparameter Pi=0.5, plus a false-positive regularizer when the model is correct.

Key results

Measured effects reported in the paper

What was changedWhereMeasuredEffect
Apply IDK-tuning on The PileMistral-7B-v0.1 · LAMA Google-REF1: 50.0 -> 51.7Helped
Apply IDK-tuning on The PileMistral-7B-v0.1 · LAMA T-RexF1: 75.0 -> 75.3Helped
Apply IDK-tuning on The PileMistral-7B-v0.1 · LAMA SQuADF1: 54.4 -> 54.9Helped
Apply IDK-tuning on The PileMistral-7B-v0.1 · TriviaQAF1: 59.4 -> 60.6Helped
Apply IDK-tuning on The PileMistral-7B-v0.1 · PopQAF1: 33.7 -> 32.5Hurt
Apply IDK-tuning on The PileMistral-7B-v0.1 · PopQAPrecision: 40.3 -> 78.1Helped
Apply IDK-tuningMistral-7B-v0.1 · lm-eval-harnessPrecision: 45.0 -> 48.8Helped
Apply IDK-tuningMistral-7B-v0.1 · lm-eval-harnessRecall: 28.2 -> 20.8Hurt
Use adaptive rather than fixed λ (ablation)IDK-tuning with Mistral-7B-v0.1 · closed-book factual sentence completionIDK error rateHelped
Add the L_FP-reg regularization term (ablation)IDK-tuning with Mistral-7B-v0.1 · closed-book factual sentence completionIDK error rateHelped
Increase Π (ablation)IDK-tuning with Mistral-7B-v0.1 · closed-book factual sentence completionIDK recallHelped
Increase Π (ablation)IDK-tuning with Mistral-7B-v0.1 · closed-book factual sentence completionIDK error rateHurt

Surprising findings

Limitations

Experimental setup

Datasets: The Pile (uncopyrighted version) LAMA TriviaQA PopQA ARC HellaSwag MMLU TruthfulQA WinoGrande GSM8k Legal Plain English TLDR

Benchmarks: LAMA Google-RE LAMA T-Rex LAMA SQuAD TriviaQA PopQA ARC HellaSwag MMLU TruthfulQA WinoGrande GSM8k Legal Plain English

Metrics: precision recall F1 IDK recall IDK error rate RougeL

Baselines: Confidence Threshold P(True) Semantic Entropy The Pile continued training ablation

Builds on: cross-entropy language modeling P(True) semantic entropy

Closest papers by meaning

Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.

Open scime draft

Get the next scime Dispatch and a monthly field update

Free. We send a confirmation link first; every email has a one-click unsubscribe. Privacy