本文第一作者顧煜賢（https://t1101675.github.io/）為清華大學(xué)計(jì)算機(jī)系四年級直博生，師從黃民烈教授，研究方向?yàn)檎Z言模型的高效訓(xùn)練與推理方法。他曾在 ACL，EMNLP，ICLR 等會議和期刊上發(fā)表近 20 篇論文，多次進(jìn)行口頭報(bào)告，Google Scholar 引用數(shù) 2600+，曾獲 2025 年蘋果學(xué)者獎(jiǎng)學(xué)金。本篇論文為他在微軟亞洲研究院實(shí)習(xí)期間所完成。

近年來，大語言模型（LLMs）在自然語言理解、代碼生成與通用推理等任務(wù)上取得了顯著進(jìn)展，逐步成為通用人工智能的核心基石。然而，訓(xùn)練此類模型所需的計(jì)算資源和數(shù)據(jù)成本正以驚人的速度增長。面對高質(zhì)量語料日益枯竭、訓(xùn)練預(yù)算持續(xù)上升的雙重挑戰(zhàn)，如何以更少的資源實(shí)現(xiàn)更高效的學(xué)習(xí)，成為當(dāng)前語言模型發(fā)展的關(guān)鍵問題。

針對這一挑戰(zhàn)，清華大學(xué)、北京大學(xué)聯(lián)合微軟亞洲研究院，提出了一種全新的預(yù)訓(xùn)練數(shù)據(jù)選擇范式 ——PMP-based Data Selection（PDS）。該方法首次將數(shù)據(jù)選擇建模為一個(gè)最優(yōu)控制問題，并基于經(jīng)典的龐特里亞金最大值原理（PMP）推導(dǎo)出一組理論上的必要條件，為預(yù)訓(xùn)練階段中 “哪些數(shù)據(jù)更值得學(xué)” 提供了明確的數(shù)學(xué)刻畫。

在理論基礎(chǔ)之上，研究團(tuán)隊(duì)設(shè)計(jì)了可在大規(guī)模語料中高效運(yùn)行的 PDS 算法框架，并在多個(gè)模型規(guī)模和任務(wù)設(shè)置中進(jìn)行了系統(tǒng)驗(yàn)證。實(shí)驗(yàn)結(jié)果表明：

• PDS 在不修改模型訓(xùn)練框架的前提下，通過一次離線選擇，即可實(shí)現(xiàn)訓(xùn)練加速達(dá) 2 倍；
• 在多項(xiàng)下游任務(wù)中，PDS 顯著優(yōu)于現(xiàn)有數(shù)據(jù)選擇方法，且對大模型訓(xùn)練具有良好的泛化能力；
• 在數(shù)據(jù)受限條件下，PDS 可減少約 1.8 倍的訓(xùn)練數(shù)據(jù)需求，提升數(shù)據(jù)利用效率。

PDS 不僅在實(shí)際效果上具備顯著優(yōu)勢，更重要的是，它建立了一套以控制論為基礎(chǔ)的數(shù)據(jù)選擇理論框架，為理解預(yù)訓(xùn)練動態(tài)、提升模型可解釋性與可控性提供了全新視角。目前，該成果已被機(jī)器學(xué)習(xí)頂級會議 ICLR 2025 正式接收，并入選口頭報(bào)告（Oral, top 1.8%）。

• 論文標(biāo)題：Data Selection via Optimal Control for Language Models
• 論文地址：https://openreview.net/forum?id=dhAL5fy8wS
• 開源代碼：https://github.com/microsoft/LMOps/tree/main/data_selection

研究背景：訓(xùn)練大模型，不只是 “多喂數(shù)據(jù)” 這么簡單

近年來，大語言模型（LLM）不斷刷新下游任務(wù)性能的記錄。但與此同時(shí)，一個(gè)關(guān)鍵問題也日益突出：訓(xùn)練這些模型所需的數(shù)據(jù)和計(jì)算資源呈指數(shù)級增長。面對海量的互聯(lián)網(wǎng)文本，如何挑選 “更有價(jià)值” 的數(shù)據(jù)，成為提升模型效率與性能的關(guān)鍵一步。

