# Understanding Token Probability Encoding in Output Embeddings

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 49 references that are in the map.
- First posted on arXiv 3 Jun 2024.

## Builds on (most cited first, 30 of 38)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [RealToxicityPrompts: Evaluating Neural Toxic Degeneration in Language Models](https://mlsanity.com/paper/2009.11462/) (2020)
- [Visualizing the Loss Landscape of Neural Nets](https://mlsanity.com/paper/1712.09913/) (2017)
- [SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot](https://mlsanity.com/paper/2301.00774/) (2023)
- [Aligning Books and Movies: Towards Story-like Visual Explanations by Watching Movies and Reading Books](https://mlsanity.com/paper/1506.06724/) (2015)
- [How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddings](https://mlsanity.com/paper/1909.00512/) (2019)
- [MAUVE: Measuring the Gap Between Neural Text and Human Text using Divergence Frontiers](https://mlsanity.com/paper/2102.01454/) (2021)
- [Editing Large Language Models: Problems, Methods, and Opportunities](https://mlsanity.com/paper/2305.13172/) (2023)
- [What is being transferred in transfer learning?](https://mlsanity.com/paper/2008.11687/) (2020)
- [Explainability for Large Language Models: A Survey](https://mlsanity.com/paper/2309.01029/) (2023)
- [Using the Output Embedding to Improve Language Models](https://mlsanity.com/paper/1608.05859/) (2016)
- [Representation Degeneration Problem in Training Natural Language Generation Models](https://mlsanity.com/paper/1907.12009/) (2019)
- [Real Time Image Saliency for Black Box Classifiers](https://mlsanity.com/paper/1705.07857/) (2017)
- [The large learning rate phase of deep learning: the catapult mechanism](https://mlsanity.com/paper/2003.02218/) (2020)
- [All-but-the-Top: Simple and Effective Postprocessing for Word Representations](https://mlsanity.com/paper/1702.01417/) (2017)
- [Compressing BERT: Studying the Effects of Weight Pruning on Transfer Learning](https://mlsanity.com/paper/2002.08307/) (2020)
- [All Bark and No Bite: Rogue Dimensions in Transformer Language Models Obscure Representational Quality](https://mlsanity.com/paper/2109.04404/) (2021)
- [The Early Phase of Neural Network Training](https://mlsanity.com/paper/2002.10365/) (2020)
- [BERT Busters: Outlier Dimensions that Disrupt Transformers](https://mlsanity.com/paper/2105.06990/) (2021)
- [Rethinking embedding coupling in pre-trained language models](https://mlsanity.com/paper/2010.12821/) (2020)

## Built on by (most cited first, 6 of 6)

- [Probabilistic Medical Predictions of Large Language Models](https://mlsanity.com/paper/2408.11316/) (2024)
- [Measuring Intrinsic Dimension of Token Embeddings](https://mlsanity.com/paper/2503.02142/) (2025)
- [Why Does Post-Training Quantization Work?](https://mlsanity.com/paper/2609.11716/) (2026)
- [On the Tip of the Tongue: Why LLMs Hallucinate Answers They Can Decode](https://mlsanity.com/paper/2603.13911/) (2026)
- [Reading Without a Reader: Large Language Models Collapse Reading and Writing into a Single Entangled Code](https://mlsanity.com/paper/2607.24797/) (2026)
- [The Geometry of Ignorance: LLMs Know When to Temper Bayesian Priors](https://mlsanity.com/paper/2609.02959/) (2026)

## Abstract (arXiv)

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.