Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model

TL;DR

This paper introduces WTA-CRS, an unbiased estimator for matrix multiplication that reduces memory usage during transformer fine-tuning, enabling larger batch sizes and faster training with minimal accuracy loss.

Contribution

The paper proposes WTA-CRS, a novel unbiased estimator for matrix multiplication that reduces variance and memory consumption in transformer training.

Findings

Achieves up to 2.7× memory reduction with minimal accuracy loss.

Enables up to 6.4× larger batch sizes during training.

Improves downstream task performance with larger models and faster training.

Abstract

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.