新智元報(bào)道
編輯：犀牛
【新智元導(dǎo)讀】來(lái)自英偉達(dá)和UIUC的華人團(tuán)隊(duì)提出一種高效訓(xùn)練方法，將LLM上下文長(zhǎng)度從128K擴(kuò)展至驚人的400萬(wàn)token SOTA紀(jì)錄！基于Llama3.1-Instruct打造的UltraLong-8B模型，不僅在長(zhǎng)上下文基準(zhǔn)測(cè)試中表現(xiàn)卓越，還在標(biāo)準(zhǔn)任務(wù)中保持頂尖競(jìng)爭(zhēng)力。
大語(yǔ)言模型（LLM）在文本和多模態(tài)任務(wù)上已經(jīng)展現(xiàn)出驚艷的表現(xiàn)。
像是最新的Gemini 2.5 Pro在文本及代碼上的頂尖性能，以及GPT-4o的原生生圖能力都很好的證明了這點(diǎn)。
然而，很多實(shí)際應(yīng)用場(chǎng)景，比如長(zhǎng)文檔和視頻理解、上下文學(xué)習(xí)以及推理時(shí)擴(kuò)展，都需要模型能夠處理超長(zhǎng)的token序列。
在這些場(chǎng)景中，模型的上下文窗口受限往往成為一大瓶頸，因?yàn)榉植荚陂L(zhǎng)文檔中的關(guān)鍵信息可能會(huì)被忽略。
為了解決這些問(wèn)題，來(lái)自英偉達(dá)和UIUC的研究者提出了一種高效的訓(xùn)練方法。
這種方法可以從現(xiàn)有的指令微調(diào)模型出發(fā)，構(gòu)建超長(zhǎng)上下文的LLM，最高可將上下文長(zhǎng)度推向400萬(wàn)token的極限！
論文地址：https://arxiv.org/pdf/2504.06214
研究人員利用上面方法訓(xùn)練的UltraLong-8B模型在長(zhǎng)上下文任務(wù)上達(dá)到了頂尖水平，同時(shí)在標(biāo)準(zhǔn)任務(wù)上也保持了競(jìng)爭(zhēng)力。
主要貢獻(xiàn)：
高效且可擴(kuò)展的訓(xùn)練方法。
關(guān)鍵技術(shù)創(chuàng)新：研究者引入了特殊文檔分隔符和基于YaRN的位置編碼擴(kuò)展技術(shù)，通過(guò)消融實(shí)驗(yàn)證明這些技術(shù)對(duì)長(zhǎng)上下文建模至關(guān)重要。
高效的單步預(yù)訓(xùn)練策略：研究者發(fā)現(xiàn)，相比多步擴(kuò)展方法，單步持續(xù)預(yù)訓(xùn)練在上下文擴(kuò)展上更高效，在合成和真實(shí)世界長(zhǎng)上下文基準(zhǔn)測(cè)試中始終表現(xiàn)出色。
全面的實(shí)驗(yàn)驗(yàn)證：研究者在多個(gè)基準(zhǔn)測(cè)試上進(jìn)行了廣泛實(shí)驗(yàn)，包括RULER、LV-Eval、InfiniteBench、MMLU、MMLU-Pro、MATH、GSM-8K和HumanEval，證明UltraLong-8B模型在長(zhǎng)上下文和標(biāo)準(zhǔn)任務(wù)上均優(yōu)于現(xiàn)有基線。
實(shí)驗(yàn)方法
如圖1所示，本文方法主要分為兩個(gè)階段：持續(xù)預(yù)訓(xùn)練和指令微調(diào)。
以Llama 3.1-8B-Instruct為基礎(chǔ)，持續(xù)預(yù)訓(xùn)練階段將模型的上下文窗口逐步擴(kuò)展到目標(biāo)長(zhǎng)度（比如100萬(wàn)、200萬(wàn)、400萬(wàn)token）。隨后，指令微調(diào)階段優(yōu)化模型的指令遵循能力和推理能力。
這兩個(gè)階段結(jié)合，讓模型既能高效處理超長(zhǎng)輸入，又能在長(zhǎng)短上下文任務(wù)中表現(xiàn)出色。
第一階段通過(guò)持續(xù)預(yù)訓(xùn)練擴(kuò)展模型的上下文窗口，采用特殊文檔分隔符和基于YaRN的縮放技術(shù)來(lái)處理超長(zhǎng)序列。第二階段使用精心挑選的數(shù)據(jù)集進(jìn)行指令微調(diào)，提升模型的指令遵循和推理能力
持續(xù)預(yù)訓(xùn)練：擴(kuò)展上下文長(zhǎng)度
在第一階段，研究者通過(guò)持續(xù)預(yù)訓(xùn)練將Llama-3.1-8B-Instruct的上下文窗口擴(kuò)展到目標(biāo)長(zhǎng)度。
研究者對(duì)少于4000 token的短文檔進(jìn)行下采樣，對(duì)超過(guò)8000 token的長(zhǎng)文檔進(jìn)行上采樣，最終形成一個(gè)包含10億token的語(yǔ)料庫(kù)。
這些文檔被拼接成對(duì)應(yīng)目標(biāo)上下文長(zhǎng)度的更長(zhǎng)序列（比如100萬(wàn)、200萬(wàn)、400萬(wàn)token）。拼接時(shí)，他們使用特殊字符分隔不同文檔，而不是用保留的開(kāi)始和結(jié)束標(biāo)記。
此外，在持續(xù)預(yù)訓(xùn)練中，研究人員沒(méi)有使用跨文檔注意力掩碼，從而允許模型關(guān)注整個(gè)輸入序列。
為了支持超長(zhǎng)上下文，研究人員采用了基于YaRN的縮放方法，而不是之前工作中常用的NTK感知縮放策略。他們固定超參數(shù)α=1和β=4，并根據(jù)目標(biāo)上下文長(zhǎng)度計(jì)算縮放因子s。
當(dāng)輸入長(zhǎng)度接近最大限制時(shí)，Llama-3.1模型的性能會(huì)下降。為解決這個(gè)問(wèn)題，他們?yōu)镽oPE嵌入采用了更大的縮放因子，從而更好地適應(yīng)超長(zhǎng)序列。
研究者針對(duì)三種上下文長(zhǎng)度（100萬(wàn)、200萬(wàn)和400萬(wàn)token）構(gòu)建了長(zhǎng)上下文模型，并將RoPE縮放因子分別設(shè)置為128、256和512。
每個(gè)模型在10億token的語(yǔ)料上訓(xùn)練一個(gè)epoch，學(xué)習(xí)率為3×10??。
為了提升訓(xùn)練的可擴(kuò)展性，他們使用了Megatron-LM框架。為了處理超長(zhǎng)輸入序列，采用了張量并行和上下文并行。
訓(xùn)練在256個(gè)NVIDIA H100 GPU上進(jìn)行，1M、2M和4M模型的訓(xùn)練時(shí)間分別約為5小時(shí)、6小時(shí)和13小時(shí)。
指令微調(diào)
在第二階段，研究者通過(guò)監(jiān)督微調(diào)（SFT）提升長(zhǎng)上下文模型的指令遵循和推理能力，使用的是一些精心挑選的數(shù)據(jù)集。
他們整合并優(yōu)化了多個(gè)開(kāi)源SFT數(shù)據(jù)集，覆蓋三個(gè)關(guān)鍵領(lǐng)域：通用領(lǐng)域、數(shù)學(xué)和代碼。
為了進(jìn)一步提升SFT數(shù)據(jù)集的質(zhì)量，他們利用GPT-4o和4o-mini優(yōu)化了這些數(shù)據(jù)集的回答內(nèi)容。
值得注意的是，研究者的SFT數(shù)據(jù)集僅包含上述短上下文數(shù)據(jù)（少于8000 token的樣本），沒(méi)有加入合成長(zhǎng)上下文指令數(shù)據(jù)。
他們發(fā)現(xiàn)，僅依靠短上下文數(shù)據(jù)就足以取得優(yōu)異效果，這與之前研究的觀察一致。
最終，研究者構(gòu)建了一個(gè)包含10萬(wàn)個(gè)樣本的SFT數(shù)據(jù)集。對(duì)于每種目標(biāo)上下文長(zhǎng)度的模型，他們使用128的批大小和5×10??的學(xué)習(xí)率。
訓(xùn)練依然基于Megatron-LM 框架，在256個(gè)NVIDIA H100 GPU上進(jìn)行，張量并行度設(shè)為tp=8。每次訓(xùn)練大約需要 30 分鐘。
基線模型與評(píng)估基準(zhǔn)
研究者將他們的模型與基于Llama家族的最先進(jìn)（SOTA）長(zhǎng)上下文模型進(jìn)行對(duì)比，以確保對(duì)訓(xùn)練方法公平且可控的評(píng)估。
Llama-3.1 (Llama-3.1-8B-Instruct)：這是他們的基礎(chǔ)模型，支持128K的上下文窗口。
ProLong (Llama-3-8B-ProLong-512k-Instruct)：基于Llama-3構(gòu)建的長(zhǎng)上下文模型，擁有512K的上下文窗口。
Gradient (Llama-3-8B-Instruct-Gradient-1048k)：另一個(gè)基于Llama的長(zhǎng)上下文模型，支持高達(dá)1M的上下文窗口。
本文研究者專注于Llama家族的模型，這樣可以更清晰地展示他們擴(kuò)展上下文長(zhǎng)度訓(xùn)練方法的有效性，同時(shí)確保在標(biāo)準(zhǔn)任務(wù)上的性能依然具有競(jìng)爭(zhēng)力。
他們通過(guò)以下基準(zhǔn)測(cè)試來(lái)評(píng)估模型的長(zhǎng)上下文能力：
RULER：這是一個(gè)專門評(píng)估長(zhǎng)上下文語(yǔ)言模型的基準(zhǔn)，通過(guò)生成不同序列長(zhǎng)度的合成樣本，覆蓋四個(gè)任務(wù)類別。
LV-Eval：這是一個(gè)長(zhǎng)上下文基準(zhǔn)，包含最高256K token的五個(gè)長(zhǎng)度級(jí)別，重點(diǎn)測(cè)試兩種任務(wù)：?jiǎn)翁鴨?wèn)答（single-hop QA）和多跳問(wèn)答（multi-hop QA）。
InfiniteBench：這是一個(gè)長(zhǎng)上下文基準(zhǔn)，平均輸入長(zhǎng)度約200K token，最大長(zhǎng)度超過(guò)2M token，包含合成任務(wù)和現(xiàn)實(shí)世界任務(wù)。
實(shí)驗(yàn)結(jié)果
研究人員首先從「大海撈針」（Needle in a Haystack，NIAH）這一測(cè)試開(kāi)始，然后再探討長(zhǎng)上下文和標(biāo)準(zhǔn)基準(zhǔn)的評(píng)估。
研究人員通過(guò)NIAH密碼檢索測(cè)試，評(píng)估模型在長(zhǎng)上下文檢索方面的能力。在這個(gè)任務(wù)中，模型需要在一大段毫無(wú)意義的文本中，找到一個(gè)簡(jiǎn)單密碼，比如一個(gè)隨機(jī)的六位數(shù)字。
為了量化檢索的準(zhǔn)確性，他們測(cè)試了40種不同的輸入序列長(zhǎng)度。對(duì)于每種長(zhǎng)度，密碼會(huì)被隨機(jī)插入到10個(gè)均勻分布的文檔深度中。
結(jié)果如圖2所示。對(duì)于本文的模型，測(cè)試了高達(dá)100萬(wàn)、200萬(wàn)和400萬(wàn)個(gè)token的輸入長(zhǎng)度；而對(duì)于基準(zhǔn)模型，只測(cè)試了最高100萬(wàn)個(gè)token。
如圖2a到2c所示，在基準(zhǔn)模型中，只有Llama-3-8B-Instruct和Gradient-1048k通過(guò)了NIAH測(cè)試，而Llama-3.1-8B-Instruct和Llama-3-8B-ProLong-512k-Instruct即使在它們聲稱的上下文長(zhǎng)度內(nèi)也出現(xiàn)了錯(cuò)誤。
相比之下，如圖2d到2f所示，研究者的超長(zhǎng)（UltraLong）模型在所有輸入長(zhǎng)度和深度上都達(dá)到了100%的準(zhǔn)確率，展現(xiàn)了強(qiáng)大的長(zhǎng)上下文檢索能力。
研究者在RULER、LV-Eval和InfiniteBench上的評(píng)估結(jié)果如表1所示。加粗的數(shù)字表示性能超過(guò)了所有基準(zhǔn)模型。
總體來(lái)說(shuō)，他們的三個(gè)模型在大多數(shù)情況下都取得了最高分。
在RULER基準(zhǔn)測(cè)試中，UltraLong模型在512K和100萬(wàn)個(gè)token的輸入長(zhǎng)度上表現(xiàn)最佳。在LV-Eval中，他們的模型在128K和256K token長(zhǎng)度內(nèi)的平均F1分?jǐn)?shù)最高。
此外，他們?cè)贗nfiniteBench上也取得了最佳表現(xiàn)。
這些結(jié)果表明，研究者的訓(xùn)練方法有效擴(kuò)展了語(yǔ)言模型的上下文窗口到超長(zhǎng)輸入，同時(shí)保持了原有輸入長(zhǎng)度的性能。
相比之下，基準(zhǔn)模型中，Llama-3.1是為128K輸入長(zhǎng)度設(shè)計(jì)的，當(dāng)輸入超過(guò)128K token時(shí)，性能顯著下降。ProLong是為512K上下文設(shè)計(jì)的，但即使它訓(xùn)練了更多token（410億對(duì)比10億），在512K長(zhǎng)度上的表現(xiàn)也不如他們的模型。
Gradient是基準(zhǔn)模型中支持最長(zhǎng)上下文的（100萬(wàn)個(gè)token），但在LV-Eval和InfiniteBench上的表現(xiàn)較差，說(shuō)明它的設(shè)計(jì)可能過(guò)于偏向人工任務(wù)，犧牲了現(xiàn)實(shí)任務(wù)的效果。
而本文的模型在人工（RULER）和混合（LV-Eval和InfiniteBench）基準(zhǔn)測(cè)試中始終保持更高的分?jǐn)?shù)，凸顯了方法的高效性和可擴(kuò)展性。
研究者還通過(guò)通用、數(shù)學(xué)和代碼領(lǐng)域的標(biāo)準(zhǔn)基準(zhǔn)測(cè)試評(píng)估了模型，以確保擴(kuò)展上下文長(zhǎng)度不會(huì)影響短上下文任務(wù)的性能。
如表2所示，他們的模型性能與基礎(chǔ)模型Llama-3.1-8B-Instruct相當(dāng)甚至更高，平均分?jǐn)?shù)分別為62.47、61.06和60.95，而Llama-3.1-8B-Instruct為61.45。
特別值得一提的是，他們的模型在MMLU和MATH基準(zhǔn)上表現(xiàn)出明顯提升，同時(shí)在GSM8K和HumanEval等其他基準(zhǔn)上的表現(xiàn)也極具競(jìng)爭(zhēng)力。
相比之下，基準(zhǔn)長(zhǎng)上下文模型Gradient和ProLong在這些標(biāo)準(zhǔn)任務(wù)上的性能大幅下降，平均分?jǐn)?shù)僅為37.36和40.81。
這些結(jié)果表明，研究者的方法不僅有效擴(kuò)展了上下文窗口，還保持甚至提升了模型的通用任務(wù)能力。
而Llama-3-8B-Instruct-Gradient-1048k和Llama-3-8B-ProLong-512k-Instruct的顯著性能下降，表明它們的超長(zhǎng)上下文方法可能存在局限性。
結(jié)論
在這項(xiàng)工作中，研究人員提出了一種高效且系統(tǒng)化的訓(xùn)練方法，用于超長(zhǎng)上下文語(yǔ)言模型，將上下文窗口擴(kuò)展到100萬(wàn)、200萬(wàn)和400萬(wàn)個(gè)token，同時(shí)在標(biāo)準(zhǔn)基準(zhǔn)測(cè)試中保持了競(jìng)爭(zhēng)力。
這種結(jié)合了高效的持續(xù)預(yù)訓(xùn)練和指令微調(diào)，不僅提升了模型對(duì)長(zhǎng)上下文的理解能力，還增強(qiáng)了其遵循指令的能力。
這一框架為可擴(kuò)展的長(zhǎng)上下文建模樹(shù)立了新標(biāo)桿，也為未來(lái)在實(shí)際應(yīng)用中提升長(zhǎng)上下文性能的研究鋪平了道路。
作者介紹
Chejian Xu
伊利諾伊大學(xué)厄巴納-香檳分校（UIUC）計(jì)算機(jī)科學(xué)博士研究生，導(dǎo)師是Bo Li教授。浙江大學(xué)計(jì)算機(jī)科學(xué)與技術(shù)專業(yè)學(xué)士學(xué)位，就讀于CKC榮譽(yù)學(xué)院，導(dǎo)師是Shouling Ji教授和Siliang Tang教授。
專注于提升基礎(chǔ)模型的安全性、可靠性和一致性，包括LLMs、多模態(tài)模型以及基于LLM的智能體。
Wei Ping
NVIDIA應(yīng)用深度學(xué)習(xí)研究團(tuán)隊(duì)的資深研究科學(xué)家，專注于大型語(yǔ)言模型和生成模型的研究。
加州大學(xué)歐文分校機(jī)器學(xué)習(xí)博士學(xué)位，熱衷于構(gòu)建用于文本、音頻和多模態(tài)數(shù)據(jù)的尖端生成模型。此前，曾擔(dān)任百度硅谷人工智能實(shí)驗(yàn)室（由吳恩達(dá)創(chuàng)立）的文本到語(yǔ)音團(tuán)隊(duì)負(fù)責(zé)人。
參考資料：
https://arxiv.org/abs/2504.06214
熱門跟貼