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) — accepted at the Sci-FM and MOSS workshops @ COLM 2026
• Code: github.com/katop1234/ember — PyTorch implementation, integrates with existing ZeRO/FSDP setups