arXiv：cs.AI· Kathan Shah·· 20 小时前
Ember：面向 Token 接口的千字节级优化器
Ember: A Kilobyte-State Optimizer for Token Interfaces
AI 导读
研究者提出轻量优化器 Ember，专用于 embedding 与 LM-head 矩阵，优化器显存占用为 O(V+D)，而 Adam 为 O(2VD)，且无需对 token 表优化器状态做分片。
正文
Abstract:Language models learn continuous programs over discrete symbols, with the embedding table and LM-head acting as the read/write interface between them. We show that this interface has gradient geometry distinct from dense hidden weights which can be exploited to improve the Pareto frontier across supervised finetuning, RL, and pretraining, while only utilizing kilobytes of optimizer state. We introduce Ember, a lightweight optimizer for embedding and LM-head matrices that utilizes $\mathcal{O}(V + D)$ optimizer VRAM, instead of Adam's $\mathcal{O}(2VD)$, and forgoes the need to shard both token table optimizer states. On 50k-step pretraining across two datasets and three paired seeds, Ember lowers val loss by 0.08-0.09 nats relative to tuned AdamW; we also provide empirical evidence that Ember scales across batch size and parameter count. Finally, we provide a distributed Ember implementation that merges cleanly with existing ZeRO/FSDP setups.
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2607.01455 [cs.LG] |
| (or arXiv:2607.01455v4 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2607.01455 arXiv-issued DOI via DataCite |
Submission history
From: Kathan Shah [view email]
[v1]
Wed, 1 Jul 2026 20:21:03 UTC (2,105 KB)
[v2]
Tue, 7 Jul 2026 07:29:48 UTC (2,105 KB)
[v3]
Wed, 15 Jul 2026 02:52:22 UTC (2,108 KB)
[v4]
Sun, 11 Oct 2026 23:28:50 UTC (2,098 KB)
来源：arXiv：cs.AI · arxiv.org