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

arXiv: https://arxiv.org/abs/2503.16511 · HTML: https://mlsanity.com/paper/2503.16511/

## 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.

## Placed in

- Knowledge distillation / [Knowledge distillation for LLMs](https://mlsanity.com/distill/knowledge-distillation-for-llms/)

## Builds on (most cited first, 17 of 17)

- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Self-Instruct: Aligning Language Models with Self-Generated Instructions](https://mlsanity.com/paper/2212.10560/) (2022)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Aleatoric and Epistemic Uncertainty in Machine Learning: An Introduction to Concepts and Methods](https://mlsanity.com/paper/1910.09457/) (2019)
- [Beyond neural scaling laws: beating power law scaling via data pruning](https://mlsanity.com/paper/2206.14486/) (2022)
- [BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learning](https://mlsanity.com/paper/1906.08158/) (2019)
- [Early-Learning Regularization Prevents Memorization of Noisy Labels](https://mlsanity.com/paper/2007.00151/) (2020)
- [Calibration of Pre-trained Transformers](https://mlsanity.com/paper/2003.07892/) (2020)
- [Learning with Noisy Labels Revisited: A Study Using Real-World Human Annotations](https://mlsanity.com/paper/2110.12088/) (2021)
- [Competence-based Curriculum Learning for Neural Machine Translation](https://mlsanity.com/paper/1903.09848/) (2019)
- [Fact-Checking the Output of Large Language Models via Token-Level Uncertainty Quantification](https://mlsanity.com/paper/2403.04696/) (2024)
- [Language Model Cascades: Token-level uncertainty and beyond](https://mlsanity.com/paper/2404.10136/) (2024)
- [High-Quality Prediction Intervals for Deep Learning: A Distribution-Free, Ensembled Approach](https://mlsanity.com/paper/1802.07167/) (2018)
- [Curriculum Learning by Transfer Learning: Theory and Experiments with Deep Networks](https://mlsanity.com/paper/1802.03796/) (2018)
- [An Empirical Exploration of Curriculum Learning for Neural Machine Translation](https://mlsanity.com/paper/1811.00739/) (2018)
- [On the Convergence of Learning-based Iterative Methods for Nonconvex Inverse Problems](https://mlsanity.com/paper/1808.05331/) (2018)

## 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.