現(xiàn)有的數(shù)據(jù)選擇方法大多依賴啟發(fā)式規(guī)則，如去重、n-gram 匹配、影響函數(shù)等，缺乏理論指導(dǎo)，效果難以穩(wěn)健推廣。而另一方面，部分方法嘗試?yán)糜?xùn)練過程中的反饋動態(tài)進(jìn)行在線數(shù)據(jù)篩選，卻需修改訓(xùn)練流程、增加訓(xùn)練時(shí)的計(jì)算開銷，實(shí)用性有限。

這項(xiàng)工作跳出常規(guī)視角，借助控制論中經(jīng)典的龐特里亞金最大值原理（Pontryagin’s Maximum Principle, PMP），首次將數(shù)據(jù)選擇建模為一個(gè)可解析的最優(yōu)控制問題，為理解和實(shí)現(xiàn)最優(yōu)數(shù)據(jù)選擇提供了系統(tǒng)的數(shù)學(xué)框架，并基于此框架設(shè)計(jì)出了一套離線數(shù)據(jù)選擇算法，在不增加訓(xùn)練開銷的情況下提升性能。

理論創(chuàng)新：數(shù)據(jù)選擇是一個(gè) “控制” 問題

作者們提出，將訓(xùn)練過程看作一個(gè)動態(tài)系統(tǒng)，數(shù)據(jù)的選擇權(quán)重作為控制變量，模型參數(shù)作為系統(tǒng)狀態(tài)，而最終下游任務(wù)的表現(xiàn)則是目標(biāo)函數(shù)。在這個(gè)框架下，預(yù)訓(xùn)練的每一步都對應(yīng)狀態(tài)的變化，而合理分配每條數(shù)據(jù)的 “重要性權(quán)重”γ，即是在有限預(yù)算下尋找最優(yōu)控制策略。基于經(jīng)典的龐特里亞金最大值原理（PMP），他們進(jìn)一步推導(dǎo)出最優(yōu)數(shù)據(jù)選擇策略所需滿足的必要條件（PMP 條件）。根據(jù)此條件來選擇數(shù)據(jù)可以很大程度上保證選擇結(jié)果的最優(yōu)性。

圖 1: PMP 條件的圖形化解釋

PMP 條件最關(guān)鍵的思想是：給出了最優(yōu)的訓(xùn)練樣本應(yīng)該具有的梯度方向（，如上左圖所示），并選擇那些梯度方向與最優(yōu)梯度高度一致的數(shù)據(jù)點(diǎn)，在數(shù)學(xué)上表現(xiàn)為梯度與  的內(nèi)積最大（如上右圖所示）。

該理論的核心價(jià)值在于：它不僅提供了選擇高質(zhì)量數(shù)據(jù)的明確準(zhǔn)則，而且揭示了目標(biāo)任務(wù)性能、模型訓(xùn)練動態(tài)與最優(yōu)數(shù)據(jù)選擇之間深層次的聯(lián)系。

算法設(shè)計(jì)：構(gòu)建高效實(shí)用的 PDS 框架

為了將理論應(yīng)用于實(shí)際的大規(guī)模語言模型訓(xùn)練，作者設(shè)計(jì)了 PMP-Based Data Selection (PDS) 算法框架，如下圖所示：

圖 2: PDS 數(shù)據(jù)選擇框架

該算法分三步進(jìn)行：

1. 在代理環(huán)境中解 PMP 方程組：在一個(gè)小規(guī)模代理模型（如 160M 參數(shù)）和代理數(shù)據(jù)集（如 0.2B tokens）上迭代求解 PMP 方程組，得到代理數(shù)據(jù)集上的最優(yōu)數(shù)據(jù)選擇策略 γ*；

2. 訓(xùn)練數(shù)據(jù)打分器（data scorer）：用一個(gè)小模型在代理數(shù)據(jù)集上擬合 γ*，根據(jù)輸入樣本輸出其質(zhì)量分?jǐn)?shù)，然后為全量數(shù)據(jù)集打分；

3. 選擇高質(zhì)量數(shù)據(jù)用于大模型訓(xùn)練：根據(jù)打分結(jié)果，對于任意的數(shù)據(jù)閾值（如 50%），選擇得分較高的樣本，用于訓(xùn)練目標(biāo)模型。

該方法完全離線進(jìn)行，僅需運(yùn)行一次，即可支持任意規(guī)模模型訓(xùn)練，且無需修改已有訓(xùn)練框架，對于高度優(yōu)化的預(yù)訓(xùn)練代碼來說，只用更換數(shù)據(jù)源，具有高度實(shí)用性與工程友好性。

