明星創(chuàng)業(yè)公司Thinking Machines，第二篇研究論文熱乎出爐！

公司創(chuàng)始人、OpenAI前CTO Mira Murati依舊親自站臺，翁荔等一眾大佬也紛紛轉(zhuǎn)發(fā)支持：

論文主題為“Modular Manifolds”，通過讓整個網(wǎng)絡(luò)的不同層/模塊在統(tǒng)一框架下進(jìn)行約束和優(yōu)化，來提升訓(xùn)練的穩(wěn)定性和效率。

之所以進(jìn)行這項研究，主要是為了解決神經(jīng)網(wǎng)絡(luò)訓(xùn)練中的一個基本挑戰(zhàn)：

網(wǎng)絡(luò)內(nèi)部的張量（如權(quán)重、激活、梯度）若數(shù)值過大或過小，會引發(fā)不穩(wěn)定、梯度爆炸/消失、訓(xùn)練效率低下等問題。

因此，論文唯一作者Jeremy Bernstein提出了一種新的優(yōu)化思路——模塊化流形（Modular Manifolds），不僅對單個權(quán)重張量施加約束，還能把整個網(wǎng)絡(luò)視為組合的流形結(jié)構(gòu)，從而統(tǒng)一設(shè)計學(xué)習(xí)率、約束方式與優(yōu)化邏輯。

網(wǎng)友們的反應(yīng)be like：

• 對神經(jīng)網(wǎng)絡(luò)訓(xùn)練有深刻的見解。

• 將權(quán)重更新限制在流形上可以帶來穩(wěn)定性和可解釋性，好奇能否擴(kuò)展到非常大的模型？

帶著這些肯定和疑惑，咱們接著看這到底是一項怎樣的研究——

從向量在球面上優(yōu)化→提出模塊化流形

為什么需要流形約束？

原因正如開頭所提到的，在訓(xùn)練大模型時，如果權(quán)重、激活值、梯度太大或太小，就會出現(xiàn)訓(xùn)練不穩(wěn)定（溢出、消失、收斂速度慢等問題）。

之前為了解決這些問題，歸一化方法逐漸成為“黃金標(biāo)準(zhǔn)”，包括激活歸一化（如層歸一化LayerNorm）、梯度歸一化（優(yōu)化器內(nèi)部的規(guī)范化）等，但很少直接對權(quán)重矩陣本身進(jìn)行歸一化處理。

而作者認(rèn)為，對權(quán)重矩陣進(jìn)行歸一化處理具有相當(dāng)價值，其好處可能包括：

• 更容易把握優(yōu)化過程中更新量的大??；
• 能夠避免權(quán)重的范數(shù)變得過大，也就是防止權(quán)重“爆炸”的問題；
• 讓研究員可以把調(diào)整超參數(shù)的精力，更多地放在那些對模型影響最大的張量上；
• 讓矩陣的條件數(shù)變小，這樣矩陣的行為就會更加穩(wěn)定和可預(yù)測；

一句話，給權(quán)重矩陣歸一化，可以讓模型訓(xùn)練更穩(wěn)定、更容易調(diào)整、行為更可預(yù)測，并且對外界干擾更有抵抗力。

基于此，作者希望設(shè)計一個幾何化框架，把神經(jīng)網(wǎng)絡(luò)的權(quán)重參數(shù)約束在特定的Stiefel流形上，從而在訓(xùn)練時能夠聯(lián)合優(yōu)化器與這些約束。

其核心研究過程大致有以下幾個步驟：

Step 1：提供一個基礎(chǔ)示例

假設(shè)要訓(xùn)練的參數(shù)是一個向量W，并且強(qiáng)制它始終在單位球面上：||W||=1。

如果用普通的Adam（自適應(yīng)矩估計）/SGD（隨機(jī)梯度下降）更新，更新后的向量可能就跑出球面。

為了解決這一問題，流形優(yōu)化自有其“套路”：

