ICLR 2025 | PKU-DAIR實驗室論文多項成果

被ICLR 2025錄用

今日，第十三屆國際學習表征會議（International Conference on Learning Representations，ICLR 2025）正式公布錄用論文名單，PKU-DAIR實驗室論文7篇論文被機器學習頂級會議ICLR 2025錄用。

ICLR 2025將于2025年4月24日至28日在新加坡博覽中心（Singapore EXPO）舉行，展示人工智能與深度學習領域的最新進展和突破性研究。

1. NetMoE: Accelerating MoE Training through Dynamic Sample Placement

作者：Xinyi Liu, Yujie Wang, Fangcheng Fu, Xupeng Miao, Shenhan Zhu, Xiaonan Nie, Bin CUI

論文鏈接：https://openreview.net/forum?id=1qP3lsatCR

Mixture of Experts (MoE) 是一種常用技術，用于在保持計算成本不變的情況下擴展模型規模以提高模型質量。MoE 模型的每一層包含多個專家，并將訓練數據僅路由到固定數量的專家，而不是全部。在分布式訓練中，由于專家分布在不同的 GPU 上，因此需要進行 All-to-All 通信，以便在每次專家路由后在 GPU 之間交換訓練數據。由于頻繁且大量的數據交換，All-to-All 通信已成為訓練效率的一個顯著挑戰。在本文中，我們從訓練樣本的角度加速 MoE 模型中的 All-to-All 通信，這是一個尚未被探索的領域。具體而言，我們發現同一訓練樣本中的數據在專家路由中具有一定程度的局部性。受此啟發，我們開發了 NetMoE，它考慮到這種局部性并動態重新排列訓練樣本的放置，以最小化 All-to-All 通信成本。我們對樣本放置進行了建模，并通過整數規劃問題在多項式時間內推導出最佳放置方案。實驗結果表明， NetMoE 相較于現有 MoE 訓練框架實現了顯著效率的提升。

2. SysBench: Can Large Language Models Follow System Messages?

作者：Yanzhao Qin, Tao Zhang, Tao Zhang， Yanjun Shen, Wenjing Luo, Haoze Sun, Yan Zhang, Yujing Qiao, Weipeng Chen, Zenan Zhou, Wentao Zhang, Bin Cui

論文鏈接：https://arxiv.org/pdf/2408.10943v1

代碼鏈接：https://github.com/PKU-BaichuanMLSystemLab/SysBench

合作單位：百川智能

大型語言模型（LLMs）在各種應用場景中發揮著重要作用，因此通過定制化模型以適應特定場景的需求變得愈加重要。System Message（系統消息）是LLMs用于引導模型的行為以實現預期目標的一個關鍵組成部分，由精心設計的指令組成。盡管系統消息實際場景中被廣泛應用，但目前仍然缺少一個全面的基準來評估不同LLMs對系統消息的實際遵循能力。為填補這一空白，我們提出了SysBench，一個分析模型對系統消息遵循能力的基準測試，重點分析了該場景中以下三個具有挑戰性的問題：約束復雜性、指令對齊性和多輪交互穩定性。為了實現有效評估，SysBench覆蓋了現實場景系統消息六種常見的約束類型，并包含各種指令對齊關系，組成多輪交互用戶對話。具體而言，我們的數據集包含來自多個領域的500條系統消息，每條系統消息與5輪用戶對話相關聯，經過人工精心編寫和審核以確保數據質量，在實驗中表現出了良好的評估一致性。我們基于SysBench對多個LLMs進行了廣泛評估，評估結果解釋了現有模型在系統消息遵循方面的優勢與不足，且為未來研究提供了關鍵見解和方向。

3. IterComp: Iterative Composition-Aware Feedback Learning from Model Gallery for Text-to-Image Generation

作者：Xinchen Zhang*, Ling Yang* (co-first), Guohao Li, Yaqi Cai, Jiake Xie, Yong Tang, Yujiu Yang, Mengdi Wang, Bin Cui

論文鏈接：https://arxiv.org/abs/2410.07171

代碼鏈接：https://github.com/YangLing0818/IterComp

