---
license: mit
tags:
- optimizer
- pytorch
- memory-efficient-training
---

# Ember — an O(V+D) optimizer for token interfaces

Ember is a lightweight optimizer for embedding tables and LM-head matrices. It replaces
Adam's dense first- and second-moment state on those layers — O(2VD) — with row/column
factored second moments, O(V+D): kilobytes of optimizer state instead of gigabytes, and no
sharding of token-table optimizer state in distributed setups.

Across supervised finetuning, RL, and pretraining, Ember matches Adam's validation loss on
these layers while carrying ~1500× less optimizer state.

- **Paper:** [Token Geometry (arXiv:2607.01455)](https://arxiv.org/abs/2607.01455) —
  accepted at the Sci-FM and MOSS workshops @ COLM 2026
- **Code:** [github.com/katop1234/ember](https://github.com/katop1234/ember) — PyTorch
  implementation, integrates with existing ZeRO/FSDP setups