On the Predictive Power of Representation Dispersion in Language Models

TL;DR

This paper demonstrates that the dispersion of a language model's embeddings correlates with its predictive accuracy, and shows how this insight can be used for model evaluation, retrieval, and training improvements.

Contribution

It introduces the concept of representation dispersion as a predictor of model performance and practical utility, with methods to measure, leverage, and enhance dispersion.

Findings

Higher dispersion correlates with lower perplexity across models and domains.

Dispersion measurement helps rank example difficulty and identify hard data slices.

Increasing dispersion through training improves model perplexity.

Abstract

We show that a language model's ability to predict text is tightly linked to the breadth of its embedding space: models that spread their contextual representations more widely tend to achieve lower perplexity. Concretely, we find that representation dispersion--the average pairwise cosine distance among hidden vectors--strongly and negatively correlates with perplexity across diverse model families (LLaMA, Qwen, and others) and domains (Wikipedia, news, scientific abstracts). Beyond illustrating this link, we show how dispersion can be leveraged for a range of practical tasks--without requiring labeled data. First, measuring dispersion on unlabeled text allows us to rank examples by difficulty and identify hard slices in new domains, offering a data-efficient tool for screening and prioritizing models before full evaluation. Next, we find that identifying layers with higher dispersion pinpoints the best representations for retrieval-based methods such as kNN-LM, bypassing exhaustive layer-by-layer searches. Finally, we integrate a simple "push-away" objective into training, which increases dispersion in both single-domain and cross-domain scenarios and directly improves perplexity in each. Code is available at https://github.com/yanhong-lbh/rep_dispersion.