Reproduce “Token-Level Uncertainty-Aware Objective for Language Model Post-Training” on OpenResearch

About this paper

Tingkai Liu, Ari S. Benjamin, Anthony M. Zador

Researchers at Cold Spring Harbor Laboratory developed a post-training objective for language models that leverages token-level uncertainty to improve adaptation and generalization. This approach identifies tokens with high epistemic uncertainty via their predictive loss, focusing training on these informative tokens while using self-distillation for low-uncertainty tokens. The method demonstrates improved performance on both in-distribution and out-of-distribution tasks, for example, achieving 7.5% higher accuracy on IF-Eval for Gemma-2-2b.

Key insights

• Not all tokens in a sequence are equally important for learning; focusing training on high-uncertainty tokens can be more efficient.
• Predictive loss serves as an effective and computationally efficient proxy for token-level epistemic uncertainty, indicating where the model has knowledge gaps.
• A combined training objective, leveraging masked MLE on high-uncertainty tokens and self-distillation on low-uncertainty tokens, leads to improved generalization and naturally implements an automatic curriculum.

Summary by alphaXiv.