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

要約

モデルのサイズが急速に増大するにつれて、メモリ使用量が膨大になるため、大規模な事前トレーニング済み言語モデルを微調整することがますます困難になってきています。
以前の研究は通常、ネットワーク内のトレーニング可能なパラメータの数を減らすことに焦点を当てていました。
モデル パラメーターはメモリ使用量に影響しますが、トレーニング中の主なメモリ ボトルネックは、勾配計算に不可欠な特徴マップ (アクティベーションとも呼ばれます) の保存から発生します。
特に、ニューラル ネットワークは通常、確率的勾配降下法を使用してトレーニングされます。
確率的最適化では、勾配推定量が適度な分散で偏っていない限り、モデルはノイズの多い勾配を処理できると主張します。
この動機に従って、我々は、分散を低減した行列生成のために、WTA-CRS と呼ばれる不偏推定量の新しいファミリーを提案します。これは、勾配を計算するためにサブサンプリングされた活性化を保存することのみを必要とします。
私たちの研究は、変圧器を調整するという文脈において、私たちが提案する推定量が既存の推定量と比較してより低い分散を示すという理論的および実験的証拠の両方を提供します。
線形演算をトランスフォーマーの近似演算に置き換えることで、精度をほとんど低下させることなく最大 2.7$\times$ のピーク メモリ削減を達成し、最大 $6.4\times$ 大きなバッチ サイズを実現できます。
同じハードウェアの下で、WTA-CRS は、より大きなモデルを適用したり、より大きなバッチ サイズでより高速なトレーニング速度を適用したりすることで、ダウンストリーミング タスクのパフォーマンスを向上させることができます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google