# Ember: A Kilobyte-State Optimizer for Token Interfaces

arXiv: https://arxiv.org/abs/2607.01455 · HTML: https://mlsanity.com/paper/2607.01455/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 34 references that are in the map.
- First posted on arXiv 1 Jul 2026.

## Builds on (most cited first, 18 of 18)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence](https://mlsanity.com/paper/2606.19348/) (2026)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [signSGD: Compressed Optimisation for Non-Convex Problems](https://mlsanity.com/paper/1802.04434/) (2018)
- [Symbolic Discovery of Optimization Algorithms](https://mlsanity.com/paper/2302.06675/) (2023)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Gradient Descent Happens in a Tiny Subspace](https://mlsanity.com/paper/1812.04754/) (2018)
- [Limitations of the Empirical Fisher Approximation for Natural Gradient Descent](https://mlsanity.com/paper/1905.12558/) (2019)
- [Qualitatively characterizing neural network optimization problems](https://mlsanity.com/paper/1412.6544/) (2014)
- [Revisiting Natural Gradient for Deep Networks](https://mlsanity.com/paper/1301.3584/) (2013)
- [Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models](https://mlsanity.com/paper/2402.19449/) (2024)
- [Modular Duality in Deep Learning](https://mlsanity.com/paper/2410.21265/) (2024)
- [An elementary introduction to information geometry](https://mlsanity.com/paper/1808.08271/) (2018)
- [Natural Neural Networks](https://mlsanity.com/paper/1507.00210/) (2015)
- [Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers](https://mlsanity.com/paper/2605.18106/) (2026)

## Built on by (most cited first, 1 of 1)

- [MuonIO: Principled Norm-Aware Descent for Embedding Tables and Language Model Heads](https://mlsanity.com/paper/2610.02705/) (2026)

## Abstract (arXiv)

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) optimizer VRAM, instead of Adam's O(2VD), and forgoes the need to shard both token table optimizer states. On 50k-step pretraining across two datasets and three paired seeds, Ember lowers val loss by 0.08-0.09 nats relative to tuned AdamW; we also provide empirical evidence that Ember scales across batch size and parameter count. Finally, we provide a distributed Ember implementation that merges cleanly with existing ZeRO/FSDP setups.