[論文レビュー] Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup

Luyu Gao, Yunyi Zhang+2 more|arXiv (Cornell University)|Jan 18, 2021

Domain Adaptation and Few-Shot Learning参考文献 16被引用数 5

ひとこと要約

この論文では、エンコーダの誤差逆伝播をミニバッチ単位の対照的損失から分離することで、メモリ制限のあるGPU上で任意の大きなバッチサイズでの学習を可能にするGradient Cacheという手法を紹介している。この手法により、表現の勾配をキャッシュすることで、メモリ使用量をバッチサイズにかかわらず一定に保ちつつ、1台のコンsumerクラスGPUでSOTA性能を達成できる。バッチサイズが大きくても性能が低下せず、複数のハイエンドGPUで学習したモデルと同等の性能を発揮するが、実行時間は20%程度長くなる。

ABSTRACT

Contrastive learning has been applied successfully to learn vector representations of text. Previous research demonstrated that learning high-quality representations benefits from batch-wise contrastive loss with a large number of negatives. In practice, the technique of in-batch negative is used, where for each example in a batch, other batch examples' positives will be taken as its negatives, avoiding encoding extra negatives. This, however, still conditions each example's loss on all batch examples and requires fitting the entire large batch into GPU memory. This paper introduces a gradient caching technique that decouples backpropagation between contrastive loss and the encoder, removing encoder backward pass data dependency along the batch dimension. As a result, gradients can be computed for one subset of the batch at a time, leading to almost constant memory usage.

研究の動機と目的

• 対照的学習において、特にリソースが限られた研究者にとってバッチサイズを制限するGPUのメモリ制限を克服すること。
• 高品質な表現学習に不可欠な大規模バッチでの学習を、1台のコンsumerクラスGPUで可能にすること。
• 近似を用いず、完全バッチ学習と同等の正確な勾配更新を維持することで、マルチGPU環境と同等のモデル性能を確保すること。
• 学術的・リソース制限のある環境でも利用可能な大規模バッチ対照的学習へのアクセスを民主化すること。

提案手法

• 2段階の誤差逆伝播プロセスを導入：まず損失関数から表現への勾配を計算し、次に表現からモデルパラメータへの勾配を計算する。
• 表現に関する損失の勾配（表現勾配キャッシュ）をキャッシュすることで、エンコーダの更新におけるバッチ依存性を解消する。
• 表現のミニバッチごとにエンコーダの最適化を部分的な更新として行い、固定されたGPUメモリ内で大規模な有効バッチサイズを実現する。
• パラメトリックな類似度関数に対しては、別個の距離勾配キャッシュを導入することで、深い距離関数に対しても同様の原理を適用する。
• キャッシュされた勾配を用いて、全バッチの正確な勾配をメモリ効率よく計算し、全バッチの活性化値をすべて保持する必要がなくなる。
• 類似度関数Φ(·,·)を深くする場合にも、距離勾配と表現勾配の両方をキャッシュすることで、同時最適化を可能にする。

実験結果

リサーチクエスチョン

• RQ1 1台のGPU上で、メモリ制限を超えない範囲で大規模バッチでの対照的モデルの学習は可能か？
• RQ2 エンコーダの誤差逆伝播を対照的損失から分離することで、完全バッチ学習と同等の正確な勾配更新が可能か？
• RQ3 提案手法により、1台のコンsumer GPUを用いて、密度的レトリーバベンチマークでSOTA性能を再現できるか？
• RQ4 勾配蓄積法や逐次学習と比較して、本手法のトレーニング速度とメモリ使用量のスケーリング特性はいかがなっているか？
• RQ5 シンプルなドット積類似度を超えた深い距離関数に対しても、本手法は一般化可能か？

主な発見

• Gradient Cacheにより、1台のRTX 2080tiでバッチサイズ512での学習が可能となり、8台のV100 GPUで学習したモデルと同等の性能を達成した。
• 20%程度の追加実行時間で、MS MARCOパッセージランク付けベンチマークでSOTA結果を達成し、標準的および勾配蓄積ベースラインを上回った。
• バッチサイズにかかわらずピークメモリ使用量が一定を保つため、ハードウェアの限界を超えてバッチサイズを拡大可能である。
• 勾配蓄積法は、1更新あたりのインバッチネガティブ例が少ないので性能が劣り、逐次学習は高い分散と低い精度を示した。
• 二重キャッシュ機構により、深い距離関数に対しても成功裏に一般化され、類似度および表現パラメータの両方の正確な勾配計算を維持した。