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