核心概念
本文提出了一種新的可擴展梯度和 HVP 草圖算法框架,專為現代硬件設計,並在訓練數據歸因、Hessian 谱分析和預訓練語言模型的內在維度計算等應用中展示了其強大功能,挑戰了關於預訓練語言模型的內在維度和 Hessian 特性的假設。
摘要
論文資訊
- 標題:用於訓練數據歸因和損失地形研究的高效草圖算法
- 作者:Andrea Schioppa
- 機構:Google DeepMind
- 會議:NeurIPS 2024
研究背景
現代機器學習模型的研究通常需要存儲大量的梯度或 Hessian 向量積 (HVP),傳統的草圖方法在這些内存限制下難以擴展。訓練數據歸因 (TDA) 和 Hessian 特徵值估計等應用需要存儲與網絡參數相同維度的向量,這導致了巨大的内存需求。
研究方法
本文提出了一種新的可擴展梯度和 HVP 草圖算法框架,專為現代硬件設計。具體來說,文章提出了三種新穎的算法:AFFD、AFJL 和 QK。這些算法通過以下設計選擇來解決現有算法的性能瓶頸:
主要發現
- AFFD 和 QK 算法在理論上被證明是有效的草圖算法,而 Fastfood Transform (FFD) 缺乏作為草圖算法的理論基礎。
- 顯式草圖算法在性能上優於隱式草圖算法。
- 修改預處理器可以顯著提高 GPU 性能。
- 對於生成任務,內在維度可能很大,這挑戰了先前關於內在維度的假設。
- 預訓練語言模型的 Hessian 谱可能與在較小網絡中觀察到的行為顯著不同。
研究結論
本文提出的草圖算法為模型分析提供了一個有效的工具包。這些方法揭示了重新思考基於層選擇的影響函數的必要性,生成任務的高內在維度,以及大型語言模型的 Hessian 谱與在較小網絡中觀察到的偏差。
研究限制
- 本文主要關注 Transformer 模型,結果可能因模型架構而異。
- 雖然本文展示了一個具有較大內在維度的生成任務示例,但仍需要對其進行深入研究。
統計資料
使用 AFFD 算法,目標維度 D = 2^20,可以為 BART、Roberta 和 GPT-2L 等預訓練語言模型構建 1k 維的 Krylov 子空間。
對於 TDA 分數,實現相關性 r ≥ 0.95 需要以下維度:FJL、FFD、AFFD:D = 2^10;AFJL:D = 2^12;QK:D = 2^14。
要達到 r ≥ 0.99,只需將每個維度增加 8 倍。
使用以下維度,相對平均絕對誤差低於 5%:AFFD:D = 2^10;AFJL:D = 2^12;QK:D = 2^13。
引述
"These limitations, both theoretical and practical, underscore the need for novel sketching algorithms."
"Our results suggest a need for non-linear projection algorithms to achieve lower intrinsic dimensionality for generative tasks."
"Crucially, we find that Hessian spectra in these models may deviate significantly from behaviors observed in smaller networks."