合作單位：清華大學，牛津大學，普林斯頓大學

先進的擴散模型，如RPG、Stable Diffusion 3和FLUX，在組合文本到圖像生成方面取得了顯著進展。然而，這些方法通常在組合生成方面表現出不同的優勢，有些在處理屬性綁定方面表現出色，而另一些則在空間關系上更為突出。這種差異突顯了需要一種方法，能夠利用各種模型的互補優勢，以全面提升組合能力。為此，我們提出了IterComp，這是一種新穎的框架，聚合來自多個模型的組合感知模型偏好，并采用迭代反饋學習方法來增強組合生成。具體而言，我們整理了六個強大的開源擴散模型的庫，并評估它們的三個關鍵組合指標：屬性綁定、空間關系和非空間關系。基于這些指標，我們開發了一個組合感知模型偏好數據集，包含大量圖像排名對，以訓練組合感知獎勵模型。然后，我們提出了一種迭代反饋學習方法，以閉環方式增強組合性，使基礎擴散模型和獎勵模型在多個迭代中逐步自我優化。理論證明表明了其有效性，廣泛的實驗結果顯示我們在多類別物體組合和復雜語義對齊方面相比之前的最先進方法（如Omost和FLUX）具有顯著優勢。IterComp為擴散模型和組合生成的獎勵反饋學習開辟了新的研究方向。

4. SuperCorrect: Supervising and Correcting Language Models with Error-Driven Insights

作者：Ling Yang, Zhaochen Yu, Tianjun Zhang, Minkai Xu, Joseph E. Gonzalez, Bin Cui, Shuicheng Yan

論文鏈接：https://arxiv.org/abs/2410.09008

代碼鏈接：https://github.com/YangLing0818/SuperCorrect-llm

合作單位：斯坦福大學， 伯克利大學

大型語言模型（LLMs），如GPT-4、PaLM和LLaMA，在各種推理任務中表現出顯著改善。然而，較小的模型如Llama-3-8B和DeepSeekMath-Base在復雜數學推理方面仍然存在困難，因為它們無法有效識別和糾正推理錯誤。最近的基于反思的方法旨在通過實現自我反思和自我糾正來解決這些問題，但在獨立檢測推理步驟中的錯誤方面仍面臨挑戰。為克服這些局限性，我們提出了SuperCorrect，這是一種新穎的兩階段框架，利用大型教師模型來監督和糾正較小學生模型的推理和反思過程。在第一階段，我們從教師模型中提取層次化的高層和詳細思維模板，以指導學生模型引出更細致的推理思維。在第二階段，我們引入跨模型協作直接偏好優化（DPO），通過在訓練過程中遵循教師的糾正軌跡來增強學生模型的自我糾正能力。這種跨模型DPO方法教會學生模型有效定位和解決錯誤思維，并借助教師模型的改錯思路，打破其思維瓶頸，獲取應對挑戰性問題的新技能和知識。大量實驗結果一致表明我們優于之前的方法。值得注意的是，我們的SuperCorrect-7B模型在MATH/GSM8K基準測試中分別比強大的DeepSeekMath-7B超出7.8%/5.3%和比Qwen2.5-Math-7B超出15.1%/6.3%，在所有7B模型中實現了新的最先進性能。

5. SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget

作者： Zihao Wang, Bin CUI, Shaoduo Gan

論文鏈接：https://openreview.net/forum?id=9HK2rHNAhd

代碼鏈接：https://github.com/hetailang/SqueezeAttention

合作單位：明之幾何

優化LLM的KV-cache被廣泛認為是降低推理成本的關鍵。大多數現有的KV-cache壓縮算法依據句子中不同token的重要性的差異，在token sequence維度上進行稀疏化壓縮。然而，這些方法大多對所有模型層一視同仁，即為每一層分配相同的KV cache budget。而我們發現，通過識別attention layer的重要性，可以從兩個維度聯合優化KV-cache。基于我們在推理過程中對各層重要性的觀察分析，我們提出了一種名為 SQUEEZEATTENTION 的方法，能夠動態調整各層KV-cache budget的分配，而后每層再根據所分配的cache budget進行sequence維度的壓縮。通過從sequence和layer兩個維度共同優化KV-cache，SQUEEZEATTENTION 在各種LLM和基準測試中能夠節省約30%至70%的推理內存，并將推理吞吐率提高了最多2.2倍。

