词元几何

语言模型通过离散符号学习连续程序，嵌入表和 LM-head 充当它们之间的读/写接口。我们表明，该接口具有与密集隐藏权重不同的梯度几何形状，可用于改进监督微调、强化学习和预训练的帕累托前沿，同时仅利用数千字节的优化器状态。我们引入了 Ember，这是一种用于嵌入和 LM-head 矩阵的轻量级优化器，它利用 O(V + D) VRAM，而不是 Adam 的 O(2VD)，并且无需对两个词元表优化器状态进行分片。我们提供的经验证据表明，Ember 可有效地跨批量大小和参数计数进行扩展。我们证明，词元的优化轨迹可以通过简单的一维射线很好地描述，这与神经网络参数在严重非凸景观中导航的流行观点相反。我们对足以满足 Transformer 训练的优化器空间提供了一个原则性的观点。最后，我们开源了我们的分布式 Ember 实现，它与现有的 ZeRO/FSDP 设置完全合并，以支持进一步的研究，网址为 https://github.com/katop1234/ember