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

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

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

在理論基礎之上，研究團隊設計了可在大規模語料中高效運行的 PDS 算法框架，并在多個模型規模和任務設置中進行了系統驗證。實驗結果表明：

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

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

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

研究背景：訓練大模型，不只是 “多喂數據” 這么簡單

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

現有的數據選擇方法大多依賴啟發式規則，如去重、n-gram 匹配、影響函數等，缺乏理論指導，效果難以穩健推廣。而另一方面，部分方法嘗試利用訓練過程中的反饋動態進行在線數據篩選，卻需修改訓練流程、增加訓練時的計算開銷，實用性有限。

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

理論創新：數據選擇是一個 “控制” 問題

作者們提出，將訓練過程看作一個動態系統，數據的選擇權重作為控制變量，模型參數作為系統狀態，而最終下游任務的表現則是目標函數。在這個框架下，預訓練的每一步都對應狀態的變化，而合理分配每條數據的 “重要性權重”γ，即是在有限預算下尋找最優控制策略。基于經典的龐特里亞金最大值原理（PMP），他們進一步推導出最優數據選擇策略所需滿足的必要條件（PMP 條件）。根據此條件來選擇數據可以很大程度上保證選擇結果的最優性。

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

PMP 條件最關鍵的思想是：給出了最優的訓練樣本應該具有的梯度方向（，如上左圖所示），并選擇那些梯度方向與最優梯度高度一致的數據點，在數學上表現為梯度與  的內積最大（如上右圖所示）。

該理論的核心價值在于：它不僅提供了選擇高質量數據的明確準則，而且揭示了目標任務性能、模型訓練動態與最優數據選擇之間深層次的聯系。

算法設計：構建高效實用的 PDS 框架

為了將理論應用于實際的大規模語言模型訓練，作者設計了 PMP-Based Data Selection (PDS) 算法框架，如下圖所示：

圖 2: PDS 數據選擇框架

該算法分三步進行：

1. 在代理環境中解 PMP 方程組：在一個小規模代理模型（如 160M 參數）和代理數據集（如 0.2B tokens）上迭代求解 PMP 方程組，得到代理數據集上的最優數據選擇策略 γ*；

2. 訓練數據打分器（data scorer）：用一個小模型在代理數據集上擬合 γ*，根據輸入樣本輸出其質量分數，然后為全量數據集打分；

3. 選擇高質量數據用于大模型訓練：根據打分結果，對于任意的數據閾值（如 50%），選擇得分較高的樣本，用于訓練目標模型。

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

實驗效果

在實驗中，作者基于 Redpajama CommonCrawl 中 125B token 的數據，使用 PDS 方法選出其中 50B tokens 用于訓練 160M 至 1.7B 規模的語言模型。評估任務覆蓋 9 個主流下游以及語言建模任務。

性能提升

在不同模型規模下，PDS 訓練出的模型在 9 個下游任務上的整體性能優于隨機選擇（Conventional）、RHO-Loss、DSIR、影響函數（IF-Score）等方法，并且性能提升趨勢隨著模型規模的擴大依然可以保持：

圖 3: PDS 和其他數據選擇方法的性能對比。

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

圖 4: PDS 方法訓練出的模型在語言建模任務上的性能

訓練加速

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

圖 5: PDS 對于模型預訓練的加速效果

數據利用率提升

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

圖 6: 數據有限情況下的數據選擇

未來展望：為 “數據驅動的 AI” 構建理論框架

當前大模型的預訓練過程仍高度依賴經驗與啟發式規則，模型如何學習、數據如何影響學習的機制長期處于 “黑箱” 之中。本工作通過最優控制理論建立起數據選擇與訓練動態之間的數學聯系，為理解預訓練數據的價值提供了理論抓手。

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