What Models Know, How Well They Know It: Knowledge-Weighted Fine-Tuning for Learning When to Say "I Don't Know"

1. Problem Statement

Large Language Models (LLMs) frequently suffer from hallucinations, generating confident but incorrect answers. The authors identify a primary cause as knowledge misalignment between the model's pre-training (parametric knowledge) and the supervised fine-tuning (SFT) datasets.

• The Core Issue: When SFT datasets contain facts the model did not encounter or learn during pre-training, the model is forced to generate labeled tokens for facts it inherently "does not know."
• The Consequence: Instead of explicitly acknowledging uncertainty (saying "I don't know"), the model attempts to answer unanswerable queries, leading to hallucinations.
• Limitations of Prior Work: Existing methods often treat knowledge as binary (known/unknown) or rely on token-level refusal mechanisms that create a trade-off between accuracy and the ability to refuse answers.

2. Methodology: Knowledge-Weighted Fine-Tuning (KWT)

The proposed solution, KWT, aims to estimate the model's knowledge at a fine-grained, instance-level and use this estimate to modulate the training signal.

A. Knowledge Estimation (Pre-training Phase)

Before fine-tuning, the model's knowledge of each training instance is quantified using multi-sampled inference:

1. Multi-Sampling: For each question $q$, the pre-trained model ($M_{base}$) generates $S$ responses (e.g., $S=5$) using 3-shot prompting with resampled demonstrations.
2. Scoring: Each generated response is compared against the gold standard answer ($r^*$) using three metrics:
   • Exact Match (EM): Strict string matching.
   • ROUGE: Lexical overlap allowing for surface variation.
   • LLM-as-a-Judge: An LLM evaluates semantic correctness. (The paper finds this metric correlates best with human judgment).
3. Knowledge Score ($KS$): The final score is the average correctness across the $S$ samples:
   $$KS(q) = \frac{1}{S} \sum_{i=1}^{S} f(\hat{r}_i, r^*)$$
   This results in a continuous score between 0 (no knowledge) and 1 (high confidence knowledge).

B. Fine-Tuning Strategy

The training process uses these scores to weight the loss function and modify the target responses:

1. Weighted Loss: The training weight $w_i$ for the $i$-th sample is determined by its knowledge score. The authors propose a Familiarity (F) weighting strategy:
   $$w_i = \frac{S \times KS_i + 1}{S + 1}$$
   This increases the contribution of samples the model already knows well, reinforcing existing knowledge, while down-weighting unknown samples.
2. Explicit Uncertainty Token: For samples where $KS_i = 0$ (the model has no knowledge), a special token <IDK> ("I Don't Know") is appended to the target response. This trains the model to explicitly express uncertainty rather than hallucinating an answer.

3. Key Contributions

1. Fine-Grained Knowledge Estimation: Demonstrated that estimating knowledge as a continuous score via multi-sampled inference is more effective than binary classification for reducing hallucinations.
2. Knowledge-Weighted Training: Introduced a strategy that scales the learning signal based on the model's prior knowledge, allowing the model to focus on reinforcing known facts while learning to refuse unknown ones.
3. New Evaluation Metrics: Proposed nAUPC (Normalized Area Under the Product Curve), which jointly evaluates Uncertainty-Aware Accuracy (rewarding correct abstention) and Certainty-Aware Accuracy (penalizing unnecessary abstention). They also introduced A-FPR (Abstain False Positive Rate) and IDK Precision.

4. Experimental Results

The method was evaluated on diverse datasets (HaluEval, MedQA, SciQ) and compared against baselines like standard SFT, FT-TOP (filtering known data), R-Tuning, and SEAL.

• Uncertainty Expression: KWT (specifically using the LLM-as-a-judge metric) significantly outperformed baselines in nAUPC, achieving a better balance between answering correctly and refusing unknown queries.
  • Example: On SciQ, KWT (LLM) achieved an nAUPC of 38.1 compared to R-Tuning's 9.9.
  • It reduced A-FPR (unnecessary refusal of correct answers) and increased IDK Precision (correctly identifying when to say "I don't know").
• Accuracy Trade-off: KWT maintained accuracy comparable to standard SFT on answerable questions. While adding the <IDK> token slightly reduced standard accuracy compared to SFT, the Familiarity weighting alone (without the token) actually improved accuracy over SFT, suggesting the weighting strategy itself is beneficial.
• Generalization: The method generalized well to out-of-domain (OOD) datasets (RefuNQ, SelfAware), showing improved uncertainty handling on data not seen during training.
• Distribution Alignment: Analysis of KL divergence showed that KWT preserves the base model's predictive distribution better than SFT or SEAL, leading to more stable training.

5. Significance and Implications

• Mitigating Hallucinations: KWT provides a systematic way to reduce hallucinations by aligning the fine-tuning process with the model's actual internal knowledge state, rather than forcing it to learn facts it cannot reliably recall.
• Reliable AI: By enabling models to explicitly say "I don't know," KWT increases the reliability of LLMs in high-stakes domains (e.g., medical or scientific queries) where guessing is dangerous.
• Efficiency vs. Cost: The paper acknowledges a limitation: the multi-sampled inference required for knowledge estimation increases computational costs during the preparation phase. However, the authors argue this is a necessary overhead for achieving calibrated uncertainty.

In conclusion, the paper argues that knowing what a model knows is as critical as teaching it new information. By weighting training data based on the model's existing knowledge and explicitly training it to refuse unknown queries, KWT achieves a superior balance between accuracy and safety.