中科院、百度提出新架構：突破參數限制，實現高效推理
有不少研究證明，大模型在不斷增加參數的情況下，其性能提升卻逐漸趨于飽和，特別是在受限的數據可用性和計算資源預算下，這種情況非常明顯。
為了解決這一難題，中國科學院信息工程研究所、百度、北京師范大學AI學院的研究人員聯合提出了Inner Thinking Transformer架構（簡稱ITT），通過動態分配計算資源給單個標記，增強了測試性能而無需增加參數。
ITT 的設計基于一個關鍵假設：Transformer 的每一層可以被視為對輸入信息的一次隱式推理。研究人員提出將單個token的生成過程分解為多個內部思考步驟，模型從初始狀態開始，通過一系列的思考步驟逐步更新隱藏狀態，最終生成輸出token。
這種設計允許模型在遇到復雜推理任務時，通過增加思考步驟來提升性能，而無需擴展模型的參數規模。
為了實現這一目標，ITT 引入了 Residual Thinking Connection (RTC) 機制。RTC 的作用是通過迭代累積每個思考步驟的輸出，逐步細化token的表示。與傳統的循環神經網絡相比，RTC 不僅能夠實現深度思考，還能通過學習每個步驟的重要性編碼來衡量和整合不同思考步驟的結果。這種機制使得模型能夠在訓練過程中逐步優化思考過程，并在推理時根據需要動態調整思考步驟的數量。
為了進一步提升計算效率，ITT 引入了 Adaptive Token Routing (ATR) 機制。ATR 的核心思想是動態選擇需要進一步思考的關鍵token，避免對所有token進行不必要的深度處理。模型通過一個路由網絡為每個token分配一個權重，權重較高的token會被選中進行額外的思考步驟。這種選擇性處理不僅減少了計算開銷，還能使模型更加專注于對任務至關重要的信息。
在訓練過程中，路由網絡會根據當前token的表示和思考步驟的編碼來決定是否對其進行進一步處理，使得模型能夠在訓練時學習到最優的思考模式，并在推理時根據需要靈活調整計算資源的分配。通過這種方式，ITT 實現了在不增加模型參數的情況下，對關鍵token進行深度處理的目標。
除了動態思考步驟和計算資源分配，ITT 還引入了 Thinking Step Encoding 機制。該機制通過為每個思考步驟分配一個獨特的編碼，幫助模型區分不同的思考階段。這種編碼不僅能夠衡量每個步驟的重要性，還能指導模型如何將不同步驟的結果進行整合。
在優化過程中，ITT 的設計使得模型能夠通過多步優化來替代傳統的單步優化。每個思考步驟的參數更新是基于全局損失的梯度乘以局部映射的導數。這種設計使得模型在反向傳播時能夠更穩定地收斂，并避免了梯度消失或爆炸的問題。通過這種方式，ITT 實現了在不增加模型參數的情況下，通過動態擴展思考步驟來提升模型的推理能力和性能。
為了測試ITT性能，研究人員使用了 RedPajama 數據集，這是一個包含 500 億訓練標記和 200 萬驗證標記的多領域數據集。模型訓練基于 Sheared-LLaMA 代碼框架，在 8 個 NVIDIA A100 GPU 上進行，序列長度為 4096，全局批量大小為 256。
實驗結果顯示，ITT 在不同參數規模（162M、230M和 466M 參數）的模型上均顯著優于基線Transformer 和 Loop 模型。
例如，在 162M 參數規模下，ITT ×4 模型通過在 50% 的層中進行 4 步思考，相比基線模型提升了 1.7% 的性能，而 Loop 模型僅提升了 0.3%。這種性能提升在更大規模的模型中也得到了體現，表明 ITT 的設計能夠有效擴展到不同規模的模型。
ITT 的另一個顯著優勢是其數據效率。在訓練過程中，ITT只需使用基線模型 56.8% 的數據量，即可達到與基線模型相同的性能水平。這一結果表明，ITT 通過動態思考步驟和計算資源分配，能夠更高效地利用訓練數據，減少對大規模數據集的依賴
論文地址：https://arxiv.org/abs/2502.11089
本文轉自 AIGC開放社區 ，作者：AIGC開放社區
