pith. sign in

arxiv: 2607.01455 · v1 · pith:SUOGEZM2new · submitted 2026-07-01 · 💻 cs.LG · cs.AI

Token Geometry

Kathan Shah This is my paper

Pith reviewed 2026-07-03 21:06 UTC · model grok-4.3

classification 💻 cs.LG cs.AI
keywords language model embeddingsLM head optimizergradient geometrymemory-efficient trainingtoken trajectoryTransformer optimizervocabulary matrixPareto frontier
0
0 comments X

The pith

The embedding and LM-head interface in language models has distinct gradient geometry that a lightweight optimizer can exploit to reduce memory use while matching or exceeding standard performance.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper argues that embeddings and the language model head form a special read-write interface whose gradients differ markedly from those of dense hidden layers. This difference lets a new optimizer called Ember store only O(V + D) state instead of the full Adam requirement, cutting optimizer memory to kilobytes while still supporting supervised finetuning, reinforcement learning, and pretraining. Experiments show the method scales with batch size and model size and integrates with existing sharding systems. The work also reports that individual token parameters move along simple one-dimensional rays rather than through the expected non-convex landscape. These observations together suggest that only a narrow family of optimizers is actually needed for Transformer training.

Core claim

The embedding table and LM-head exhibit gradient geometry distinct from hidden weights; this geometry can be exploited by an optimizer that stores only O(V + D) VRAM and still improves the Pareto frontier across supervised finetuning, RL, and pretraining. Token trajectories are well described by a one-dimensional ray, and the same geometry explains why a small set of optimizer designs suffices for Transformer training.

What carries the argument

Ember, a lightweight optimizer for embedding and LM-head matrices that treats each token's update as motion along a one-dimensional ray and requires only O(V + D) optimizer state.

If this is right

  • Ember reduces total optimizer state for the vocabulary matrices from O(2VD) to O(V + D), eliminating the need to shard both embedding and head optimizer states.
  • The same optimizer works across supervised finetuning, reinforcement learning from human feedback, and pretraining without architecture changes.
  • Token-level updates follow a one-dimensional trajectory, which accounts for why only a narrow class of optimizers performs well on Transformer models.
  • The implementation merges directly with ZeRO and FSDP, allowing immediate use in existing large-scale training stacks.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • If the one-dimensional ray description holds at larger scales, then token-level learning rates or momentum could be set independently of the rest of the network without destabilizing training.
  • The geometry distinction may explain why vocabulary size often becomes a memory bottleneck even when hidden-layer optimizers are heavily sharded.
  • Similar interface-specific optimizers might be worth testing on other discrete-to-continuous mappings, such as code tokenizers or multimodal projection layers.

Load-bearing premise

The observed difference in gradient geometry between token tables and hidden weights is large enough that a specialized optimizer can improve efficiency without hidden costs to convergence speed or training stability.

What would settle it

Running Ember on a standard pretraining run and finding that final loss or downstream accuracy falls behind an otherwise identical Adam baseline by more than a small constant margin.

Figures

Figures reproduced from arXiv: 2607.01455 by Kathan Shah.

