🤖 machine learning

Toward a First-Principles Update Geometry for the Language-Model Head

This paper proposes a first-principles update geometry for language-model heads by treating the softmax and weight matrix as a single module under Hilbert's projective distance, leading to an optimization strategy that maximizes the smallest token-row separation while constraining their Euclidean diameter.

Original authors: Aditya Somasundaram

Published 2026-08-25

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Aditya Somasundaram

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the vast landscape of modern artificial intelligence, the most visible achievements often come from systems that generate human-like text. Behind these conversations lies a complex machine learning model, a digital brain trained on enormous amounts of data. At the very end of this machine's processing chain sits a specific component known as the language-model head. Its job is deceptively simple: it takes a hidden internal state—a mathematical representation of what the computer is thinking—and converts it into a list of probabilities for every word in its vocabulary. If the model is trying to finish the sentence "The sky is," this component calculates how likely it is that the next word will be "blue," "cloudy," or "falling," and assigns a score to each possibility. For years, engineers have treated the mathematics that drive this final conversion as a standard problem, applying the same general rules to adjust the model's weights as they do for the earlier layers of the network.

However, a new perspective suggests that this final step is fundamentally different from the rest of the machine. Because the output is a probability distribution—a set of percentages that must add up to one hundred—the rules for changing it should respect the unique geometry of probabilities, not just the rules for changing raw numbers. A recent paper by Aditya Somasundaram from Columbia University explores this idea by treating the final layer and the probability conversion as a single, unified module. The author argues that to understand how to improve this part of the system, we must look at how the update changes the relative odds between words, rather than just how much the numbers shift. By using a specific mathematical distance that measures the change in these relative odds, the study derives a new way to measure the size of an update. The result is a fresh geometric picture where the "size" of a change is determined by the spread of the model's internal representations for different words, leading to a proposal for a new kind of optimizer that could make these models learn more efficiently.

The core of this investigation begins with a question about how we measure change. When a computer updates its internal settings to learn from a mistake, it makes a tiny adjustment to a massive table of numbers. In standard training methods, engineers often measure the magnitude of this adjustment by looking at the largest possible change it could cause in the raw numbers before they are converted to probabilities. But the author points out that for the final layer, the raw numbers are not the final product; the probabilities are. The conversion process, known as softmax, has a special property: if you add the same amount to every single number in the list, the resulting probabilities do not change at all. This means that measuring the raw size of the update is misleading, because it counts changes that have no effect on the final output. To solve this, the paper turns to a concept called Hilbert's projective distance. This is a way of measuring the distance between two sets of probabilities that focuses entirely on how the ratio between any two words changes. It ignores the absolute size of the numbers and looks only at the relative standing of one word against another.

By applying this specific distance measure, the researcher discovered a surprising and clean relationship between the geometry of the update and the behavior of the model. The study shows that the maximum change an update can cause in the probability distribution is directly tied to the physical spread of the rows in the update matrix. Imagine the update as a collection of vectors, where each vector corresponds to a specific word in the vocabulary. The "size" of the update, in terms of how much it can shake the probabilities, is determined by the distance between the two vectors that are farthest apart. If the vectors for different words are clustered tightly together, the update is small and safe. If they are spread far apart, the update is large and could cause wild swings in the model's predictions. This finding reframes the problem of updating the language-model head: instead of worrying about the overall magnitude of the numbers, the goal becomes managing the diameter of the cloud of points representing the words.

This geometric insight leads to a new proposal for how to construct these updates, drawing inspiration from a recent method called Muon that has been successful in other parts of neural networks. Muon works by balancing the different directions of an update so that they are all equally strong, preventing the model from getting stuck in narrow valleys of the learning landscape. The author suggests that a similar principle should apply to the language-model head, but with a twist. Instead of balancing the strength of singular directions, the goal should be to make the distances between all pairs of word-vectors as equal as possible. The ideal update would spread the vectors out so that every word is roughly the same distance from every other word, creating a perfectly balanced cloud. This would ensure that the model treats the distinction between any pair of words with equal sensitivity.

However, the paper also identifies a hard physical limit to this ideal. In the language-model head, the number of words in the vocabulary is vastly larger than the number of dimensions available to represent them. It is mathematically impossible to arrange a huge number of points in a small space so that they are all exactly the same distance from each other. Just as you cannot fit a hundred points on a flat sheet of paper so that they are all equidistant, you cannot make the word-vectors perfectly equidistant when the vocabulary is huge and the hidden space is small. The study acknowledges this constraint and suggests that the goal should be to find the best possible approximation. The proposed optimizer would try to maximize the smallest distance between any two words while keeping the largest distance within a safe limit. This approach aims to create an update that is as evenly distributed as the geometry allows, ensuring that no pair of words is ignored or treated as indistinguishable while others are pushed too far apart.

The implications of this work are primarily theoretical and architectural, offering a new lens through which to view the final step of language generation. The author does not claim to have built a fully trained model that proves this method works better in practice, but rather provides a rigorous derivation of what the update geometry should look like if one follows the first principles of probability. The paper argues that current methods, which often leave this final layer to standard optimization techniques, are missing an opportunity to respect the specific function of the module. By treating the conversion to probabilities as an integral part of the update process, and by using a distance measure that respects the invariance of the softmax function, the proposed geometry offers a more natural way to navigate the learning landscape. The study concludes that while the perfect equidistant arrangement is impossible for large vocabularies, the pursuit of an approximately equidistant configuration provides a clear, principled direction for designing future optimizers that are tailored specifically to the unique demands of the language-model head.

