Kimi K3的成功藏在4個月前楊植麟的這場演講當中

來源：AI寒武紀

kimi k3為什么可以取得突破，我們可以從4個月前（3月21）月之暗面ceo 楊植麟在英偉達GTC大會上的演講回憶一下。

演講的主題是Kimi最新模型K2.5背后的擴展方法論。整場演講幾乎全程在講技術細節,涉及優化器、長上下文架構、智能體訓練范式,以及一個剛剛公布的新架構預告。

但從這次演講中可以窺見楊植麟的LLM擴展的三大核心技術野心：1.主要Adam優化器誕生于2014年,團隊現在做出了一個可以直接替換Adam的開源版本Muon Clip,用它訓練Transformer大語言模型,效果會明顯更好;2.注意力機制誕生于8年多之前,團隊現在推出了它的線性版本Kimi Linear,不必在每一層都用全注意力,線性注意力在長短上下文任務上都能表現得更好;3.殘差連接,團隊推出了開源版本的注意力殘差架構。

現在大家看到了，僅僅過去了4個月，K3就是最好的答案

擴展的三個維度

整場演講的主線,是團隊在三個維度上做擴展。

第一個維度是最經典的scaling law:橫軸是訓練token數量,縱軸是loss,token越多,loss越低。團隊關注的重點不是單純堆token數量,而是提升token的利用效率,讓同樣數量的token能換來更低的loss,這靠的是更好的架構和優化器。

第二個維度是擴展上下文長度。上下文越長,模型在給定位置預測下一個token的準確率就越高,這意味著模型能完成更復雜的任務。

第三個維度是智能體數量。團隊提出了一種新的訓練范式,叫智能體蜂群,讓模型不再單純依賴單個智能體處理任務,轉而編排一群智能體并行完成子任務,從而提升處理復雜任務的能力。

這三個維度可以直接映射到智能體的語境里:token效率對應更強的先驗知識,讓智能體在做強化學習搜索答案時更高效;長上下文對應能持續運行更久的智能體,可能連續跑上幾天、幾周甚至幾個月去完成復雜任務;智能體蜂群則是疊加在前兩者之上的又一個維度。最終團隊想要的,是一群擁有超長上下文、又具備強先驗的智能體,在整個強化學習系統里協同搜索答案。

維度一:怎么讓每個token更值錢

楊植麟先展示了一張機器學習史上很經典的圖,來自Kaplan等人的scaling law論文:只要按比例擴大訓練token數量、模型參數量和算力,loss就會持續下降,這是過去幾年整個行業的重要基礎。

但團隊真正在意的是提升token效率。楊植麟特別強調,token效率不只是效率問題,它直接關系到智能的上限。他舉了個例子:假設手上有50萬億條高質量token,應用一個新優化器之后效率提升兩倍,效果上就相當于憑空多出了100萬億條token。眼下整個行業正在逼近所謂的數據墻,高質量數據總量有限,如果把數據量看作一個常數,提升token效率就意味著能榨取出更強的智能,而不只是省了算力這么簡單。

團隊為此重點投入的是Muon優化器,一種二階優化器。它的核心做法是,每一次梯度更新都會被轉換,使更新的每一項彼此正交,這和傳統的Adam優化器差異很大。如果實現得當,可以帶來兩倍的token效率提升。團隊是首個證明Muon優化器可以擴展用于大語言模型訓練的團隊,論文里提出了兩項關鍵技術:一是權重衰減,這對擴展到更大模型至關重要;二是保證更新幅度和Adam保持一致的水平,團隊引入了一個可調系數應用在每次更新上,讓最終的更新幅度和Adam保持可比。為了在英偉達GPU集群上做到內存高效,團隊還專門開發了一套分布式Muon優化器實現,把優化器狀態切分到不同的數據并行組里。在相同參數量和相同訓練token數的對照實驗里,把AdamW換成Muon優化器之后,各項指標全面明顯提升。

不過把Muon擴展到一萬億參數規模時,團隊遇到了新的麻煩:訓練不穩定。具體表現是,注意力機制里的最大logit值會迅速爆炸,超過1000,而正常情況下這個值一般在50甚至不到100。與此同時訓練會發散,loss先往下走一段,隨后突然爆炸,沒法按預期收斂。

