400萬token新SOTA！英偉達UIUC聯手：兼顧長短上下文頂尖效能

2025-05-01 06:10 新智元

新智元報道  

編輯：犀牛
【新智元導讀】來自英偉達和UIUC的華人團隊提出一種高效訓練方法，將LLM上下文長度從128K擴充套件至驚人的400萬token SOTA紀錄！基於Llama3.1-Instruct打造的UltraLong-8B模型，不僅在長上下文基準測試中表現卓越，還在標準任務中保持頂尖競爭力。
大語言模型（LLM）在文字和多模態任務上已經展現出驚豔的表現。
像是最新的Gemini 2.5 Pro在文字及程式碼上的頂尖效能，以及GPT-4o的原生生圖能力都很好的證明了這點。
然而，很多實際應用場景，比如長文件和影片理解、上下文學習以及推理時擴充套件，都需要模型能夠處理超長的token序列。
在這些場景中，模型的上下文視窗受限往往成為一大瓶頸，因為分佈在長文件中的關鍵資訊可能會被忽略。
為了解決這些問題，來自英偉達和UIUC的研究者提出了一種高效的訓練方法。
這種方法可以從現有的指令微調模型出發，構建超長上下文的LLM，最高可將上下文長度推向400萬token的極限！
論文地址：https://arxiv.org/pdf/2504.06214
研究人員利用上面方法訓練的UltraLong-8B模型在長上下文任務上達到了頂尖水平，同時在標準任務上也保持了競爭力。
主要貢獻：
  • 高效且可擴充套件的訓練方法。
  • 關鍵技術創新：研究者引入了特殊文件分隔符和基於YaRN的位置編碼擴充套件技術，透過消融實驗證明這些技術對長上下文建模至關重要。
  • 高效的單步預訓練策略：研究者發現，相比多步擴充套件方法，單步持續預訓練在上下文擴充套件上更高效，在合成和真實世界長上下文基準測試中始終表現出色。
  • 全面的實驗驗證：研究者在多個基準測試上進行了廣泛實驗，包括RULER、LV-Eval、InfiniteBench、MMLU、MMLU-Pro、MATH、GSM-8K和HumanEval，證明UltraLong-8B模型在長上下文和標準任務上均優於現有基線。

實驗方法
如圖1所示，本文方法主要分為兩個階段：持續預訓練和指令微調。
以Llama 3.1-8B-Instruct為基礎，持續預訓練階段將模型的上下文視窗逐步擴充套件到目標長度（比如100萬、200萬、400萬token）。隨後，指令微調階段最佳化模型的指令遵循能力和推理能力。
這兩個階段結合，讓模型既能高效處理超長輸入，又能在長短上下文任務中表現出色。
第一階段透過持續預訓練擴充套件模型的上下文視窗，採用特殊文件分隔符和基於YaRN的縮放技術來處理超長序列。第二階段使用精心挑選的資料集進行指令微調，提升模型的指令遵循和推理能力

持續預訓練：擴充套件上下文長度
在第一階段，研究者透過持續預訓練將Llama-3.1-8B-Instruct的上下文視窗擴充套件到目標長度。
研究者對少於4000 token的短文件進行下采樣，對超過8000 token的長文件進行上取樣，最終形成一個包含10億token的語料庫。
這些文件被拼接成對應目標上下文長度的更長序列（比如100萬、200萬、400萬token）。拼接時，他們使用特殊字元分隔不同文件，而不是用保留的開始和結束標記。
此外，在持續預訓練中，研究人員沒有使用跨文件注意力掩碼，從而允許模型關注整個輸入序列。
為了支援超長上下文，研究人員採用了基於YaRN的縮放方法，而不是之前工作中常用的NTK感知縮放策略。他們固定超引數α=1和β=4，並根據目標上下文長度計算縮放因子s。
當輸入長度接近最大限制時，Llama-3.1模型的效能會下降。為解決這個問題，他們為RoPE嵌入採用了更大的縮放因子，從而更好地適應超長序列。
研究者針對三種上下文長度（100萬、200萬和400萬token）構建了長上下文模型，並將RoPE縮放因子分別設定為128、256和512。
每個模型在10億token的語料上訓練一個epoch，學習率為3×10⁻⁵。
為了提升訓練的可擴充套件性，他們使用了Megatron-LM框架。為了處理超長輸入序列，採用了張量並行和上下文並行。
訓練在256個NVIDIA H100 GPU上進行，1M、2M和4M模型的訓練時間分別約為5小時、6小時和13小時。

指令微調
在第二階段，研究者透過監督微調（SFT）提升長上下文模型的指令遵循和推理能力，使用的是一些精心挑選的資料集。
他們整合並優化了多個開源SFT資料集，覆蓋三個關鍵領域：通用領域、數學和程式碼。
為了進一步提升SFT資料集的質量，他們利用GPT-4o和4o-mini優化了這些資料集的回答內容。
值得注意的是，研究者的SFT資料集僅包含上述短上下文資料（少於8000 token的樣本），沒有加入合成長上下文指令資料。
他們發現，僅依靠短上下文資料就足以取得優異效果，這與之前研究的觀察一致。
最終，研究者構建了一個包含10萬個樣本的SFT資料集。對於每種目標上下文長度的模型，他們使用128的批大小和5×10⁻⁶的學習率。
訓練依然基於Megatron-LM 框架，在256個NVIDIA H100 GPU上進行，張量並行度設為tp=8。每次訓練大約需要 30 分鐘。

