Token Geometry

Abstract: LLMs learn continuous programs over discrete symbols, with the embedding table and LM-head acting as the read/write interface between them. We show that this interface has gradient geometry distinct from dense hidden weights which can be exploited to improve the Pareto frontier across supervised finetuning, RL, and pretraining, while only utilizing kilobytes of optimizer state. We introduce Ember, a lightweight optimizer for embedding and LM-head matrices that utilizes O(V + D) VRAM, instead of Adam's O(2VD), and forgoes the need to shard both token table optimizer states. We provide empirical evidence that Ember scales effectively across batch size and parameter count. We show that the optimization trajectory of tokens can be well described by a simple 1D ray, counter to the popular belief that neural net parameters navigate a heavily nonconvex landscape. We provide a principled view on the surprisingly narrow space of optimizers that suffice for Transformer training. Finally, we open-source our distributed Ember implementation that merges cleanly with existing ZeRO/FSDP setups to support further research at https://github.com/katop1234/ember

• The paper introduces Ember, which leverages token-specific Fisher information geometry to reduce optimizer state from O(VD) to O(V+D) while maintaining performance.
• It employs factored second-moment estimation using row/column EMAs, eliminating the need for first moment accumulation due to low gradient autocorrelation.
• Empirical results show Ember matches or outperforms Adam in language and RL tasks, enabling training of larger models with minimal VRAM usage.

Token Geometry: A Principled and Memory-Efficient Optimizer for Transformer Token Interfaces

Problem Setting and Motivations

The optimization of embedding tables and LM-head matrices—collectively the "token interface" parameters—in Transformer architectures is a core scalability and efficiency bottleneck. The field’s de facto standard, Adam, entails a $\mathcal{O}(2VD)$ optimizer state per token table (with $V$ vocabulary size and $D$ embedding dimension), resulting in severe memory constraints as both $V$ and $D$ scale. Despite their structural and activation-sparsity differences from dense hidden weights, token tables typically use Adam without dedicated geometric consideration, and nearly all proposed Adam alternatives in the literature focus exclusively on dense linear layer parameters, not token interfaces.

This work addresses this discrepancy. It identifies the unique Fisher information geometry of the token interface and exploits this structure to derive Ember, a factored optimizer that requires only $\mathcal{O}(V+D)$ state while outperforming or matching the performance of Adam across language modeling, reinforcement learning (RL), and autoregressive generation tasks.

Theoretical Insights and Derivation

The Fisher metric is central to the analysis, grounded in information geometry. For cross-entropy loss, the Fisher Information Matrix precisely captures the expected local curvature, providing a mathematically canonical preconditioner for parameter updates. For embedding tables, the row-wise gradient energy serves as an unbiased estimator of token activation probability, $p_i$, leading to the observation that the second moment’s state can be accurately captured by row-wise and column-wise statistics rather than full $V \times D$ matrices. Adam’s update, in effect, computes a square-root Fisher preconditioned step per-parameter.

Ember’s derivation is anchored in these principles. The optimizer:

• Drops the first-moment accumulation entirely for token tables—backed by empirical near-zero temporal autocorrelation of gradients—eliminating $\mathcal{O}(VD)$ state.
• Approximates per-parameter curvature with an outer product of row- and column-wise exponential moving averages (EMAs) of squared gradients, requiring only $\mathcal{O}(V+D)$ VRAM.
• Normalizes this estimate with a geometric mean to maintain correct units and optimal variance properties.
• Implements only a single hyperparameter, $V$0, making it plug-and-play with existing Adam setups.

Empirical Results

A series of rigorous experiments in both supervised and RL domains comprehensively evaluate Ember’s efficacy. Notably, across model scales from distilgpt2 to Pythia-2.8B, Ember universally matches or marginally improves over Adam in validation loss on language modeling, while its optimizer state remains orders of magnitude smaller.

Figure 2: Ember achieves parity or better loss than Adam across transformer scales and tasks while scaling cleanly in VRAM.

The Pareto memory-quality gap is evident (Figure 3): Adam's optimizer state grows as $V$1, reaching multi-GB size for large models, while Ember's state increases only linearly and remains within hundreds of KB, even for billion-parameter transformers.

