mlsanityData to 14 Oct 2026

Token-Level Uncertainty-Aware Objective for Language Model Post-Training

arXiv 2503.16511 ↗First posted 15 Mar 2025cs.CL cs.AIcited by 0

On the map: Knowledge distillation → Knowledge distillation for LLMs

Abstract (arXiv)

In the current work, we connect token-level uncertainty in causal language modeling to two types of training objectives: 1) masked maximum likelihood (MLE), 2) self-distillation. We show that masked MLE is effective in reducing epistemic uncertainty, and serve as an effective token-level automatic curriculum learning technique. However, masked MLE is prone to overfitting and requires self-distillation regularization to improve or maintain performance on out-of-distribution tasks. We demonstrate significant performance gain via the proposed training objective - combined masked MLE and self-distillation - across multiple architectures (Gemma, LLaMA, Phi) and datasets (Alpaca, ShareGPT, GSM8K), mitigating overfitting while maintaining adaptability during post-training. Our findings suggest that uncertainty-aware training provides an effective mechanism for enhancing language model training.

Key facts. No citing papers in the map yet (data through 14 Oct 2026). Placed in Knowledge distillation for LLMs, one of 12 directions in Knowledge distillation. Builds on 17 of its 33 references that are in the map. First posted on arXiv 15 Mar 2025.

Builds on

17 of 33 references.

References · 17

  1. 2015Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learningcited by 4,448
  2. 2016Simple and Scalable Predictive Uncertainty Estimation using Deep Ensemblescited by 3,926
  3. 2022Self-Instruct: Aligning Language Models with Self-Generated Instructionscited by 2,724
  4. 2017What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?cited by 2,352
  5. 2019Aleatoric and Epistemic Uncertainty in Machine Learning: An Introduction to Concepts and Methodscited by 913
  6. 2022Beyond neural scaling laws: beating power law scaling via data pruningcited by 528
  7. 2019BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learningcited by 432
  8. 2020Early-Learning Regularization Prevents Memorization of Noisy Labelscited by 416
  9. 2020Calibration of Pre-trained Transformerscited by 315
  10. 2021Learning with Noisy Labels Revisited: A Study Using Real-World Human Annotationscited by 256
  11. 2019Competence-based Curriculum Learning for Neural Machine Translationcited by 202
  12. 2024Fact-Checking the Output of Large Language Models via Token-Level Uncertainty Quantificationcited by 162
5 more
  1. 2024Language Model Cascades: Token-level uncertainty and beyondcited by 121
  2. 2018High-Quality Prediction Intervals for Deep Learning: A Distribution-Free, Ensembled Approachcited by 117
  3. 2018Curriculum Learning by Transfer Learning: Theory and Experiments with Deep Networkscited by 114
  4. 2018An Empirical Exploration of Curriculum Learning for Neural Machine Translationcited by 58
  5. 2018On the Convergence of Learning-based Iterative Methods for Nonconvex Inverse Problemscited by 11

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2026Distilling What Matters: Confidence-Aware Selective Distillation for Large Language Models
  2. 2023Uncertainty-aware Parameter-Efficient Self-training for Semi-supervised Language Understanding
  3. 2025Revisiting Uncertainty Estimation and Calibration of Large Language Models
  4. 2024Improving Instruction Following in Language Models through Proxy-Based Uncertainty Estimation
  5. 2026Semantic Self-Distillation for Language Model Uncertainty
  6. 2024Uncertainty Aware Learning for Language Model Alignment
  7. 2026Entropy-aware Masking for Masked Language Modeling
  8. 2025Dist2ill: Distributional Distillation for One-Pass Uncertainty Estimation in Large Language Models