归一化层 per-example 梯度足以预测变换器中的梯度噪声规模
机器学习
2024-11-05 v1 机器学习
摘要
per-example 梯度范数是估计梯度噪声规模 (GNS) 的关键要素，可实现最小方差。观察计算它们所需的张量收缩操作，我们提出一种在 3D 或更高维张量场景下以最小 FLOPs 计算该方法的方案，同时计算参数梯度。使用该方法，我们能够比以往更准确地观测变换器模型中各层的 GNS。我们发现，现代变换器模型的总 GNS 可由仅归一化层的 GNS 良好预测。因此，我们开发了用于在 LayerNorm 反向传播期间计算 per-example 梯度范数且无额外吞吐量开销的自定义内核。通过仅跟踪这些层的 GNS，我们能够指导一个实用的 batch size 计划，使在 Chinchilla-optimal 语言模型上的训练时间缩短 18%。
引用
@article{arxiv.2411.00999,
title = {Normalization Layer Per-Example Gradients are Sufficient to Predict Gradient Noise Scale in Transformers},
author = {Gavia Gray and Aman Tiwari and Shane Bergsma and Joel Hestness},
journal= {arXiv preprint arXiv:2411.00999},
year = {2024}
}
备注
23 pages, 16 figures, to be published in the proceedings of the 2024 Conference on Neural Information Processing Systems (NeurIPS), code is available at: https://github.com/CerebrasResearch/nanoGNS