With the rapid growth in model size, fine-tuning the large pre-trained language model has become increasingly difficult due to its extensive memory usage. Previous works usually focus on reducing the number of trainable parameters in the network. While the model parameters do contribute to memory usage, the primary memory bottleneck during training arises from storing feature maps, also known as activations, as they are crucial for gradient calculation. Notably, neural networks are usually trained using stochastic gradient descent. We argue that in stochastic optimization, models can handle noisy gradients as long as the gradient estimator is unbiased with reasonable variance. Following this motivation, we propose a new family of unbiased estimators called WTA-CRS, for matrix production with reduced variance, which only requires storing the sub-sampled activations for calculating the gradient. Our work provides both theoretical and experimental evidence that, in the context of tuning transformers, our proposed estimators exhibit lower variance compared to existing ones. By replacing the linear operation with our approximated one in transformers, we can achieve up to 2.7$\times$ peak memory reduction with almost no accuracy drop and enables up to $6.4\times$ larger batch size. Under the same hardware, WTA-CRS enables better down-streaming task performance by applying larger models and/or faster training speed with larger batch sizes.

翻译：随着模型规模的快速增长，由于内存消耗巨大，微调大型预训练语言模型变得越来越困难。以往的研究通常侧重于减少网络中可训练参数的数量。尽管模型参数确实会占用内存，但训练过程中内存的主要瓶颈源于存储特征图，即激活值，因为它们在梯度计算中至关重要。值得注意的是，神经网络通常使用随机梯度下降进行训练。我们认为，在随机优化中，只要梯度估计器无偏且方差合理，模型能够处理噪声梯度。基于这一动机，我们提出了一类新的无偏估计器WTA-CRS，用于降低方差的矩阵乘积计算，该方法仅需存储子采样后的激活值来计算梯度。我们的工作提供了理论和实验证据，表明在微调Transformer的背景下，所提出的估计器相比现有方法具有更低的方差。通过用我们的近似估计替换Transformer中的线性运算，我们能够实现高达2.7倍的内存峰值降低，且准确率几乎无下降，并支持高达6.4倍更大的批量大小。在相同硬件条件下，WTA-CRS通过应用更大模型和/或更大批量实现更快训练速度，从而提升下游任务性能。