• 先把梯度投影到切空間（球面上某點(diǎn)的切平面），保證更新方向合理；
• 然后更新參數(shù)；
• 最后用Retraction投影，把更新后的向量“拉回”到球面上。

而在這個過程中，我們核心需要思考兩個問題：一是應(yīng)該采用哪種流形約束，二是應(yīng)該如何測量長度。通過做出不同的選擇，最后可以創(chuàng)建出不同的優(yōu)化算法，如下表所示：

Step 2：推廣到矩陣參數(shù)

緊接著，作者將上述思路從向量（球面）推廣到矩陣。

由于Transformer的權(quán)重矩陣（把輸入的向量轉(zhuǎn)換成輸出的向量）動輒成千上萬維，直接優(yōu)化容易失控，因此作者提出把矩陣參數(shù)放在Stiefel流形上。

在Stiefel流形上的矩陣，它們的列向量都是正交的，而且條件數(shù)（一個衡量矩陣穩(wěn)定性的指標(biāo)）為1。

這樣做的好處在于：

其一，向量正交可以確保不同方向上的信號不會相互干擾，這有助于模型學(xué)習(xí)到更加獨(dú)立的特征；

其二，條件數(shù)為1意味著矩陣不會放大或縮小向量太多，這有助于保持?jǐn)?shù)值穩(wěn)定性，避免在計算過程中出現(xiàn)大的數(shù)值誤差。

選定流形后，作者通過一系列步驟最終設(shè)計了一個適用于Stiefel流形的優(yōu)化算法——流形Muon算法。

其流程大致包括：

• 在Stiefel流形的切空間上算梯度更新；
• 用一些矩陣運(yùn)算（奇異值分解、sign函數(shù)等）確保更新合法；
• 再把結(jié)果投影回Stiefel流形。

Step 3：小規(guī)模實(shí)驗(yàn)驗(yàn)證

為了驗(yàn)證算法的合理性，作者隨即進(jìn)行了一個小規(guī)模實(shí)驗(yàn)（每次訓(xùn)練運(yùn)行不到一分鐘即可完成）。

他在CIFAR-10數(shù)據(jù)集上（一個包含10個類別的彩色圖像數(shù)據(jù)集）訓(xùn)練了一個小MLP，對比了manifold Muon和AdamW這兩種算法。

結(jié)果發(fā)現(xiàn)，前者在訓(xùn)練/測試準(zhǔn)確率上略優(yōu)于AdamW，且權(quán)重矩陣的奇異值都更集中、幅度穩(wěn)定。

不過，雖然效果有提升，但每一步的運(yùn)行時間倒是比AdamW稍慢一點(diǎn)（額外開銷在dual ascent計算、投影等環(huán)節(jié)）。

作者表示，后續(xù)通過優(yōu)化dual ascent步數(shù)、引入動量等方法，可以進(jìn)一步減少額外開銷。（意思是問題不大）

Step 4：最終推出“模塊化流形”概念

以上驗(yàn)證了單個矩陣放到Stiefel流形上的有效性，那么問題來了——

多個層、多個模塊怎么辦？

對此，作者最終提出了模塊化流形（Modular Manifolds）這一概念。

具體而言，作者把神經(jīng)網(wǎng)絡(luò)里的每一層或每一個模塊都看成是一個單獨(dú)的“流形”，它有自己定義的范數(shù)和專屬的優(yōu)化方法。

當(dāng)這些模塊組合成一個完整的網(wǎng)絡(luò)時，就把它們的流形通過笛卡爾積拼接在一起，形成一個大的流形空間。

在這個大空間里，作者采用最大范數(shù)（max norm）的方式來統(tǒng)一分配學(xué)習(xí)率，相當(dāng)于設(shè)置一個全局的更新上限，避免不同模塊的步伐不一致。

這樣一來，每一層仍然可以按照自己的規(guī)則在小空間里更新參數(shù)，但整個網(wǎng)絡(luò)的更新過程又被全局機(jī)制約束住，從而保持協(xié)調(diào)和穩(wěn)定。

