🤖 machine learning

Token Geometry

The paper introduces Ember, a lightweight optimizer that exploits the unique gradient geometry of embedding and LM-head matrices to significantly reduce VRAM usage and improve training efficiency across various tasks, while challenging the conventional view of neural network optimization landscapes.

Original authors: Kathan Shah

Published 2026-07-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kathan Shah

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

Imagine you are teaching a giant, super-smart robot to speak human language. To do this, you give the robot a massive dictionary (the embedding table) and a translation guide (the LM-head) that connects its internal thoughts to actual words.

For a long time, the standard way to teach this robot was using a method called Adam. Think of Adam as a very thorough, but heavy-handed teacher. For every single word in the dictionary, Adam keeps a massive notebook of "study notes" (optimizer state) to remember how the robot reacted to that word in the past.

The problem? As the dictionary grows to include millions of words, these notebooks become so huge they don't fit in the robot's memory (VRAM). This forces researchers to split the work across many computers, which is slow, expensive, and complicated.

Enter Ember, a new, lightweight teacher introduced in this paper.

The Big Idea: A Lighter Touch

The authors of the paper discovered that the "dictionary" part of the robot learns differently than the rest of its brain. While the rest of the brain needs complex, heavy notes, the dictionary only needs a simple, streamlined approach.

The Analogy: The "Z-Score" Check
Imagine you are grading a student's test.

Adam looks at every single answer, remembers every time the student got it right or wrong, and keeps a detailed file on every specific mistake. It's like keeping a 100-page biography for every single word.
Ember asks a simpler question: "How often does this student get this word right, and how confident are they?" It essentially converts the student's performance into a simple standard score (a "z-score"). It strips away the heavy baggage and just focuses on the core signal: Is the robot getting better at this word, or not?

How Ember Saves Space

The paper claims that Ember is incredibly efficient because it stops keeping those massive 100-page biographies.

Adam's Memory: If you have a dictionary with 100,000 words, Adam needs a notebook for every single word that is huge. The paper says this takes up gigabytes of space (like a whole library).
Ember's Memory: Ember realizes it only needs a tiny list of "how often" and "how confident" for each word. It shrinks that library down to a single post-it note (kilobytes).

The paper shows that even though Ember is so much lighter, it teaches the robot just as well as the heavy-handed Adam. In fact, in some tricky situations (like when the robot is learning from very small batches of data), Ember actually learns faster and more stably.

The Surprising Discovery: A Straight Line

One of the most fascinating findings in the paper is about how the robot learns.

Old Belief: Scientists thought the robot's learning path was like a chaotic hike through a foggy mountain range, zig-zagging up and down, getting stuck in small valleys, and taking a very long, winding route to the top.
Ember's Reality: The authors found that when using Ember, the robot's learning path is actually a straight, smooth highway. If you plotted the robot's progress on a graph, it looked like a simple, straight line moving from "beginner" to "expert."

This suggests that the "landscape" of learning for these specific dictionary parts isn't as messy as we thought. It's a direct path, and Ember is the vehicle that drives straight down it without getting lost.

Why This Matters (According to the Paper)

It's Cheap: You can train these models on a single computer instead of needing a massive supercomputer cluster just to hold the memory.
It's Fast: Because the memory is so small, researchers can test new ideas much faster.
It Works Everywhere: The paper tested Ember on different sizes of robots (from small to very large) and different tasks (learning to speak, learning to reason, and even generating images). In almost every case, Ember matched or beat the old standard (Adam) while using a fraction of the memory.

Summary

The paper introduces Ember, a new way to train the "vocabulary" part of AI models. It replaces the heavy, memory-hogging "Adam" method with a clever, lightweight approach that treats the learning process like a simple math problem (a z-score) rather than a complex biography. The result? You get the same (or better) intelligence, but you don't need a supercomputer to hold the notes. Plus, it turns out the robot learns in a straight line, not a chaotic zig-zag.

Technical Summary: Token Geometry and the Ember Optimizer

Problem Statement

Current large language models (LLMs) optimize embedding tables and LM-head matrices (collectively termed the "token interface") using the Adam optimizer. While Adam performs well, it imposes a significant engineering burden in distributed training. Specifically, Adam maintains first and second moment buffers for every parameter, resulting in an optimizer state complexity of $O(2VD)$ for a vocabulary size $V$ and embedding dimension $D$ . For large models, this state consumes gigabytes of VRAM, necessitating complex sharding strategies (e.g., ZeRO/FSDP) even for parameters that are not the primary bottleneck of the model's compute.

Furthermore, the paper argues that the standard application of Adam to token interfaces ignores the specific geometry of these parameters. Unlike dense linear layers, token interfaces are updated via sparse gradients (only the active token row is updated per step), and the underlying loss landscape for cross-entropy may possess distinct properties that standard optimizers fail to leverage efficiently.

