A research team says the part of a language model that maps words to numbers is being over-engineered by standard optimizers.
The paper, posted to arXiv, argues that the embedding table and LM-head — the layers that translate between raw text and a model's internal math — behave differently from the rest of a network's weights. The authors built Ember, an optimizer tuned specifically for those layers. Where Adam requires O(2VD) optimizer state memory for token tables, Ember needs only O(V + D), cutting storage to kilobytes instead of gigabytes for large vocabularies. The team claims it works across supervised fine-tuning, reinforcement learning, and pretraining, and scales with both batch size and model size.
The memory angle matters because optimizer state is often the silent budget-killer in large model training — it can dwarf the model weights themselves. If Ember's efficiency holds up under scrutiny, it could let teams run larger training runs on tighter hardware budgets, or simply drop one headache from distributed training setups that already lean on ZeRO or FSDP.
The paper also pushes back on a piece of received wisdom: that neural network training navigates a jagged, high-dimensional loss landscape. For token embeddings at least, the authors say the optimization path looks more like a straight line. That is either a useful simplification or an oversimplification — the open-source release should make it easier for others to find out.