arXiv:2410.21265cs.LGcs.NE2024-10ICML被引 71

提出模块化对偶性理论,让神经网络训练更快更可扩展。

Modular Duality in Deep Learning

  • 按层语义分配算子范数,递归构建权重空间对偶映射。
  • 推导出适配GPU的Embed、Linear和Conv2D层对偶算法。
  • 方法已用于创纪录训练NanoGPT,适合追求高效优化的开发者。

优化理论中有一个古老观点:梯度是对偶向量,不能直接从权重中减去,而应先映射到权重所在的原空间。本文认真对待这一思想,为通用神经网络构建了这样的对偶映射,称为模块化对偶化。该方法首先根据每层的语义为各层分配算子范数,再利用这些逐层范数递归诱导整个网络结构的权重空间对偶映射。最后推导出适用于GPU的嵌入(Embed)、线性(Linear)和二维卷积(Conv2D)层的对偶化算法,其中后两者基于矩形牛顿-舒尔茨迭代(Kovarik, 1970;Björck & Bowie, 1971)。该方法的变体曾用于创纪录地训练NanoGPT。我们希望模块化对偶性理论能催生下一代快速且可扩展的优化器,适用于各类神经网络架构。

原文摘要 · Abstract (English)

An old idea in optimization theory says that since the gradient is a dual vector it may not be subtracted from the weights without first being mapped to the primal space where the weights reside. We take this idea seriously in this paper and construct such a duality map for general neural networks. Our map, which we call modular dualization, forms a unifying theoretical basis for training algorithms that are a) fast and b) scalable. Modular dualization involves first assigning operator norms to layers based on the semantics of each layer, and then using these layerwise norms to recursively induce a duality map on the weight space of the full neural architecture. We conclude by deriving GPU-friendly algorithms for dualizing Embed, Linear and Conv2D layers -- the latter two methods are based on a rectangular Newton-Schulz iteration (Kovarik, 1970; Björck & Bowie, 1971). A variant of our methods was used to set speed records for training NanoGPT. Overall, we hope that our theory of modular duality will yield a next generation of fast and scalable optimizers for general neural architectures.

深度学习优化器对偶性可扩展

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。