Methodology: The Ember Optimizer

The authors introduce Ember, a lightweight optimizer designed specifically for embedding and LM-head matrices. The core methodology is grounded in Information Geometry, specifically the Fisher Information metric.

Theoretical Foundation

The paper posits that for cross-entropy loss, the Fisher metric (the expected outer product of gradients) is the exact second-order term of the Kullback–Leibler divergence. Near the optimum, the Fisher metric approximates the Hessian.

Square-Root Fisher: Instead of applying a full Newton step (inverting the Hessian), which is computationally expensive and unstable, the authors argue that conditioning gradients by the square root of the Fisher metric is sufficient. This effectively normalizes the gradient by its standard deviation (a "z-score" transformation).
Gradient Sparsity: In token interfaces, the gradient for a token $i$ is non-zero only when that token is selected. The paper observes that the row-wise gradient energy is an unbiased estimator of the token's participation probability ( $p_i$ ).

Algorithmic Design

Ember replaces Adam's $O(2VD)$ memory footprint with an $O(V + D)$ approach by factorizing the second-moment estimate:

Removal of First Moment: The authors found that removing Adam's first-moment (momentum) EMA entirely favors performance and saves memory.
Factorized Second Moment: Instead of maintaining a full $V \times D$ $V \times D$ second-moment buffer, Ember maintains two 1D vectors:
- $r_t$ : A row-wise second-moment estimate (tracking token frequency/energy).
- $c_t$ : A column-wise second-moment estimate (tracking feature-wise energy).
Preconditioning: The update rule forms a lightweight outer product preconditioner. The effective denominator for a parameter $(i, j)$ is derived from the geometric mean of the row and column estimates ( $\sqrt{\hat{r}_i \hat{c}_j}$ ).
Bias Correction: Standard bias correction is applied to the row and column estimates before forming the preconditioner.

The resulting update rule effectively divides the gradient by an estimate of the square root of the Fisher metric, normalized to be independent of token frequency while reducing variance for frequent tokens.

Key Contributions

Memory Efficiency: Ember reduces the optimizer state for token interfaces from $O(2VD)$ to $O(V + D)$ . This eliminates the need to shard optimizer states for embedding tables, allowing the entire state to fit on a single GPU even for massive models.
Geometric Insight: The paper provides a principled view that competitive optimizers for Transformers (including Adam, Lion, and Muon) can be unified under the framework of conditioning gradients by the square root of the Fisher metric (or its matrix analogues).
Optimization Trajectory Analysis: Through Singular Value Decomposition (SVD) of token trajectories, the authors demonstrate that embedding tokens optimize along a surprisingly simple 1D ray (monomial basis). This challenges the prevailing view that neural network optimization navigates a highly non-convex, chaotic landscape, suggesting instead a smooth, linear convergence path for token interfaces.
Open-Source Implementation: A distributed implementation compatible with existing ZeRO/FSDP setups is released, facilitating immediate adoption.

Results

The authors provide empirical evidence across multiple regimes:

Supervised Fine-Tuning (SFT): On GPT-2 and Pythia models, Ember matches or slightly improves upon AdamW validation loss while reducing optimizer state from gigabytes to kilobytes (e.g., 2.06 GB $\to$ 413 KB for Pythia-2.8B).
Reinforcement Learning (RL): In post-training Qwen2.5-3B with GRPO, Ember achieves parity with AdamW across various rollout batch sizes while using 4,000x less optimizer state.
Pretraining: On FineWeb, Ember matches AdamW and outperforms Adafactor, particularly in low-batch regimes where the rank-1 gradient decomposition provides stability against sparse/noisy gradients.
Scaling Laws: Ember scales effectively with model size and batch size. The performance gap between Ember and AdamW vanishes as model size increases, suggesting that Adam's full second-moment buffer is well-approximated by Ember's rank-1 factorization at scale.
Stability: Ember exhibits no early-training loss spikes at small batch sizes, a failure mode observed in AdamW.

Significance and Claims

The paper claims that Ember offers a "Pareto frontier" improvement, achieving state-of-the-art performance while drastically reducing memory overhead. The significance is twofold:

Practical Engineering: By removing the need to shard embedding optimizer states, Ember simplifies distributed training setups and enables faster iteration on single devices for models up to 7B parameters (and beyond), where Adam's state was previously a binding constraint.
Theoretical Understanding: The discovery that token optimization trajectories decompose into a clean 1D monomial basis suggests that the "non-convex" nature of the loss landscape may be an artifact of the optimizer or the parameterization, rather than an intrinsic property of the token interface. This implies that the search space for optimal optimizers for Transformers is narrower than previously thought, centered around the square-root Fisher metric.

The authors conclude that while Ember is a practical tool for memory-constrained training, its theoretical framing provides a new lens for understanding how LLMs learn, specifically highlighting the simplicity of the optimization path for token representations.