论文

语言模型预训练的贪婪局部学习:差距和目标设计

Greedy Local Learning for Language Model Pretraining: Gaps and Objective Design

模型训练预训练

摘要

贪婪的逐块局部学习将网络分割为由局部辅助损失训练的梯度隔离块,删除块之间的后向传递:级间通信变为仅前向,每个块都可以独立地步进其优化器,这些属性与去中心化模型并行训练直接相关。本地学习在图像分类方面与端到端反向传播具有竞争力,并且在小型 Transformer 上,众所周知,它会以更差的最佳损失换取并行加速。这种损失差距在更大规模的自回归语言模型 (LM) 预训练中的表现如何,以及哪些辅助设计可以减少这种损失差距,尚未得到测量。我们提出了token预算匹配实证研究,在 125M 和 400M 参数上使用 $K \in \{1,2,4\}$ 块,在 Chinchilla 最佳预算下,将辅助设计分解为网络架构和训练目标。我们观察到:(i)从 $K=2$ 到 $K=4$ 到端到端训练的差距增加了一倍多,但在 $K=4$ 则从 125M 缩小到 400M; (ii) 用基于 Transformer 的 MLP 辅助替代是强有力的网络侧干预,弥补了 22-41% 的差距; (iii) 多token预测 (MTP) 辅助目标有助于第一个块边界,而在更深的边界处添加它会产生伤害,并将其限制在第一个块会产生最佳的 $K=4$ 配置($+0.062$ 与 $+0.075$ nat 在 400M 处); (iv) 部署式每块执行可将激活记忆减少多达 $2.2\times$。我们将这些结果视为基于经验的方法方向,而不是最终的方法:局部目标应有选择地跨越边界施加未来预测压力,同时抵制绕过预测内容的捷径。