總之，這條研究路線更強(qiáng)調(diào)整個模型訓(xùn)練流程的設(shè)計耦合，而不是“建一個模型+套一個優(yōu)化器”。

一旦這套方法能在大型Transformer/LLM上成功應(yīng)用，今后訓(xùn)練模型的效率和穩(wěn)定性無疑都將獲得極大提升。

論文唯一作者Jeremy Bernstein

最后簡單介紹一下論文作者Jeremy Bernstein。

個人主頁顯示，他本科和碩士階段都在劍橋大學(xué)三一學(xué)院接受物理理論與實(shí)驗(yàn)物理方面的教育，曾獲得NVIDIA研究生獎學(xué)金。

后來去了加州理工攻讀博士學(xué)位，研究方向轉(zhuǎn)為計算與神經(jīng)系統(tǒng)。

目前他既在Thinking Machines擔(dān)任機(jī)器學(xué)習(xí)研究員，還在MIT從事博士后研究工作，致力于揭示自然與人工智能的計算與統(tǒng)計規(guī)律，從而設(shè)計出更加高效、自動、實(shí)用的學(xué)習(xí)系統(tǒng)。

從谷歌學(xué)術(shù)論文頁面可以看到，其研究也大多集中在機(jī)器學(xué)習(xí)中的算法優(yōu)化、模型魯棒性、學(xué)習(xí)表示、以及優(yōu)化器設(shè)計等方面。

另外經(jīng)翁荔提醒，我們也去翻了翻Jeremy Bernstein給論文標(biāo)下的一些小腳注。

原來，雖然論文署名只有一人，但其背后實(shí)打?qū)嶋x不開很多研究員的跨界合作。比如這項工作就涉及到數(shù)學(xué)、計算機(jī)科學(xué)以及工程學(xué)的大融合。

包括Jeremy Bernstein本人，從當(dāng)初的物理專業(yè)逐漸轉(zhuǎn)向數(shù)學(xué)、機(jī)器學(xué)習(xí)，怎么不算跨界帶頭人呢（doge）~

One More Thing

目前Thinking Machines已發(fā)布兩篇研究論文。

首項研究成果發(fā)表于今年9月10日，研究主題為“Defeating Nondeterminism in LLM Inference”，克服大語言模型推理中的不確定性。

主要討論的內(nèi)容就是，為什么大模型每次的推理結(jié)果總是難以復(fù)現(xiàn)？根源在于批次不變性。

不過上次的作者名單相當(dāng)簡略，由一位華人面孔小哥Horace He主導(dǎo)，尚不清楚背后有沒有Jeremy Bernstein的身影。

以及再補(bǔ)充一點(diǎn)，之前被曝加入Thinking Machines的清華姚班校友、普林斯頓教授陳丹琦，曾經(jīng)也轉(zhuǎn)發(fā)了這篇論文。

并且陳丹琦團(tuán)隊最新也帶來了一篇關(guān)于可驗(yàn)證獎勵強(qiáng)化學(xué)習(xí)的論文，標(biāo)題是《Language Models that Think, Chat Better》，主要觀點(diǎn)為：

可驗(yàn)證獎勵強(qiáng)化學(xué)習(xí)（RLVR）在可驗(yàn)證領(lǐng)域之外也是有效的，并提出了“基于模型獎勵思考的強(qiáng)化學(xué)習(xí)”（RL with Model-rewarded Thinking，RLMT），以實(shí)現(xiàn)通用聊天能力。

論文傳送門已放文末，這里不再過多展開。

最后，肉眼可見Thinking Machines這家明星創(chuàng)業(yè)公司正在加速轉(zhuǎn)動，隨著研究成果不斷涌出，第一個產(chǎn)品是不是也快了？

嗯，0產(chǎn)出估值就已沖破120億美元（約合人民幣840億元），毫無疑問值得期待一波。

論文：
https://thinkingmachines.ai/blog/modular-manifolds/
陳丹琦團(tuán)隊新作：
https://arxiv.org/pdf/2509.20357