Understanding Token Probability Encoding in Output Embeddings

TL;DR

This paper uncovers an approximate log-linear encoding of token probabilities in language model output embeddings, demonstrating its accuracy, sparsity, and implications for model efficiency and understanding pre-training dynamics.

Contribution

It reveals the sparse, log-linear structure of output token probabilities in embeddings and shows that many embedding dimensions are redundant, enabling potential model compression.

Findings

Output probabilities are encoded in a sparse, log-linear manner.

Over 30% of output embedding dimensions can be removed without affecting performance.

Output embeddings capture token frequency information early in training.

Abstract

In this paper, we investigate the output token probability information in the output embedding of language models. We find an approximate common log-linear encoding of output token probabilities within the output embedding vectors and empirically demonstrate that it is accurate and sparse. As a causality examination, we steer the encoding in output embedding to modify the output probability distribution accurately. Moreover, the sparsity we find in output probability encoding suggests that a large number of dimensions in the output embedding do not contribute to causal language modeling. Therefore, we attempt to delete the output-unrelated dimensions and find more than 30% of the dimensions can be deleted without significant movement in output distribution and sequence generation. Additionally, in the pre-training dynamics of language models, we find that the output embeddings capture the corpus token frequency information in early steps, even before an obvious convergence of parameters starts.