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