Logit Distance Bounds Representational Similarity

For a broad family of discriminative models that includes autoregressive language models, identifiability results imply that if two models induce the same conditional distributions, then their internal representations are equal up to an invertible linear transformation. We ask whether an analogous conclusion holds approximately when the distributions are close instead of equal. Building on the observation of Nielsen et al. (2025) that closeness in KL divergence need not imply high linear representational similarity, we study a distributional distance based on logit differences and show that closeness in this distance does yield linear similarity guarantees. Specifically, we define a representational dissimilarity measure based on the models' identifiability class and prove that it is bounded by the logit distance. We further show that, when model probabilities are bounded away from zero, KL divergence upper-bounds logit distance; yet the resulting bound fails to provide nontrivial control in practice. As a consequence, KL-based distillation can match a teacher’s predictions while failing to preserve linear representational properties, such as linear-probe recoverability of human-interpretable concepts. In distillation experiments on synthetic and image datasets, logit-distance distillation yields students with higher linear representational similarity and better preservation of the teacher’s linearly recoverable concepts.

AI models are already being used for many tasks like writing code for software, drafting emails and analysing medical images. However, the models are still mostly black boxes, in the sense that we can seldom properly explain why a model gives a certain output. In this article, we explore the connection between a model's output (for example the probabilities a language model assigns to next words), and the internal representations of input and words, that is, the vectors which a model will make text into in order to do calculations. We find that different ways of measuring the similarity of outputs tells us different things about the representations. This means that if we want a model to have linearly similar representations to another model, then we need to use the right output similarity measure for training. This is especially relevant if we want to make a smaller model which both makes the same predictions as a larger model and also shares the same internal (linear) concepts.