提出轻量级优化器Ember，提升语言模型训练效率与效果。
Token Geometry
- 针对嵌入层和LM头设计专用优化器，利用独特梯度几何。
- 仅需千字节显存，实现超大模型高效训练，支持批量与参数规模扩展。
- 发现词元优化轨迹呈一维射线，为高效训练提供新视角，适合大规模模型研究者。
语言模型在离散符号上学习连续程序，嵌入表和LM头作为读写接口。我们发现该接口的梯度几何不同于密集隐藏权重，可被用于改进监督微调、强化学习和预训练中的帕累托前沿，且仅需千字节优化器状态。我们提出Ember，一种针对嵌入和LM头矩阵的轻量级优化器，显存占用仅为O(V + D)，远低于Adam的O(2VD)，并避免对词表优化状态分片。实验证明Ember在批大小和参数量上均具有良好的可扩展性。我们发现词元优化轨迹可用简单的一维射线描述，反驳了神经网络参数在高度非凸空间中游走的普遍认知。本文从理论上揭示了适用于Transformer训练的优化器空间极为狭窄。最后，我们开源了可无缝集成现有ZeRO/FSDP架构的分布式Ember实现，网址为https://github.com/katop1234/ember。
原文摘要 · Abstract (English)
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 O(V + D) VRAM, instead of Adam's O(2VD), and forgoes the need to shard both token table optimizer states. We provide empirical evidence that Ember scales effectively across batch size and parameter count. We show that the optimization trajectory of tokens can be well described by a simple 1D ray, counter to the popular belief that neural net parameters navigate a heavily nonconvex landscape. We provide a principled view on the surprisingly narrow space of optimizers that suffice for Transformer training. Finally, we open-source our distributed Ember implementation that merges cleanly with existing ZeRO/FSDP setups to support further research at https://github.com/katop1234/ember
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。