團隊的解法叫QK clip。做法是對網絡里的每一個注意力頭,在前向傳播過程中計算出當前的最大logit值,再據此算出一個縮放系數,分別應用到key和query的投影上,把query和key的數值限制在一個給定范圍內,從而避免爆炸。效果很直接:應用clip前后兩條訓練loss曲線幾乎完全重合,說明這項技術不會拖累loss下降的速度;而如果單獨看最大logit這個中間指標,可以看到它一開始照常上升,在數值100左右被卡住維持一段時間,之后隨著訓練步數增加自然回落。網絡自己找到了辦法把最大logit值控制住,同時完全不影響收斂。團隊把這項技術用到了K2模型的訓練上,成功把參數量擴展到一萬億,這也是機器學習歷史上第一次把Muon優化器用在這種規模的訓練上。

維度二:把上下文拉得更長

第二個維度的引子,是一張相對少被提及但很關鍵的圖,對比的是Transformer和LSTM。左邊的圖顯示,在相同參數量和相同訓練token數量下,Transformer的訓練loss更低,這也是Transformer會成為如今主流架構的原因。右邊的圖更有意思:橫軸是上下文里的token位置,可以看到Transformer的訓練loss會隨著上下文變長持續下降,而LSTM的loss在某個token數量之后就不再下降,進入飽和狀態。這說明Transformer在捕捉長上下文信息上明顯更強。

這項能力在智能體時代格外重要,因為任務正在變得越來越復雜,需要的上下文也越來越長,比如理解一整個代碼倉庫,或者跑很長的智能體軌跡去完成任務,像是從零開始寫一個操作系統內核,這些都是LSTM做不到的。團隊的研究目標,是做出一種架構,既能高效擴展到更長的上下文,又能在更靠后的token位置上取得更低的單token loss。這正是團隊推出Kimi Linear架構的出發點。

Kimi Linear里包含一種新的線性注意力變體,叫Kimi Delta Attention,是在原始delta rule的基礎上,通過改進循環記憶機制得到的。同時團隊把線性注意力層和全注意力層按1比3的比例混合使用,兼顧長上下文能力和實現效率。

原始線性注意力的問題在于,它的記憶是全局性的,只有一個統一的衰減因子在起作用,導致模型要么幾乎遺忘掉所有歷史信息,要么幾乎把所有信息都保留下來,但沒辦法在長上下文里有選擇地丟掉不重要的部分。Kimi Delta Attention引入了細粒度的衰減因子,把原本的標量衰減系數換成了一個對角矩陣,分別控制每個通道各自的衰減速度。一部分通道可以衰減得很慢,從而在很長的范圍內保留信息,另一部分通道則可以很快遺忘掉舊信息,及時接收新信息,整體上提升了模型的表達能力。

為了適配現代GPU,團隊采用了分塊的方式來做并行計算。由于衰減項從標量變成了矩陣,不能像以前那樣簡單地提出來單獨處理,團隊因此重寫了整套計算公式,引入了矩陣求逆運算以及累積衰減因子,在不犧牲效率的前提下實現并行計算。更重要的是,這次重寫在數學上和原始公式完全等價,屬于精確改寫,效率提升的同時也不會帶來任何性能損失。

從結果看,在偏短上下文的任務比如MMLU上,Kimi Linear的表現優于MLA和GDN;在偏長上下文的任務比如ruler上,Kimi Linear同樣更好,而且比MLA更省算力。當上下文進一步擴展到一百萬token甚至更長時,Kimi Linear相比基線效率提升更加明顯,這也是第一個能在短上下文任務、長輸入任務和長輸出任務上全面超過全注意力架構的方案。

維度三:一群智能體一起干活

團隊為智能體蜂群設計了一個編排者,也可以理解為主智能體,負責協調整體任務。它可以生成一批子智能體并給它們分配任務,也可以收集子智能體返回的結果,通過反復迭代這個過程,最終完成比單個智能體能力范圍更大的復雜任務。

楊植麟拿人類社會做了類比:辦一家公司需要不同的角色分工,需要一個類似CEO的角色把任務拆解、分配給不同崗位,最終整個組織朝著同一個目標推進。放到智能體蜂群里,可能同時存在負責AI研究的智能體、負責網頁開發的智能體、負責物理研究的智能體,各自研究不同的課題,最后再生成一批事實核查員、網頁開發者和文件下載員,把所有結果匯總成一份完整報告。

