大多數語言模型都是在短文本序列上進行訓練的,然後透過額外的長文件訓練(稱為上下文擴充套件)來處理更長的輸入。由於Llama 3的預訓練資料是專有的,很難判斷其擴充套件的易用性來自架構還是訓練資料。研究者假設相同的擴充套件方法可以遷移到其他架構,但OlmPool的研究表明,事實往往並非如此。
OlmPool是艾倫人工智慧研究所(Ai2)開發的受控模型套件,包含26個7B模型,所有模型在相同資料上預訓練1400億個token,然後使用相同的長上下文資料混合和過程擴充套件到64K上下文。唯一變化的是架構。研究聚焦於四個影響注意力機制的架構選擇:QK歸一化、分組查詢注意力(GQA)、滑動視窗注意力和預訓練上下文長度。
研究發現,這些選擇單獨作用時影響溫和:QK歸一化影響最大,移除它並切換歸一化順序可在HELMET基準上帶來6分提升;GQA和較短的預訓練上下文長度各自造成較小下降;滑動視窗注意力單獨影響約1分。但當它們組合時,效果遠超各部分之和——例如,在已有GQA的模型上新增滑動視窗注意力平均導致9分下降。最差的配置組合了兩種或更多限制注意力靈活性的選擇。
更關鍵的是,標準訓練訊號幾乎無法預測長上下文效能。訓練損失、驗證困惑度和16個短上下文基準都無法識別哪些模型在32K或64K上下文中表現更好。即使在同一基準的8K分割上,分數也無法預測擴充套件後兩位數的波動。擴充套件後表現看似相同的模型,在HELMET上可能相差26分以上。
研究還發現,Llama 3的架構組合表現強勁,但並非在所有情況下最優——其他幾種模型明顯超越它。這表明Llama 3的長上下文成功主要源於架構,驗證於Llama的擴充套件方法可能需要針對其他模型族進行調整。此外,架構導致的差距不會隨著更多資料而消失:即使經過500億token的上下文擴充套件(佔訓練總量的26%),最差架構仍無法達到Llama架構在10億token後的效能。
注意力模式分析揭示了原因:沒有QK歸一化的模型會產生更強的注意力匯聚——輸入早期位置持續獲得大量關注,即使它們與當前預測無關。更強的匯聚與更好的長上下文表現相關,表明它是模型在沒有QK歸一化時支援長距離檢索的預設策略。在針在該測試中,有QK歸一化的模型對目標資訊的關注度較低,與其較弱的長上下文表現一致。
OlmPool套件已完整發布,包含所有26個模型在預訓練和上下文擴充套件各階段的38個檢查點。這項研究強調,架構選擇的組合可能產生遠低於預期的長上下文效能,且這種結果無法從標準訓練訊號中預知。研究者希望該資源能幫助開發更好的上下文擴充套件方法,並研究早期預訓練中的其他現象。