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可通过部署更大模型和/或采用更大批量实现更快训练速度，从而获得更优的下游任务性能。