Figure 1
Figure 1. Figure 1: The Ember optimizer utilizes almost no optimizer state for the embedding table and [PITH_FULL_IMAGE:figures/full_fig_p001_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Embedding-optimizer ablation at two batch sizes for GPT-2 small / FineWeb. Interestingly, [PITH_FULL_IMAGE:figures/full_fig_p004_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Model-scaling results for Ember on token tables. Ember scales cleanly with model size, and matches or surpasses Adam regardless of scale. Training setup in Appendix 6.2 [PITH_FULL_IMAGE:figures/full_fig_p006_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: The memory–quality gap widens with model scale. The two optimizers reach the same loss at every scale, yet AdamW’s state grows as O(2V D) — crossing 1 GB by 1.4B and reaching ∼2 GB at 2.8B — while Ember’s O(V +D) state stays near 400 KB. unexpectedly observed that it decomposes separably into a monomial basis ( [PITH_FULL_IMAGE:figures/full_fig_p006_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Token optimization trajectories under Ember admit a clean polynomial structure. We show how strongly the monomials fit the first 3 principal component projections and show indeed the Taylor series decomposition holds surprisingly well. Moreover, in the energy spectrum, the linear component significantly dominates across tokens, indicating the overall trajectory can be approximated surprisingly simply by a … view at source ↗
Figure 6
Figure 6. Figure 6: Ember is competitive in reinforcement learning while removing O(2V D) optimizer state. AdamW and Ember achieve parity across rollout regardless of rollout batch size 64, 128, 192, and 256. 5 Discussion We open-source a distributed version of Ember compatible with ZeRO [22]. Because Ember’s opti￾mizer state is O(V +D), the embedding table optimizer state no longer requires sharding, simplifying distributed … view at source ↗
Figure 7
Figure 7. Figure 7: For each of four monitored tokens, we stack it over the time-axis to get [PITH_FULL_IMAGE:figures/full_fig_p012_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: Batch-size scaling. Ember follows AdamW’s log-batch trend and remains within seed noise across nine batch-size doublings [PITH_FULL_IMAGE:figures/full_fig_p012_8.png] view at source ↗
Figure 10
Figure 10. Figure 10: Ablating bias correction in Ember. Including the bias correction is crucial to performance and makes training much stabler while helping the optimizer reach lower val loss faster. (2) relative_step = min(10−2 , 1/ √ t), which replaces the learning rate (3) scale_parameter, which makes step size proportional to parameter scale (4) clip_threshold = 1.0, post-hoc RMS clipping. We observe that including the b… view at source ↗
Figure 11
Figure 11. Figure 11: Why removing Adam’s first-moment state is safe on token tables. We instrument gradients of the untied input embedding Ein ∈ R V ×d and output head Wout ∈ R V ×d during GPT￾2 small training on FineWeb. Left: Ember’s row-column factored second moment approximates the dense squared-gradient structure, where the y-axis is relative error. Right: temporal gradient autocorrelation is near zero or negative, showi… view at source ↗
Figure 12
Figure 12. Figure 12: Ember’s V +D factored second moment closely tracks the dense Adam second moment. Right: each point is one embedding parameter (i, j); the factored estimate rˆicˆj/r¯ tracks the dense gb 2 ij across many orders of magnitude. Left: log-space residuals concentrate near zero with no heavy tail, showing that most dense curvature structure is captured by the row/column outer product. Pointwise, the match holds … view at source ↗
read the original abstract

Language models 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

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

0 major / 2 minor

Summary. The paper claims that the embedding table and LM-head in language models exhibit gradient geometry distinct from dense hidden weights. This geometry is exploited by a new lightweight optimizer, Ember, which uses O(V + D) VRAM (versus Adam's O(2VD)) and only kilobytes of optimizer state to improve the Pareto frontier across supervised finetuning, RL, and pretraining. The manuscript further claims that token optimization trajectories are well-described by a simple 1D ray (contrary to typical nonconvex landscape views), provides a principled perspective on sufficient optimizers for Transformer training, demonstrates scaling across batch size and parameter count, and releases a distributed implementation compatible with ZeRO/FSDP.

Significance. If the empirical claims hold, the work could meaningfully reduce optimizer memory overhead for large-vocabulary models, enabling larger scales or training on constrained hardware without sharding both token-table optimizer states. The open-sourced implementation and explicit scaling tests across batch sizes and parameter counts are concrete strengths. The 1D-ray trajectory description, if supported by the full derivations and ablations, offers a falsifiable perspective that could influence optimizer design. The stress-test concern (whether distinct geometry yields Pareto gains without hidden convergence or stability costs) does not land as a load-bearing issue on the basis of the scaling evidence presented.

minor comments (2)
  1. The abstract states that Ember 'scales effectively across batch size and parameter count' and that the 1D ray description is 'well' supported; the main text should include explicit quantitative metrics (e.g., trajectory fit error, stability metrics, or Pareto deltas) with error bars or multiple seeds to allow readers to assess robustness.
  2. Notation for the memory complexity O(V + D) versus O(2VD) is clear in the abstract but would benefit from an explicit comparison table (e.g., Table X) showing actual VRAM usage and wall-clock overhead for representative V and D values.

Simulated Author's Rebuttal

0 responses · 0 unresolved

We thank the referee for their supportive summary, significance assessment, and recommendation of minor revision. The report contains no major comments requiring point-by-point rebuttal.

Circularity Check

0 steps flagged

No significant circularity detected

full rationale

The provided abstract and context describe empirical observations on gradient geometry for embeddings/LM-heads, a 1D ray trajectory description, and an optimizer Ember with O(V+D) state. No load-bearing derivation step is shown to reduce by construction to its own inputs via self-definition, fitted parameters renamed as predictions, or self-citation chains. The 1D ray claim is presented as an empirical finding rather than a fitted input called prediction, and the central claims rest on scaling tests and implementation details without internal reduction to tautology. The derivation chain is self-contained against external benchmarks.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Only the abstract is available, so free parameters, axioms, and invented entities cannot be enumerated from the paper's derivations or experiments.

pith-pipeline@v0.9.1-grok · 5705 in / 1039 out tokens · 19198 ms · 2026-07-03T21:06:15.968624+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

27 extracted references · 27 canonical work pages · 2 internal anchors

  1. [1]

    Kingma and Jimmy Ba

    Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. InInternational Conference on Learning Representations (ICLR), 2015

    work page 2015

  2. [2]

    Sham M. Kakade. A natural policy gradient. InAdvances in Neural Information Processing Systems 14, pages 1531–1538. MIT Press, 2001

    work page 2001

  3. [3]

    Heavy-tailed class imbalance and why Adam outperforms gradient descent on language models

    Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Heavy-tailed class imbalance and why Adam outperforms gradient descent on language models. InAdvances in Neural Information Processing Systems, 2024

    work page 2024

  4. [4]

    Modular duality in deep learning.arXiv preprint arXiv:2410.21265, 2024

    Jeremy Bernstein and Laker Newhouse. Modular duality in deep learning. InarXiv preprint arXiv:2410.21265, 2024

    work page arXiv 2024

  5. [5]

    Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y . K. Li, Y . Wu, and Daya Guo. DeepSeekMath: Pushing the limits of mathematical reasoning in open language models.arXiv preprint arXiv:2402.03300, 2024

    work page internal anchor Pith review Pith/arXiv arXiv 2024

  6. [6]

    N. N. Chentsov.Statistical Decision Rules and Optimal Inference. Translations of Mathematical Monographs, V ol. 53. American Mathematical Society, Providence, RI, 1982

    work page 1982

  7. [7]

    Natural gradient works efficiently in learning.Neural Computation, 10(2):251– 276, 1998

    Shun-ichi Amari. Natural gradient works efficiently in learning.Neural Computation, 10(2):251– 276, 1998

    work page 1998

  8. [8]

    Applied Mathematical Sciences, V ol

    Shun-ichi Amari.Information Geometry and Its Applications. Applied Mathematical Sciences, V ol. 194. Springer Japan, Tokyo, 2016

    work page 2016

  9. [9]

    Alexander Ly, Maarten Marsman, Josine Verhagen, Raoul P. P. P. Grasman, and Eric-Jan Wa- genmakers. A tutorial on Fisher information.Journal of Mathematical Psychology, 80:40–55, 2017

    work page 2017

  10. [10]

    An elementary introduction to information geometry.Entropy, 22(10):1100, 2020

    Frank Nielsen. An elementary introduction to information geometry.Entropy, 22(10):1100, 2020

    work page 2020

  11. [11]

    Revisiting natural gradient for deep networks

    Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. In International Conference on Learning Representations (ICLR), 2014

    work page 2014

  12. [12]

    Natural neural networks

    Guillaume Desjardins, Karen Simonyan, Razvan Pascanu, and Koray Kavukcuoglu. Natural neural networks. InAdvances in Neural Information Processing Systems, 2015

    work page 2015

  13. [13]

    Modular manifolds.Thinking Machines Lab: Connectionism, 2025

    Jeremy Bernstein. Modular manifolds.Thinking Machines Lab: Connectionism, 2025. https: //thinkingmachines.ai/blog/modular-manifolds/, doi:10.64434/tml.20250926

    work page doi:10.64434/tml.20250926 2025

  14. [14]

    DeepSeek-V4: Towards highly efficient million-token context intelligence

    DeepSeek-AI. DeepSeek-V4: Towards highly efficient million-token context intelligence. Technical Report, 2026

    work page 2026

  15. [15]

    Muon: An optimizer for LLMs.Keller Jordan’s Blog, 2024

    Keller Jordan. Muon: An optimizer for LLMs.Keller Jordan’s Blog, 2024. URL: https: //kellerjordan.github.io/posts/muon/

    work page 2024

  16. [16]

    Rumelhart, Geoffrey E

    David E. Rumelhart, Geoffrey E. Hinton, and Ronald J. Williams. Learning representations by back-propagating errors.Nature, 323(6088):533–536, 1986

    work page 1986

  17. [17]

    Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yoram Singer, and Quoc V . Le. Symbolic discovery of optimization algorithms.arXiv preprint arXiv:2302.06675, 2023

    work page arXiv 2023

  18. [18]

    Adafactor: Adaptive learning rates with sublinear memory cost

    Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. InInternational Conference on Machine Learning (ICML), 2018

    work page 2018

  19. [19]

    Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers

    Tim Tsz-Kit Lau and Weijie Su. Symmetry-compatible principle for optimizer design: Em- beddings, LM heads, SwiGLU MLPs, and MoE routers.arXiv preprint arXiv:2605.18106, 2026. 9

    work page internal anchor Pith review Pith/arXiv arXiv 2026

  20. [20]

    signSGD: Compressed optimisation for non-convex problems

    Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Anima Anandkumar. signSGD: Compressed optimisation for non-convex problems. InInternational Conference on Machine Learning (ICML), 2018

    work page 2018

  21. [21]

    Jorge Nocedal and Stephen J. Wright. Numerical optimization. Springer, 2nd edition, 2006

    work page 2006

  22. [22]

    ZeRO: Memory op- timizations toward training trillion parameter models

    Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. ZeRO: Memory op- timizations toward training trillion parameter models. InInternational Conference for High Performance Computing, Networking, Storage and Analysis (SC20), 2020

    work page 2020

  23. [23]

    M. S. Bartlett. Approximate confidence intervals. InBiometrika, 1953

    work page 1953

  24. [24]

    The law of error

    Francis Ysidro Edgeworth. The law of error. InTransactions of the Cambridge Philosophical Society, 20:33–66, 1905

    work page 1905

  25. [25]

    Lecture Notes in Statistics, V olume 28

    Shun-ichi Amari.Differential-Geometrical Methods in Statistics. Lecture Notes in Statistics, V olume 28. Springer, 1985

    work page 1985

  26. [26]

    Chentsov.Statistical Decision Rules and Optimal Inference

    Nikolai N. Chentsov.Statistical Decision Rules and Optimal Inference. American Mathematical Society, 1982

    work page 1982

  27. [27]

    Google AI.Gemma 3n model overview. 2025. https://ai.google.dev/gemma/docs/ gemma-3n. Accessed: 2026-06-17. 10 6 Appendix 6.1 Further work Parameter class-awareness.In general, a principled method of coming up with optimizers for different transformer parameter classes can likely save memory state and improve performance. Comfortingly, the search space of ...

    work page 2025