Figure 1: Ember dramatically reduces optimizer state requirements, widening the VRAM efficiency gap as model size increases.

Batch-size scaling is preserved; Ember matches Adam's validation trends and is robust to gradient sparsity and low-batch regimes, outperforming AdamW when batch sizes are small or gradients are noisy (Figure 4).

Figure 3: Validation loss across nine batch-size scales; Ember remains within seed noise of AdamW.

RL experiments further solidify generality: post-training Qwen2.5-3B-Instruct with GRPO on the Countdown task yields performance indistinguishable from AdamW within random seed variation, but with a $V$2 state memory reduction (Figure 5).

Figure 6: RL post-training: Ember achieves on-par results with AdamW using negligible optimizer state.

Optimization Landscape and Geometry

The optimization trajectory for individual token embeddings under Ember displays a strikingly regular structure. Principal Component Analysis (PCA) of token trajectories during training shows that over 90% of variance is explained by a single linear direction, with second and higher-order monomials capturing only minor residuals (Figure 6). This contradicts the common belief that neural parameter updates traverse a highly nonconvex or chaotic landscape; for token tables under the Ember update, parameter optimization follows a near-1D ray from initialization to solution.

Figure 5: Token update trajectories project cleanly onto a polynomial (monomial) basis, with the leading component linear.

Supplements confirm that higher-order basis components (quadratic, cubic, etc.) are present but diminish rapidly, supporting a simple, analytically tractable picture of the parameter geometry (Figure 7).

Figure 7: Across multiple tokens, trajectory projections along principal components reveal monomial structure: PC1 (linear), PC2 (quadratic), etc.

Comparative Analysis with Existing Optimizers

While Adafactor also incorporates a factored second moment and omits the first moment for efficiency, Ember’s geometric derivation from the Fisher metric unifies and clarifies why such structures excel. Bias correction is demonstrated to be critical—without it, performance and stability degrade sharply (Figure 8).

Figure 4: Ablating bias correction in Ember leads to instability and worsened validation loss.

The Muon optimizer’s conditioning in the spectral basis is shown to be mathematically equivalent to square-root Fisher preconditioning, but for embedding tables, the one-hot activation naturally eliminates the utility of such spectral conditioning. Ember, by leveraging the actual activation structure, achieves optimality with superior implementation and scalability properties.

Further, empirical analysis shows that dropping Adam’s first moment incurs no loss of performance; temporal autocorrelation of token gradients is minimal, rendering cumulative statistics redundant (Figure 9).

Figure 8: Autocorrelation of gradients for token tables is effectively zero, justifying omission of the first moment.

Finally, precise decomposition validates that token table gradients are rank-1 outer products, so the factored approach is not only memory-efficient but nearly lossless with respect to full diagonal second-moment estimates (Figure 10).

Figure 9: Ember’s rank-1 second moment factorization closely matches dense full-matrix estimates, capturing essential curvature with negligible error.

Practical and Theoretical Implications

Ember enables training and experimentation with larger vocabulary/token interface sizes on a single device, removing the necessity for distributed sharding due to VRAM constraints. This is especially impactful for recent architectures with per-layer embedding (as in Gemma 3n), multimodal models with large or multiple vocabularies, and research settings with limited hardware. Its cleaner implementation profile (all 1D buffers) streamlines engineering and integration, especially with ZeRO and FSDP.

On the theoretical front, the identification of a canonical, Fisher-metric-determined geometry for the full class of competitive optimizers provides clarity for the design space. The finding that embedding table optimization reduces to a 1D ray in parameter space challenges popular intuitions about deep optimization dynamics and suggests new analytic pathways for understanding generalization and learning in LLMs.

Conclusion

Token Geometry introduces Ember, a token interface optimizer that retains or improves the learning efficacy of Adam while virtually eliminating associated optimizer state memory costs. By centering the Fisher metric, the work unifies disparate lines of optimizer research and exposes practical and geometric principles guiding efficient deep learning. These results have immediate application for large-model training and provoke new theoretical questions regarding the dynamics and geometry of overparameterized neural networks.