基線模型與評估基準
研究者將他們的模型與基於Llama家族的最先進（SOTA）長上下文模型進行對比，以確保對訓練方法公平且可控的評估。
  • Llama-3.1 (Llama-3.1-8B-Instruct)：這是他們的基礎模型，支援128K的上下文視窗。
  • ProLong (Llama-3-8B-ProLong-512k-Instruct)：基於Llama-3構建的長上下文模型，擁有512K的上下文視窗。
  • Gradient (Llama-3-8B-Instruct-Gradient-1048k)：另一個基於Llama的長上下文模型，支援高達1M的上下文視窗。
本文研究者專注於Llama家族的模型，這樣可以更清晰地展示他們擴充套件上下文長度訓練方法的有效性，同時確保在標準任務上的效能依然具有競爭力。
他們透過以下基準測試來評估模型的長上下文能力：
  • RULER：這是一個專門評估長上下文語言模型的基準，透過生成不同序列長度的合成樣本，覆蓋四個任務類別。
  • LV-Eval：這是一個長上下文基準，包含最高256K token的五個長度級別，重點測試兩種任務：單跳問答（single-hop QA）和多跳問答（multi-hop QA）。
  • InfiniteBench：這是一個長上下文基準，平均輸入長度約200K token，最大長度超過2M token，包含合成任務和現實世界任務。

實驗結果
研究人員首先從「大海撈針」（Needle in a Haystack，NIAH）這一測試開始，然後再探討長上下文和標準基準的評估。
研究人員透過NIAH密碼檢索測試，評估模型在長上下文檢索方面的能力。在這個任務中，模型需要在一大段毫無意義的文字中，找到一個簡單密碼，比如一個隨機的六位數字。
為了量化檢索的準確性，他們測試了40種不同的輸入序列長度。對於每種長度，密碼會被隨機插入到10個均勻分佈的文件深度中。
結果如圖2所示。對於本文的模型，測試了高達100萬、200萬和400萬個token的輸入長度；而對於基準模型，只測試了最高100萬個token。
如圖2a到2c所示，在基準模型中，只有Llama-3-8B-Instruct和Gradient-1048k通過了NIAH測試，而Llama-3.1-8B-Instruct和Llama-3-8B-ProLong-512k-Instruct即使在它們聲稱的上下文長度內也出現了錯誤。
相比之下，如圖2d到2f所示，研究者的超長（UltraLong）模型在所有輸入長度和深度上都達到了100%的準確率，展現了強大的長上下文檢索能力。
研究者在RULER、LV-Eval和InfiniteBench上的評估結果如表1所示。加粗的數字表示效能超過了所有基準模型。
總體來說，他們的三個模型在大多數情況下都取得了最高分。
在RULER基準測試中，UltraLong模型在512K和100萬個token的輸入長度上表現最佳。在LV-Eval中，他們的模型在128K和256K token長度內的平均F1分數最高。
此外，他們在InfiniteBench上也取得了最佳表現。
這些結果表明，研究者的訓練方法有效擴充套件了語言模型的上下文視窗到超長輸入，同時保持了原有輸入長度的效能。
相比之下，基準模型中，Llama-3.1是為128K輸入長度設計的，當輸入超過128K token時，效能顯著下降。ProLong是為512K上下文設計的，但即使它訓練了更多token（410億對比10億），在512K長度上的表現也不如他們的模型。
Gradient是基準模型中支援最長上下文的（100萬個token），但在LV-Eval和InfiniteBench上的表現較差，說明它的設計可能過於偏向人工任務，犧牲了現實任務的效果。
而本文的模型在人工（RULER）和混合（LV-Eval和InfiniteBench）基準測試中始終保持更高的分數，凸顯了方法的高效性和可擴充套件性。
研究者還透過通用、數學和程式碼領域的標準基準測試評估了模型，以確保擴充套件上下文長度不會影響短上下文任務的效能。
如表2所示，他們的模型效能與基礎模型Llama-3.1-8B-Instruct相當甚至更高，平均分數分別為62.47、61.06和60.95，而Llama-3.1-8B-Instruct為61.45。
特別值得一提的是，他們的模型在MMLU和MATH基準上表現出明顯提升，同時在GSM8K和HumanEval等其他基準上的表現也極具競爭力。
相比之下，基準長上下文模型Gradient和ProLong在這些標準任務上的效能大幅下降，平均分數僅為37.36和40.81。
這些結果表明，研究者的方法不僅有效擴充套件了上下文視窗，還保持甚至提升了模型的通用任務能力。
而Llama-3-8B-Instruct-Gradient-1048k和Llama-3-8B-ProLong-512k-Instruct的顯著效能下降，表明它們的超長上下文方法可能存在侷限性。

結論
在這項工作中，研究人員提出了一種高效且系統化的訓練方法，用於超長上下文語言模型，將上下文視窗擴充套件到100萬、200萬和400萬個token，同時在標準基準測試中保持了競爭力。
這種結合了高效的持續預訓練和指令微調，不僅提升了模型對長上下文的理解能力，還增強了其遵循指令的能力。
這一框架為可擴充套件的長上下文建模樹立了新標杆，也為未來在實際應用中提升長上下文效能的研究鋪平了道路。

作者介紹
Chejian Xu
伊利諾伊大學厄巴納-香檳分校（UIUC）計算機科學博士研究生，導師是Bo Li教授。浙江大學計算機科學與技術專業學士學位，就讀於CKC榮譽學院，導師是Shouling Ji教授和Siliang Tang教授。
專注於提升基礎模型的安全性、可靠性和一致性，包括LLMs、多模態模型以及基於LLM的智慧體。
Wei Ping
NVIDIA應用深度學習研究團隊的資深研究科學家，專注於大型語言模型和生成模型的研究。
加州大學歐文分校機器學習博士學位，熱衷於構建用於文字、音訊和多模態資料的尖端生成模型。此前，曾擔任百度矽谷人工智慧實驗室（由吳恩達創立）的文字到語音團隊負責人。
參考資料：
https://arxiv.org/abs/2504.06214


相關文章