analiza aiqbench
Ember introduces a novel optimizer specifically targeting the gradient geometry of embedding tables and LM-heads, significantly reducing VRAM overhead from $\mathcal{O}(2VD)$ to $\mathcal{O}(V + D)$. By eliminating the need to shard token table states, this approach improves the Pareto frontier for pretraining and fine-tuning tasks while maintaining compatibility with standard distributed frameworks like ZeRO/FSDP. The method demonstrates consistent loss reductions across multiple seeds and scales effectively with parameter count and batch size.
Ember: A Kilobyte-State Optimizer for Token Interfaces
arXiv:2607.01455v4 Announce Type: replace-cross 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.
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.
Machine Learning (cs.LG); Artificial Intelligence (cs.AI)
arXiv:2607.01455 [cs.LG]
(or arXiv:2607.01455v4 [cs.LG] for this version)
https://doi.org/10.48550/arXiv.2607.01455 Focus to learn more arXiv-issued DOI via DataCite
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Text integral sub licența deschisă a publicației · 233 cuvinte
Subiecte conexe
Debugging Logprob Mismatches in LLM Reinforcement Learning
AMD · 23 sept. 2026
[AINews] DeepSeek v4.1-Flash: 763B-P8B-D16B novel causal Encoder–Decoder architecture with vision marks the Return of the Whale
Latent Space · 12 sept. 2026
Bridging intent and execution in agentic systems
Amazon Science · 8 iun. 2026
Building Blocks for Foundation Model Training and Inference on AWS
Hugging Face · 11 mai 2026