Apple Machine Learning Research
·
在大词汇量语言模型中减少损失
内容提要
随着语言模型词汇量的增加，交叉熵的内存占用显著上升。我们提出了Cut Cross-Entropy（CCE）方法，仅计算正确标记的logit，从而显著降低内存消耗。以Gemma 2模型为例，CCE将损失计算的内存从24GB降至1MB，训练时间内存从28GB降至1GB，且不影响训练速度和收敛性。
关键要点
-
随着语言模型的词汇量增加，交叉熵的内存占用显著上升。
-
提出了Cut Cross-Entropy（CCE）方法，仅计算正确标记的logit，降低内存消耗。
-
以Gemma 2模型为例，CCE将损失计算的内存从24GB降至1MB。
-
训练时间内存从28GB降至1GB，且不影响训练速度和收敛性。
-
CCE通过在闪存中执行矩阵乘法和log-sum-exp减少全局内存消耗。
-
利用softmax的稀疏性，跳过对梯度计算贡献微小的元素，以提高CCE的吞吐量。
延伸解读
内存消耗的显著降低
随着语言模型的规模扩大，交叉熵的内存占用成为一个主要瓶颈。Cut Cross-Entropy（CCE）方法通过仅计算正确标记的logit，显著降低了内存消耗。这一创新使得在训练大型模型时，内存需求从24GB降至1MB，极大地提升了模型的可用性，尤其是在资源有限的环境中。
训练效率与收敛性
CCE方法不仅减少了内存消耗，还保持了训练速度和收敛性。这意味着研究人员和开发者可以在不牺牲性能的情况下，使用更大的词汇量进行训练。这一特性使得CCE在实际应用中具有较高的实用价值，尤其是在需要快速迭代和优化的场景中。
稀疏性利用的优势
CCE方法利用softmax的稀疏性，跳过对梯度计算贡献微小的元素，从而提高了计算效率。这一策略不仅优化了内存使用，还提升了整体的计算吞吐量。这对于需要处理大规模数据集的语言模型训练尤为重要，能够有效缩短训练时间。
延伸问答
Cut Cross-Entropy（CCE）方法的主要优点是什么？
CCE方法显著降低了内存消耗，从而减少了交叉熵计算的内存占用。
CCE方法如何影响Gemma 2模型的内存使用？
CCE将Gemma 2模型的损失计算内存从24GB降至1MB，训练时间内存从28GB降至1GB。
CCE方法是如何计算交叉熵损失的？
CCE方法仅计算正确标记的logit，并实时评估所有logit的log-sum-exp。
CCE方法在训练速度和收敛性上有什么影响？
CCE方法在显著降低内存消耗的同时，不影响训练速度和收敛性。
CCE方法如何利用softmax的稀疏性？
CCE方法跳过对梯度计算贡献微小的元素，以提高计算吞吐量。
大词汇量语言模型中交叉熵的内存占用为何会增加？
随着语言模型词汇量的增加，交叉熵的内存占用显著上升，主要集中在单一层上。
