Calibrated Language Models and How to Find Them with Label Smoothing

TL;DR

This paper investigates how instruction tuning affects language model calibration, finds that label smoothing can improve calibration, and introduces a memory-efficient kernel for label smoothed training of large language models.

Contribution

It demonstrates the effectiveness of label smoothing in maintaining calibration during fine-tuning and proposes a memory-efficient implementation for large vocabulary models.

Findings

Calibration degrades after instruction tuning in open-source LLMs.

Label smoothing helps maintain calibration during fine-tuning.

A customized kernel reduces memory footprint without performance loss.

Abstract

Recent advances in natural language processing (NLP) have opened up greater opportunities to enable fine-tuned large language models (LLMs) to behave as more powerful interactive agents through improved instruction-following ability. However, understanding how this impacts confidence calibration for reliable model output has not been researched in full. In this work, we examine various open-sourced LLMs, identifying significant calibration degradation after instruction tuning in each. Seeking a practical solution, we look towards label smoothing, which has been shown as an effective method to regularize for overconfident predictions but has yet to be widely adopted in the supervised fine-tuning (SFT) of LLMs. We first provide insight as to why label smoothing is sufficient to maintain calibration throughout the SFT process. However, settings remain where the effectiveness of smoothing is severely diminished, in particular the case of large vocabulary LLMs (LV-LLMs). We posit the cause to stem from the ability to become over-confident, which has a direct relationship with the hidden size and vocabulary size, and justify this theoretically and experimentally. Finally, we address an outstanding issue regarding the memory footprint of the cross-entropy loss computation in the label smoothed loss setting, designing a customized kernel to dramatically reduce memory consumption without sacrificing speed or performance in comparison to existing solutions for non-smoothed losses.