玩AI的朋友最近一年估計(jì)都被推理成本搞瘋了。要解決數(shù)學(xué)、代碼這種燒腦題，就得用帶慢思考的大模型，結(jié)果模型思考一步就要生成一堆token，這些token都是算出來的草稿，看得見，但是貴到肉疼。一道復(fù)雜數(shù)學(xué)題，光思考過程耗的計(jì)算資源，都是普通對話的十倍往上，誰頂?shù)米“　?/p>

現(xiàn)在業(yè)內(nèi)所有人都在想辦法壓低推理成本，可只要思維鏈那一堆中間步驟還得一個(gè)個(gè)生成token，那延遲就降不下根子來。畢竟前一步?jīng)]出結(jié)果，后一步根本開始不了，推理鏈多長，你就得等多久，這是架構(gòu)層面的問題，不是調(diào)調(diào)參數(shù)就能糊弄過去的。

那能不能讓模型把草稿藏在自己腦子里，不輸出中間步驟，還能保留原來的推理能力呢？這不，來自加州伯克利和普林斯頓大學(xué)的研究團(tuán)隊(duì)，真就朝著這個(gè)方向搞出了關(guān)鍵突破，還從數(shù)學(xué)上嚴(yán)格證明了這個(gè)方法真的管用。

要明白這個(gè)新方法牛在哪，得先搞懂原來的顯式思維鏈為啥這么費(fèi)錢。舉個(gè)很簡單的例子，你教學(xué)生做多位數(shù)乘法，一種方法是讓孩子把每一步運(yùn)算都寫在紙上，算完個(gè)位算十位，最后再加起來，每一步都清清楚楚方便檢查。另一種方法就是讓孩子在心里面算完，直接給你最終答案，所有中間過程都不用寫出來省時(shí)間。

放到大模型身上，這個(gè)差別就直接體現(xiàn)在花錢多少和等待時(shí)長上。原來的顯式思維鏈，推理有多少步就要輸出多少個(gè)額外token，這些token還得嚴(yán)格按照順序一個(gè)個(gè)生成，現(xiàn)在好用點(diǎn)的推理模型，動不動就是幾百上千個(gè)中間token，算下來能不貴嗎。隱式思維鏈的思路，就是讓模型把中間步驟全內(nèi)化到自己的隱藏狀態(tài)里，推理的時(shí)候只輸出最終答案，不用多花一毛錢在中間token上。

其實(shí)隱式思維鏈這個(gè)想法好幾年前就有人提了，之前的老方法得一步步把中間token藏起來，15步推理就得搞14個(gè)訓(xùn)練階段，訓(xùn)練開銷直接跟著推理鏈長度成正比漲，太不劃算。更尷尬的是，沒人能說清這個(gè)方法為啥會有效，也沒法保證它做出來的效果和原來的顯式思維鏈一模一樣，一直就是個(gè)碰運(yùn)氣的玄學(xué)技巧。

這次伯克利團(tuán)隊(duì)搞出來的Log-ICoT，最核心的突破就是摸到了思維鏈的本質(zhì)結(jié)構(gòu)。說白了，任何多步推理的思維鏈，拆開來都是一棵樹，就拿16步推理來說，拆完就是一棵深度為4層的二叉樹，從最底層的輸入節(jié)點(diǎn)一層層往上算，到根節(jié)點(diǎn)就是最終答案。老方法一次只藏一個(gè)步驟，完全浪費(fèi)了這個(gè)天然結(jié)構(gòu)，新方法直接一次藏掉整層，原來15個(gè)訓(xùn)練階段直接壓縮成4個(gè)，剛好對上Transformer本身的分層結(jié)構(gòu)，每一層模型剛好負(fù)責(zé)處理樹的一層，完美適配。

不光訓(xùn)練效率提了飛起，這還是學(xué)術(shù)界第一個(gè)給隱式思維鏈做了嚴(yán)格數(shù)學(xué)證明的工作。團(tuán)隊(duì)用理論計(jì)算機(jī)科學(xué)經(jīng)典的k-奇偶校驗(yàn)問題做測試，最終證明了L層的Transformer用Log-ICoT訓(xùn)練，只需要log?k個(gè)訓(xùn)練階段，推理的時(shí)候不用輸出任何中間token，就能做到和顯式思維鏈一樣的準(zhǔn)確率，樣本復(fù)雜度完全沒有上漲。

研究過程里還搞定了兩個(gè)繞不開的技術(shù)難題，第一個(gè)是表示坍縮，就是Transformer堆多了層之后，不同位置的向量表示會慢慢趨同，梯度直接沒了信號。團(tuán)隊(duì)設(shè)計(jì)了門控連接，每層只激活對應(yīng)樹層級的位置，剩下的全部關(guān)閉，梯度剛好精準(zhǔn)用在該處理的任務(wù)上，不會平白浪費(fèi)消失。第二個(gè)難題是誤差傳播，早期訓(xùn)練的一點(diǎn)點(diǎn)小誤差，經(jīng)過多階段放大之后會直接毀掉整個(gè)結(jié)果，團(tuán)隊(duì)的解決辦法也很巧妙，每次梯度更新完把注意力權(quán)重量化成最近的整數(shù)，直接把訓(xùn)練好的部分鎖死，誤差根本傳不下去。

后續(xù)做的真實(shí)實(shí)驗(yàn)結(jié)果，也完全貼合之前的理論預(yù)測。團(tuán)隊(duì)用k=16也就是4個(gè)訓(xùn)練階段做測試，第一階段開放完整思維鏈，損失很快就掉到接近零。之后每一個(gè)階段都藏掉一半的思維鏈位置，損失會短暫出現(xiàn)一個(gè)小尖峰，剛好對應(yīng)模型在消化新藏起來的步驟，很快就又回落下去。等到四個(gè)階段全部結(jié)束，所有思維鏈位置都被藏起來，模型只拿到原始輸入，驗(yàn)證集準(zhǔn)確率直接沖到了100%。

就連注意力權(quán)重的分布，都剛好和理論分析的一模一樣。第一層注意力就盯著樹第一層的輸入節(jié)點(diǎn)對，第二層就盯著第二層的計(jì)算節(jié)點(diǎn)，模型真的把每一層思維鏈都妥妥放進(jìn)了對應(yīng)Transformer層里，沒有亂成一鍋粥。

之前隱式思維鏈只有實(shí)踐層面的有效結(jié)果，沒人能說清它底層為啥能跑通，這次相當(dāng)于給這個(gè)方向正式正名了。原來它不是碰巧好用的野路子，是在明確條件下就能保證有效的訓(xùn)練方法，模型的沉默思考第一次有了數(shù)學(xué)層面的合法性。

這個(gè)方向走通了，未來我們就能把大模型的長推理鏈，通過結(jié)構(gòu)化訓(xùn)練一步步全壓縮進(jìn)模型的隱藏層。到那時(shí)候，模型照樣有頂尖的推理能力，用戶用的時(shí)候直接拿到答案，不用熬半天等輸出，也不用收到動輒幾十塊的token賬單，想想都香。

當(dāng)然現(xiàn)在這個(gè)成果離落地到真實(shí)大模型還有一段距離，目前的證明還是基于不少簡化假設(shè)，用的也是理論測試用的合成任務(wù)。接下來要解決的問題，就是怎么給沒有明確層級結(jié)構(gòu)的真實(shí)任務(wù)做階段劃分，不過能邁出這最關(guān)鍵的一步已經(jīng)足夠有意義，方向?qū)α耍Ｏ碌墓こ虇栴}只是時(shí)間問題。

參考資料：新華社 人工智能推理成本優(yōu)化研究取得國際新突破