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