6.  Facilitating Multi-turn Function Calling for LLMs via Compositional Instruction Tuning

作者：Mingyang Chen, Haoze Sun, Tianpeng Li, Fan Yang, Hao Liang, Keer Lu, Bin Cui, Wentao Zhang, Zenan Zhou, Weipeng Chen

論文鏈接：https://arxiv.org/pdf/2410.12952

代碼鏈接：https://github.com/PKU-Baichuan-MLSystemLab/BUTTON

合作單位：百川智能

大型語言模型 (LLM) 在執行各種任務方面表現出巨大的潛力，包括調用函數或使用外部工具來增強其性能的能力。雖然目前對 LLM 函數調用的研究主要集中在單輪交互上，但本文討論了 LLM 參與多輪函數調用的被忽視的必要性——這對于處理需要使用函數進行規劃但不僅僅是使用函數的組合、現實世界查詢至關重要。為了促進這一點，我們引入了一種方法，BUTTON，它通過自下而上的指令構造和自上而下的軌跡生成來生成合成組合指令調整數據。在自下而上的階段，我們根據現實世界場景生成簡單的原子任務，并使用基于原子任務的啟發式策略構建組合任務。然后為這些組合任務開發相應的功能。自上而下的階段以多智能體環境為特色，其中利用模擬人、助手和工具之間的交互來收集多輪函數調用軌跡。這種方法確保了任務的組合性，并通過檢查組合任務中的原子任務來實現有效的功能和軌跡生成。我們生成了一個包含 8k 個數據點的數據集 BUTTONInstruct，并通過在各種 LLM1 中進行的大量實驗證明了其有效性。

7. GlycanML: A Multi-Task and Multi-Structure Benchmark for Glycan Machine Learning

作者：Minghao Xu, Yunteng Geng, Yihang Zhang, Ling Yang, Jian Tang, Wentao Zhang

論文鏈接：https://arxiv.org/abs/2405.16206

代碼鏈接：https://github.com/GlycanML/GlycanML

項目主頁：https://glycanml.github.io/project

合作單位：Mila 魁北克人工智能

在這個項目中，我們推出了國際上第一個針對多糖化合物的機器學習基準GlycanML。GlycanML包含了各種類型的多糖功能理解任務，包括多糖分類學預測、多糖免疫原性預測、糖基化類型預測和蛋白質-多糖相互作用預測。在GlycanML上我們評估了各種類型機器學習模型的綜合表現，包括序列編碼模型、圖神經網絡、小分子預訓練編碼器等，最終我們驗證了基于多關系建模的異構圖神經網絡在各項任務上獲得了普遍最優的性能，這為設計更強大的多糖模型提供了技術路線。此外，在8個不同尺度的多糖分類學任務上，我們建立了多任務學習基準GlycanML-MTL，旨在評估各種多任務學習算法在多糖表征學習上的有效性。通過評估代表性的多任務學習算法，我們發現基于自適應溫度放縮（Temperature Scaling）的算法能夠最有效地進行多尺度多糖分類學預測。

實驗室簡介

北京大學數據與智能實驗室（Data And Intelligence Research Lab at Peking Univeristy，PKU-DAIR實驗室）由北京大學計算機學院崔斌教授領導，長期從事數據庫系統、大數據管理與分析、人工智能等領域的前沿研究，在理論和技術創新以及系統研發上取得多項成果，已在國際頂級學術會議和期刊發表學術論文200余篇，發布多個開源項目。課題組同學曾數十次獲得包括CCF優博、ACM中國優博、北大優博、微軟學者、蘋果獎學金、谷歌獎學金等榮譽。PKU-DAIR實驗室持續與工業界展開卓有成效的合作，與騰訊、阿里巴巴、蘋果、微軟、百度、快手、中興通訊等多家知名企業開展項目合作和前沿探索，解決實際問題，進行科研成果的轉化落地。