Основні поняття
本研究提出了一種名為「對比累積」(CONTACCUM)的新方法,透過使用雙記憶庫結構來緩存查詢和文章表示,從而在記憶體受限的環境下,有效地訓練密集檢索器,並超越了現有的記憶體縮減方法,甚至優於高資源配置的訓練結果。
Анотація
論文資訊
Jaehee Kim, Yukyung Lee, & Pilsung Kang. (2024). A Gradient Accumulation Method for Dense Retriever under Memory Constraint. Thirty-eighth Conference on Neural Information Processing Systems.
研究目標
本研究旨在解決密集檢索器訓練過程中,因 InfoNCE 損失函數需要大量負樣本而導致的高記憶體需求問題,並提出一個在低資源環境下,也能有效訓練密集檢索器的方法。
方法
本研究提出了「對比累積」(CONTACCUM)方法,該方法利用雙記憶庫結構來緩存過去生成的查詢和文章表示,並將其用於當前批次的訓練中,以增加負樣本的數量。具體來說,CONTACCUM 使用兩個先進先出(FIFO)隊列作為記憶庫,分別存儲查詢和文章表示。在每個訓練步驟中,模型會從記憶庫中提取表示,並將其與當前批次的表示組合,構建一個更大的相似度矩陣,用於計算 InfoNCE 損失。
主要發現
- CONTACCUM 在低資源環境下(例如,僅有 11GB 記憶體),其效能優於高資源環境下(80GB 記憶體)的密集檢索器訓練結果。
- CONTACCUM 的訓練速度比其他記憶體縮減方法(例如,GradCache)更快。
- 相比於僅使用文章記憶庫的方法,CONTACCUM 的雙記憶庫結構有效緩解了梯度範數不平衡問題,提高了訓練穩定性。
主要結論
CONTACCUM 是一種有效且高效的密集檢索器訓練方法,尤其適用於記憶體受限的環境。透過利用雙記憶庫和梯度累積技術,CONTACCUM 能夠在不犧牲效能的情況下,顯著降低記憶體需求,並提高訓練穩定性。
研究意義
本研究對於推動密集檢索器在資源受限環境下的應用具有重要意義,例如個人電腦、行動設備等。此外,CONTACCUM 的雙記憶庫結構和梯度範數平衡分析,為其他基於對比學習的模型訓練提供了新的思路。
局限與未來研究方向
- 本研究主要關注於監督式微調階段,未來將探索 CONTACCUM 在預訓練階段的應用。
- CONTACCUM 仍然依賴於計算成本較高的 softmax 操作,未來將研究更有效的訓練策略以降低計算複雜度。
Статистика
CONTACCUM 在僅有 11GB 記憶體的環境下,其效能超越了 80GB 記憶體環境下訓練的密集檢索器。
在 NQ 資料集上,當總批次大小為 512 時,GradCache 的訓練速度比 GradAccum 慢 93%,而 CONTACCUM 僅慢 26%。
當不使用查詢記憶庫時,文章編碼器的梯度範數最高可達查詢編碼器的 30 倍。
Цитати
"Empirical studies have shown that the dense retriever performs better with large batch [28, 43, 42]."
"In this study, we propose Contrastive Accumulation (CONTACCUM), which demonstrates high performance and stable training under memory constraints."
"CONTACCUM leverages previously generated query and passage representations through a memory bank, enabling the use of more negative samples."