Initialization of Large Language Models via Reparameterization to Mitigate Loss Spikes

📄 arXiv: 2410.05052v1 📥 PDF

作者: Kosuke Nishida, Kyosuke Nishida, Kuniko Saito

分类: cs.CL

发布日期: 2024-10-07

备注: EMNLP2024 accepted

💡 一句话要点

提出WeSaR重参数化方法，缓解大语言模型预训练中的Loss Spikes问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 预训练 损失尖峰 重参数化 Transformer 权重初始化 模型优化

📋 核心要点

  1. 大语言模型预训练中存在损失尖峰问题，现有方法难以有效缓解，影响训练稳定性和效率。
  2. 提出WeSaR重参数化方法，通过引入门参数调整权重范数，使其均匀化，从而稳定训练过程。
  3. 实验结果表明，WeSaR在不同规模的Transformer模型上均能有效稳定并加速训练，优于现有初始化方法。

📝 摘要（中文）

本文关注大语言模型预训练中损失尖峰（Loss Spikes）这一根本问题，假设参数范数的不均匀性是造成损失尖峰的原因之一。在神经网络训练中，为了避免梯度消失和梯度爆炸问题，需要保持各层梯度尺度的恒定。然而，为了满足Transformer模型的这一要求，模型参数的范数必须是非均匀的，因此，范数较小的参数对参数更新更为敏感。为了解决这个问题，我们提出了一种新的技术，即权重缩放作为重参数化（WeSaR）。WeSaR为每个参数矩阵引入一个门参数，并将其调整到满足要求的值。由于门参数的存在，WeSaR统一了原始参数的范数，从而实现了稳定的训练。在包含1.3亿、13亿和130亿参数的Transformer解码器上的实验结果表明，WeSaR稳定并加速了训练，并且优于包括流行的初始化方法在内的比较方法。

🔬 方法详解

问题定义：论文旨在解决大语言模型预训练过程中出现的损失尖峰（Loss Spikes）问题。现有的初始化方法和训练策略无法完全避免损失尖峰的出现，这会导致训练不稳定，甚至中断训练过程。根本原因在于Transformer模型中，为了保证梯度在各层之间传递的尺度不变性，参数的范数需要是非均匀分布的，这使得部分参数对更新非常敏感，容易引起突变，从而导致损失尖峰。

核心思路：论文的核心思路是通过重参数化的方式，使得模型的原始参数范数分布更加均匀，从而降低参数对更新的敏感性，进而缓解损失尖峰问题。具体来说，就是为每个参数矩阵引入一个门参数，通过调整门参数的值，来控制原始参数的有效范数，使其接近均匀分布。

技术框架：WeSaR方法主要是在Transformer模型的权重矩阵上进行操作。对于每一个权重矩阵，引入一个可学习的门参数。在模型前向传播过程中，权重矩阵首先与对应的门参数相乘，然后再进行后续的计算。门参数的初始值被设置为一个合适的值，以保证初始状态下模型的性能不受影响。在训练过程中，门参数与权重参数一起进行更新。

关键创新：WeSaR的关键创新在于将权重缩放视为一种重参数化方法，通过引入门参数来动态调整权重矩阵的范数，从而实现参数范数的均匀化。与传统的初始化方法相比，WeSaR不是一次性地设置参数的初始值，而是在训练过程中动态地调整参数的范数，从而更好地适应模型的训练过程。

关键设计：WeSaR的关键设计在于门参数的初始化和更新策略。门参数的初始值需要 carefully 选择，以保证模型在初始状态下具有良好的性能。门参数的更新策略也需要 carefully 设计，以保证门参数能够有效地调整权重矩阵的范数，并且不会引入额外的训练难度。论文中具体使用了Adam优化器来更新门参数，并设置了合适的学习率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，WeSaR在1.3亿、13亿和130亿参数的Transformer解码器上均表现出优越的性能。与包括常见的初始化方法在内的基线方法相比，WeSaR能够显著减少损失尖峰的出现，加速训练过程，并最终获得更好的模型性能。具体性能提升数据未知，但结论是优于现有方法。

🎯 应用场景

该研究成果可应用于各种基于Transformer的大语言模型的预训练，尤其是在模型规模较大、训练数据复杂的情况下，能够有效提升训练的稳定性和效率，降低训练成本。此外，该方法也可以推广到其他类型的神经网络模型中，以解决类似的训练不稳定问题。

📄 摘要（原文）

Loss spikes, a phenomenon in which the loss value diverges suddenly, is a fundamental issue in the pre-training of large language models. This paper supposes that the non-uniformity of the norm of the parameters is one of the causes of loss spikes. Here, in training of neural networks, the scale of the gradients is required to be kept constant throughout the layers to avoid the vanishing and exploding gradients problem. However, to meet these requirements in the Transformer model, the norm of the model parameters must be non-uniform, and thus, parameters whose norm is smaller are more sensitive to the parameter update. To address this issue, we propose a novel technique, weight scaling as reparameterization (WeSaR). WeSaR introduces a gate parameter per parameter matrix and adjusts it to the value satisfying the requirements. Because of the gate parameter, WeSaR sets the norm of the original parameters uniformly, which results in stable training. Experimental results with the Transformer decoders consisting of 130 million, 1.3 billion, and 13 billion parameters showed that WeSaR stabilizes and accelerates training and that it outperformed compared methods including popular initialization methods.