實(shí)驗(yàn)效果

在實(shí)驗(yàn)中，作者基于 Redpajama CommonCrawl 中 125B token 的數(shù)據(jù)，使用 PDS 方法選出其中 50B tokens 用于訓(xùn)練 160M 至 1.7B 規(guī)模的語言模型。評估任務(wù)覆蓋 9 個(gè)主流下游以及語言建模任務(wù)。

性能提升

在不同模型規(guī)模下，PDS 訓(xùn)練出的模型在 9 個(gè)下游任務(wù)上的整體性能優(yōu)于隨機(jī)選擇（Conventional）、RHO-Loss、DSIR、影響函數(shù)（IF-Score）等方法，并且性能提升趨勢隨著模型規(guī)模的擴(kuò)大依然可以保持：

圖 3: PDS 和其他數(shù)據(jù)選擇方法的性能對比。

此外，如下左圖，PDS 方法訓(xùn)練出來的模型在高質(zhì)量語料（如 DCLM）上的語言建模性能也顯著優(yōu)于隨機(jī)選擇。如下表，使用語言模型的擴(kuò)展定律外推到 GPT-3，Llama 系列模型的訓(xùn)練規(guī)模之后，PDS 的性能優(yōu)勢依然明顯。

圖 4: PDS 方法訓(xùn)練出的模型在語言建模任務(wù)上的性能

訓(xùn)練加速

如下圖，在達(dá)到同等下游任務(wù)性能的情況下，PDS 能將 1.7B 模型的訓(xùn)練 FLOPs 減少約一半。值得注意的是，PDS 中對 PMP 條件的求解都是在預(yù)訓(xùn)練階段離線完成的，從而避免了引入訓(xùn)練時(shí)開銷。

圖 5: PDS 對于模型預(yù)訓(xùn)練的加速效果

數(shù)據(jù)利用率提升

作者們通過實(shí)驗(yàn)證明，在數(shù)據(jù)受限場景下，使用 PDS 選擇一部分高質(zhì)量數(shù)據(jù)并進(jìn)行多輪訓(xùn)練，要好于使用原始數(shù)據(jù)進(jìn)行一輪訓(xùn)練。如下圖，圖中藍(lán)線表示使用原始數(shù)據(jù)進(jìn)行 1 輪訓(xùn)練，而橙色線、綠色線、紅色線分別表示使用 PDS 選擇原始數(shù)據(jù)的 50%，25% 和 12.5%，并進(jìn)行 2 輪，4 輪和 8 輪的訓(xùn)練，從而保證總體訓(xùn)練 token 數(shù)一致?？梢钥吹?，使用 PDS 選擇質(zhì)量較高的 25% 數(shù)據(jù)表現(xiàn)最好，由此說明 PDS 提升了數(shù)據(jù)有限情況下模型的性能，即提升了數(shù)據(jù)利用率，緩解了 “數(shù)據(jù)枯竭” 問題。

圖 6: 數(shù)據(jù)有限情況下的數(shù)據(jù)選擇

未來展望：為 “數(shù)據(jù)驅(qū)動的 AI” 構(gòu)建理論框架

當(dāng)前大模型的預(yù)訓(xùn)練過程仍高度依賴經(jīng)驗(yàn)與啟發(fā)式規(guī)則，模型如何學(xué)習(xí)、數(shù)據(jù)如何影響學(xué)習(xí)的機(jī)制長期處于 “黑箱” 之中。本工作通過最優(yōu)控制理論建立起數(shù)據(jù)選擇與訓(xùn)練動態(tài)之間的數(shù)學(xué)聯(lián)系，為理解預(yù)訓(xùn)練數(shù)據(jù)的價(jià)值提供了理論抓手。

這一方向不僅有望替代傳統(tǒng)依賴直覺與試驗(yàn)的數(shù)據(jù)篩選流程，也為未來自動化、可解釋的大模型訓(xùn)練打開了新思路。作者們相信，圍繞 “如何選擇學(xué)什么” 這一核心問題建立理論體系，將成為推動 AI 從經(jīng)驗(yàn)工程走向科學(xué)建模的關(guān)鍵一步。