摘要：论文主题为“Modular Manifolds”，通过让整个网络的不同层/模块在统一框架下进行约束和优化，来提升训练的稳定性和效率。
明星创业公司Thinking Machines，第二篇研究论文热乎出炉！
公司创始人、OpenAI前CTO Mira Murati依旧亲自站台，翁荔等一众大佬也纷纷转发支持：
论文主题为“Modular Manifolds”，通过让整个网络的不同层/模块在统一框架下进行约束和优化，来提升训练的稳定性和效率。
之所以进行这项研究，主要是为了解决神经网络训练中的一个基本挑战：
网络内部的张量（如权重、激活、梯度）若数值过大或过小，会引发不稳定、梯度爆炸/消失、训练效率低下等问题。
因此，论文唯一作者Jeremy Bernstein提出了一种新的优化思路——模块化流形（Modular Manifolds），不仅对单个权重张量施加约束，还能把整个网络视为组合的流形结构，从而统一设计学习率、约束方式与优化逻辑。
网友们的反应be like：
对神经网络训练有深刻的见解。
将权重更新限制在流形上可以带来稳定性和可解释性，好奇能否扩展到非常大的模型？
为什么需要流形约束？
原因正如开头所提到的，在训练大模型时，如果权重、激活值、梯度太大或太小，就会出现训练不稳定（溢出、消失、收敛速度慢等问题）。
之前为了解决这些问题，归一化方法逐渐成为“黄金标准”，包括激活归一化（如层归一化LayerNorm）、梯度归一化（优化器内部的规范化）等，但很少直接对权重矩阵本身进行归一化处理。
而作者认为，对权重矩阵进行归一化处理具有相当价值，其好处可能包括：
更容易把握优化过程中更新量的大小；能够避免权重的范数变得过大，也就是防止权重“爆炸”的问题；让研究员可以把调整超参数的精力，更多地放在那些对模型影响最大的张量上；让矩阵的条件数变小，这样矩阵的行为就会更加稳定和可预测；……一句话，给权重矩阵归一化，可以让模型训练更稳定、更容易调整、行为更可预测，并且对外界干扰更有抵抗力。
基于此，作者希望设计一个几何化框架，把神经网络的权重参数约束在特定的Stiefel流形上，从而在训练时能够联合优化器与这些约束。
其核心研究过程大致有以下几个步骤：
假设要训练的参数是一个向量W，并且强制它始终在单位球面上：||W||=1。
如果用普通的Adam（自适应矩估计）/SGD（随机梯度下降）更新，更新后的向量可能就跑出球面。
为了解决这一问题，流形优化自有其“套路”：
先把梯度投影到切空间（球面上某点的切平面），保证更新方向合理；然后更新参数；最后用Retraction投影，把更新后的向量“拉回”到球面上。而在这个过程中，我们核心需要思考两个问题：一是应该采用哪种流形约束，二是应该如何测量长度。通过做出不同的选择，最后可以创建出不同的优化算法，如下表所示：
Step 2：推广到矩阵参数
紧接着，作者将上述思路从向量（球面）推广到矩阵。
由于Transformer的权重矩阵（把输入的向量转换成输出的向量）动辄成千上万维，直接优化容易失控，因此作者提出把矩阵参数放在Stiefel流形上。
在Stiefel流形上的矩阵，它们的列向量都是正交的，而且条件数（一个衡量矩阵稳定性的指标）为1。
这样做的好处在于：
其一，向量正交可以确保不同方向上的信号不会相互干扰，这有助于模型学习到更加独立的特征；
其二，条件数为1意味着矩阵不会放大或缩小向量太多，这有助于保持数值稳定性，避免在计算过程中出现大的数值误差。
选定流形后，作者通过一系列步骤最终设计了一个适用于Stiefel流形的优化算法——流形Muon算法。
其流程大致包括：
在Stiefel流形的切空间上算梯度更新；用一些矩阵运算（奇异值分解、sign函数等）确保更新合法；再把结果投影回Stiefel流形。Step 3：小规模实验验证
为了验证算法的合理性，作者随即进行了一个小规模实验（每次训练运行不到一分钟即可完成）。
他在CIFAR-10数据集上（一个包含10个类别的彩色图像数据集）训练了一个小MLP，对比了manifold Muon和AdamW这两种算法。
结果发现，前者在训练/测试准确率上略优于AdamW，且权重矩阵的奇异值都更集中、幅度稳定。
不过，虽然效果有提升，但每一步的运行时间倒是比AdamW稍慢一点（额外开销在dual ascent计算、投影等环节）。
作者表示，后续通过优化dual ascent步数、引入动量等方法，可以进一步减少额外开销。（意思是问题不大）
Step 4：最终推出“模块化流形”概念
以上验证了单个矩阵放到Stiefel流形上的有效性，那么问题来了——
多个层、多个模块怎么办？
对此，作者最终提出了模块化流形（Modular Manifolds） 这一概念。
具体而言，作者把神经网络里的每一层或每一个模块都看成是一个单独的“流形”，它有自己定义的范数和专属的优化方法。
当这些模块组合成一个完整的网络时，就把它们的流形通过笛卡尔积拼接在一起，形成一个大的流形空间。
在这个大空间里，作者采用最大范数（max norm）的方式来统一分配学习率，相当于设置一个全局的更新上限，避免不同模块的步伐不一致。
这样一来，每一层仍然可以按照自己的规则在小空间里更新参数，但整个网络的更新过程又被全局机制约束住，从而保持协调和稳定。
总之，这条研究路线更强调整个模型训练流程的设计耦合，而不是“建一个模型+套一个优化器”。
一旦这套方法能在大型Transformer/LLM上成功应用，今后训练模型的效率和稳定性无疑都将获得极大提升。
最后简单介绍一下论文作者Jeremy Bernstein。
个人主页显示，他本科和硕士阶段都在剑桥大学三一学院接受物理理论与实验物理方面的教育，曾获得NVIDIA研究生奖学金。
后来去了加州理工攻读博士学位，研究方向转为计算与神经系统。
目前他既在Thinking Machines担任机器学习研究员，还在MIT从事博士后研究工作，致力于揭示自然与人工智能的计算与统计规律，从而设计出更加高效、自动、实用的学习系统。
从谷歌学术论文页面可以看到，其研究也大多集中在机器学习中的算法优化、模型鲁棒性、学习表示、以及优化器设计等方面。
另外经翁荔提醒，我们也去翻了翻Jeremy Bernstein给论文标下的一些小脚注。
原来，虽然论文署名只有一人，但其背后实打实离不开很多研究员的跨界合作。比如这项工作就涉及到数学、计算机科学以及工程学的大融合。
包括Jeremy Bernstein本人，从当初的物理专业逐渐转向数学、机器学习，怎么不算跨界带头人呢（doge）~
目前Thinking Machines已发布两篇研究论文。
首项研究成果发表于今年9月10日，研究主题为“Defeating Nondeterminism in LLM Inference”，克服大语言模型推理中的不确定性。
主要讨论的内容就是，为什么大模型每次的推理结果总是难以复现？根源在于批次不变性。
不过上次的作者名单相当简略，由一位华人面孔小哥Horace He主导，尚不清楚背后有没有Jeremy Bernstein的身影。
以及再补充一点，之前被曝加入Thinking Machines的清华姚班校友、普林斯顿教授陈丹琦，曾经也转发了这篇论文。
并且陈丹琦团队最新也带来了一篇关于可验证奖励强化学习的论文，标题是《Language Models that Think, Chat Better》，主要观点为：
可验证奖励强化学习（RLVR）在可验证领域之外也是有效的，并提出了“基于模型奖励思考的强化学习”（RL with Model-rewarded Thinking，RLMT），以实现通用聊天能力。
论文传送门已放文末，这里不再过多展开。
最后，肉眼可见Thinking Machines这家明星创业公司正在加速转动，随着研究成果不断涌出，第一个产品是不是也快了？
嗯，0产出估值就已冲破120亿美元（约合人民币840亿元），毫无疑问值得期待一波。
论文：
https://thinkingmachines.ai/blog/modular-manifolds/
陈丹琦团队新作：
参考链接：
[1]https://x.com/miramurati/status/1971624309262373359
[2]https://jeremybernste.in/
[3]https://x.com/jxbz/status/1971703483767435446
— 完 —
量子位 QbitAI
来源：量子位一点号