Dr.LLM:大語言模型中的動態層路由
論文(Dr.LLM: Dynamic Layer Routing in LLMs)介紹了一種巧妙的技術,在提高準確率的同時降低大語言模型的計算量。路由器在提高準確率的同時,每次查詢可減少約3到11層的計算。論文的方法是可改裝的,為凍結的大語言模型添加逐層路由器,決定跳過、執行或重復每個模塊。路徑通過對層編輯進行簡短的蒙特卡洛樹搜索進行離線監督,然后在線執行時無需搜索。在平均節省層數的同時,提高了邏輯和數學任務的準確率。


啥是Dr.LLM: 附加在每一層的微型MLP路由器讀取窗口化的均值池化隱藏狀態,并輸出三種操作之一:跳過、執行一次或重復一次?;A權重保持凍結,并且與KV緩存保持兼容。

監督機制如何工作: 長度感知的MCTS探索在計算預算約束下跳過或重復層的編輯前向傳播,只保留那些保持或改善金標準答案獎勵的路徑。然后使用焦點損失和類別重新平衡在大約4000條發現的路徑上訓練路由器。

主要結果: 在六個主干網絡上的ARC和DART任務中,路由器在提高準確率的同時,每次查詢減少約3到11層。 例如:LLaMA-3B-Base在DART上從11.8%提升到15.8%,平均節省4.1層。經過指令調優的8B模型在DART上也有提升,同時節省11層。 域外泛化能力也很強。

這個方法為什么有效: 早期層保持穩定,許多中間層被跳過,后期層有時會被重復,尤其是在更難的數學問題上,這將深度重新分配到迭代優化更有價值的地方。

凍結解碼器LLM,為每個模塊附加一個2層瓶頸MLP路由器。對隱藏狀態使用窗口化均值池化。運行簡短的離線MCTS對層編輯進行搜索,使用長度懲罰來收集改進或保持的路徑。** 使用焦點損失訓練路由器以處理嚴重的類別不平衡。在推理時,執行貪婪路由,無需搜索??蛇x擇暴露單個標量來偏向更多跳過或更多重復,以實現更嚴格的成本控制。
本文轉載自??AI帝國??,作者:無影寺

















