On the Significance of Softmax Geometry: Interpretability and Token Decoding

ICLR 2026 Conference Submission21505 Authors

Keywords: Embedding geometry, Softmax, Next-token prediction

TL;DR: This paper studies the geometry induced by softmax and demonstrates its practical benefits in interpretability and top-k prediction.

Abstract: Large language models represent their internal state as high-dimensional vectors. Many tasks of practical interest require measuring similarity between these vectors. Usually, this similarity is measured with a Euclidean notion. Recent work has argued that Euclidean geometry is ill-matched to semantic structure represented by LLMs. However, it's unclear whether this mismatch actually has practical consequences. In this paper, we study the practical effect of the similarity measure in the softmax layer of large language models (where the geometry is best understood). We consider two tasks: (1) learning interpretable features using sparse autoencoders, and (2) efficiently finding the most probable next tokens given a context. In both cases, we find that using the correct geometry dramatically improves the performance.

Primary Area: interpretability and explainable AI

Submission Number: 21505