Technical Summary: Toward a First-Principles Update Geometry for the Language-Model Head

Problem Statement
Recent advancements in optimization, such as the Muon algorithm, have shifted focus from coordinate-wise updates to geometry-aware updates that respect the functional role of parameter blocks. While Muon successfully applies spectral geometry (based on the spectral norm) to hidden linear layers, its direct application to the Language-Model (LM) head is suboptimal. The LM head maps hidden states to probability distributions via a softmax function, making the relevant output a distribution rather than a Euclidean activation. Furthermore, the LM head is typically "tall" ( $V \gg d$ , where $V$ is vocabulary size and $d$ is hidden dimension), meaning an update matrix has at most $d$ non-zero singular directions, which fails to capture the geometric relationships between the $V$ token rows. Current implementations often leave the LM head to standard AdamW optimization, lacking a principled update geometry derived from the module's specific function.

Methodology
The paper proposes deriving an update geometry for the LM head by treating the linear layer and the softmax function as a single composite module. The methodology proceeds through the following steps:

Selection of Distance Metric: The author selects Hilbert's projective distance ( $H$ ) to measure changes in the probability distribution. Unlike Euclidean distances, Hilbert distance measures the largest change in pairwise log-odds between tokens. It is invariant to common shifts in logits, aligning with the invariance of the softmax function ( $\text{softmax}(z + c\mathbf{1}) = \text{softmax}(z)$ ).
Derivation of Functional Size: By analyzing the composition of the linear update $S$ and the softmax, the paper derives the maximum change in the probability distribution for a hidden state bounded by $\|h\|_2 \le H$ . The result shows that the functional size of an update $S$ is determined by the Euclidean diameter of its row cloud:
$\sup_{\|h\|_2 \le H} H(p_{U+S}(h), p_U(h)) = H \max_{i<j} \|s_i - s_j\|_2$
Here, $s_i$ and $s_j$ are the rows of the update matrix $S$ corresponding to tokens $i$ and $j$ . This establishes that the "size" of an update is not its spectral norm, but the maximum Euclidean distance between any pair of its rows.
Backward Propagation Analysis: The paper demonstrates that this row diameter constraint also bounds the backward gradient movement. Specifically, the change in the backward map is bounded by the row diameter scaled by the $L_1$ norm of the gradient, ensuring that controlling the row diameter simultaneously controls both forward probability movement and backward perturbation.
Optimization Formulation: The author formulates the update problem as a steepest descent optimization. The goal is to minimize the first-order loss change $\langle G, S \rangle_F$ (where $G$ is the gradient) subject to a constraint on the row diameter $D(S) \le \eta$ and a centering constraint ( $\mathbf{1}^\top S = 0$ ) to remove functionally invisible shifts.
Conditioning Analogue: Drawing an analogy to Muon's singular value conditioning (which equalizes singular values), the paper proposes an analogous goal for the LM head: maximizing the minimum pairwise row separation ( $\delta(S) = \min_{i<j} \|s_i - s_j\|_2$ ) while constraining the maximum separation ( $D(S)$ ). This aims to create an "approximately equidistant" update cloud.

Key Contributions

Functional Geometry Derivation: The paper derives that the natural geometry for LM head updates is the Euclidean diameter of the token-indexed rows, rather than the spectral norm of the matrix.
Hilbert Distance Connection: It establishes a direct link between Hilbert's projective distance on probability distributions and the variation seminorm of logit differences, showing that controlling Hilbert distance is equivalent to bounding the maximum pairwise log-odds change.
Approximate Equidistance Problem: For the regime where $V \gg d$ , the paper identifies that exact equidistance (a regular simplex) is geometrically impossible. It reframes the conditioning problem as finding the best approximate equidistant packing of rows in $\mathbb{R}^d$ .
Optimizer Proposal: It proposes a new optimizer framework that maximizes the ratio $\rho(S) = \delta(S)/D(S)$ (the condition number of the row cloud) to ensure uniform token separation, analogous to Muon's spectral conditioning.

Results and Claims
The paper does not present empirical experimental results or benchmarks. Instead, its "results" are theoretical derivations:

It proves that under Euclidean-bounded hidden states, the functional size of an LM head update is exactly the Euclidean diameter of its rows.
It proves that this geometry naturally bounds the Cross-Entropy loss change for all target tokens simultaneously.
It demonstrates that the proposed geometry leads to a well-defined steepest descent problem with a diameter constraint.

Significance and Future Work
The paper claims significance by extending the "first-principles" philosophy of Muon to the final layer of language models. It argues that different modules (hidden layers vs. output heads) require distinct update geometries based on their specific functional outputs (activations vs. probability distributions).

The author modestly frames their contribution as a proposal for a new geometry and an optimization problem. They acknowledge that:

The assumption of equal treatment for all token pairs may be too strong for synonyms or related subwords.
The regime $V \gg d$ necessitates an approximate solution rather than an exact geometric configuration.
The immediate next step is to determine whether this derived geometry can yield a scalable and practical optimizer that improves the condition number $\rho(S)$ while remaining efficient for large vocabularies.

The paper concludes that this approach connects the LM head update to concepts of sphere packing and neural collapse, suggesting that future work should explore whether language requires a semantic or data-dependent version of this geometry.

Technical Summary: Toward a First-Principles Update Geometry for the Language-Model Head

More like this