從另一個角度看這個范式:橫軸是任務復雜度,用一組模型在該任務上的準確率來衡量,縱軸是執行時間。用上智能體蜂群之后,執行時間比單智能體大幅縮短,效率明顯更高。這意味著可以把蜂群規模擴展到比如100個甚至1000個子智能體,在可接受的時間范圍內完成復雜任務,從而真正產生經濟價值。團隊認為這種擴展可以體現在多個方向:擴展輸入,比如并行下載、閱讀成百上千份資料;擴展輸出,比如并行寫出一百頁的文獻綜述;擴展執行的動作,比如同時對十個不同任務做數據分析;也可以擴展編排本身,也就是學會怎么拆分子任務、怎么把結果匯總起來。

技術層面,團隊為智能體蜂群系統專門設計了新的獎勵目標,相比傳統的單智能體強化學習,這里一共考慮三項獎勵。

第一項叫實例化獎勵,用來鼓勵模型生成子智能體,防止出現退化成單智能體串行執行的現象,尤其是在訓練早期要鼓勵并行執行,這項獎勵的權重會隨訓練推進逐漸衰減。

第二項叫完成獎勵。團隊觀察到,有些子任務被創建出來之后從來沒有被完成過,模型幾乎是靠不斷生成一堆子智能體去刷第一項獎勵,而這些任務本身可能太復雜,或者壓根沒有意義。因此加入完成獎勵,鼓勵每個子任務都保持相對較高的完成比例,確保生成出來的子任務是有意義的,同樣采用訓練前期權重高、后期權重低的衰減策略。

第三項是標準的結果獎勵,衡量整個任務最終是否完成。

三項獎勵加總起來,構成了整個強化學習系統的優化目標。與此同時,團隊也搭建了配套的基礎設施,用來支持并行執行、支持不同的獎勵函數,并盡可能提升整個智能體蜂群系統的輸入輸出效率。

三條線合到一起,就有了K2.5

把token效率、長上下文、智能體蜂群這三條線放在一起,就誕生了Kimi K2.5

楊植麟提到,團隊從K2.5身上觀察到不少有意思的能力,比如視覺能力和編程能力融合之后,會涌現出很多新玩法,例如看一段視頻,就能生成一個復現或者風格遷移自這段視頻的網頁,這些能力都源于預訓練階段成功且穩定的訓練過程。

他還展示了K2.5基礎模型的訓練曲線,稱這是自己見過最漂亮的曲線之一。模型經歷了超過15萬億token的訓練,在K2.5階段又額外訓練了15萬億token,整個訓練過程非常穩定,尤其是引入新的Muon優化器之后,沒有出現任何loss尖峰。這種平滑穩定的訓練過程,產出了一個很強的基座模型,團隊可以在此基礎上繼續微調,獲得視頻里展示的那些新能力。整個訓練在英偉達H800 GPU集群上完成,集群里每個節點配備2TB內存,GPU之間通過NVLink連接。

K2.5的另一項關鍵創新,是它成為首個原生實現視覺和文本聯合訓練的開源模型。以往開源模型通常的做法是先在純文本上訓練,比如先用20萬億token訓練出一個文本模型,再額外用2萬億token左右的后訓練,把視覺能力疊加上去。K2.5的做法不一樣,團隊從訓練的第一天就把視覺和文本訓練過程融合在一起,這被稱為早期融合,從進度0%開始就把視覺token和文本token合并訓練。初步實驗顯示,這種方式效果優于先文本后視覺的后期融合做法,團隊觀察到的一些新能力也正是來自這種訓練方式。舉個例子,如果想讓模型看圖生成代碼,必須把視覺和文本放進同一個分支里訓練才能實現,如果兩個分支各自獨立訓練是做不到的,必須把兩種模態對齊進同一個共享表示空間里才行。

另一個有意思的發現是,這兩種模態其實可以互相增強。行業里此前普遍的認知是,給文本模型加上視覺能力往往會拖累文本本身的表現,但團隊發現,只要訓練方式得當,兩種模態反而能夠相互促進,這是這次訓練過程里的一個重要發現。

具體來看,視覺訓練能提升文本表現。團隊做了一組對比實驗,只用純視覺任務做強化學習,不涉及任何文本任務,比如教模型數數、回答視覺問答,不包含任何數學或編程題目,結果發現即便是偏重推理的文本任務,表現也跟著提升了。反過來,文本訓練也能提升視覺表現。如果文本基礎足夠扎實,訓練過程中其實不需要額外準備視覺方向的SFT數據。團隊采用的方法叫zero vision SFT,也就是視覺方向幾乎不用SFT數據,唯一使用的SFT數據是文本SFT數據,再對文本和視覺做聯合強化學習。結果顯示,在完全沒有專門視覺數據的情況下,視覺類任務的表現依然能做到接近業內最好水平。這說明只要在預訓練階段把兩種模態對齊進同一個共享空間,足夠強的文本基座本身就能帶動視覺能力一起提升。視頻里展示的強視覺設計和前端編程能力,同樣來自這套視覺文本聯合訓練方式。

