File size: 908 Bytes
31be279 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | ---
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
|