下一代架構可能長什么樣

演講快結束時,楊植麟順帶介紹了團隊前一天剛剛放出技術報告的新架構,叫注意力殘差,算是對下一代架構方向的一次預告。

出發點很簡單:團隊此前在時間維度上用到的一些技術思路,能不能搬到深度這個維度上來用一遍。

這要從殘差連接說起。楊植麟回憶起十年前在ICML 2016的一場分享上第一次聽到這個想法,稱它是個很精彩的點子。在ResNet出現之前,沒人能訓練很深的網絡,一旦層數增加就會遇到梯度爆炸、梯度消失之類的訓練穩定性問題;而ResNet提出之后,理論上可以堆疊任意多的層數,不用再擔心訓練穩定性。

他引用了Ilya兩年前一次演講里的說法:殘差連接本質上是LSTM的一個變體,只是把方向旋轉了90度。這句話可以這樣理解:LSTM是循環網絡的一種,處理過程是循環的,每一步都會把上一步的隱藏狀態拿過來,通過某種門控機制生成當前狀態;而殘差連接在深度這個維度上做的其實是同一件事,把上一層的輸出拿過來,通過某個函數處理后生成當前層的輸出,只是具體的計算形式不一樣。比如殘差連接用的是固定的加法,把上一層的隱藏狀態和當前層的輸出直接相加,本質上還是把循環網絡那套邏輯用在了深度維度上。

由此團隊想到一個新方向:既然深度維度上可以套用LSTM的邏輯,那能不能換成用注意力機制來做同一件事,替代原本類似LSTM的門控函數,畢竟注意力機制已經在Transformer時代證明了自己足夠強大。具體做法是,每一層的輸出不再只依賴上一層的隱藏狀態,而是把之前所有層的隱藏狀態都納入考慮,通過注意力操作對這些歷史隱藏狀態做聚合,計算出當前層的狀態。團隊把這個思路稱為把注意力旋轉90度,可以看作是殘差連接在這套類比之下的一種自然延伸。

為了控制通信和內存開銷,團隊還設計了一個效率更高的變體,叫分塊注意力殘差。做法是把神經網絡的所有層劃分成若干個塊,比如每塊16層,或者每塊4層,只在每個塊的輸出上應用注意力殘差,塊內部依然沿用標準的殘差連接。這樣能大幅降低開銷,同時訓練精度幾乎不受影響。

效果上,這套新架構在scaling law層面把token效率提升了24%,原本50萬億條高質量token能達到的效果,現在相當于憑空多出了60多萬億條token的效果。驗證集loss也持續低于原來的曲線,說明這套優化過程更加穩定。在編程、數學和推理密集型任務上提升最明顯,團隊在GPQA、數學和HumanEval等基準測試上都觀察到了相應的提升。

楊植麟認為,這個時代做研究的思路和過去不太一樣。十年前的研究更偏重發表新想法,但很難做到足夠嚴謹的實驗驗證,難以得出扎實的結論。而現在有了所謂的擴展階梯,團隊有足夠的資源在不同規模上訓練和驗證模型,配合一整套基準測試來衡量進展,因此更容易得出有把握的結論。這也是為什么很多看起來已經很成熟的老技術,最近又重新出現新進展的原因之一。

最后他總結,團隊會繼續在token效率、長上下文、智能體數量這三個維度上擴展模型,未來還會看到不同的架構和優化器同時在這三個維度上做優化,也會不斷發現新的擴展維度,智能體蜂群不會是終點。

k3將于7月27之前開源，到時候各個推理商提供k3后你再看A廠的敘事還能不能撐得住。

閱讀最新前沿科技趨勢報告，請訪問21世紀關鍵技術研究院的“未來知識庫”

未來知識庫是 “21世紀關鍵技術研究院”建 立的在線知識庫平臺，收藏的資料范圍包括人工智能、腦科學、互聯網、超級智能，數智大腦、能源、軍事、經濟、人類風險等等領域的前沿進展與未來趨勢。目前擁有超過8000篇重要資料。每周更新不少于100篇世界范圍最新研究資料。 歡迎掃描二維碼或訪問https://wx.zsxq.com/group/454854145828進入。

截止到2月28日 ”未來知識庫”精選的百部前沿科技趨勢報告

（加入未來知識庫，全部資料免費閱讀和下載）