AI在线 AI在线

比Adam更有效,POET从谱不变原理出发,让LLM训练又稳又快

Zeju Qiu和Tim Z. Xiao是德国马普所博士生,Simon Buchholz和Maximilian Dax担任德国马普所博士后研究员,Bernhard Schölkopf是德国马普所所长,Weiyang Liu是香港中文大学计算机系助理教授。 随着大型语言模型(LLM)推动人工智能领域取得突破性进展,如何实现高效、稳定的超大规模模型训练,始终是该领域最富挑战性的核心议题之一。

Zeju Qiu和Tim Z. Xiao是德国马普所博士生,Simon Buchholz和Maximilian Dax担任德国马普所博士后研究员,Bernhard Schölkopf是德国马普所所长,Weiyang Liu是香港中文大学计算机系助理教授。

随着大型语言模型(LLM)推动人工智能领域取得突破性进展,如何实现高效、稳定的超大规模模型训练,始终是该领域最富挑战性的核心议题之一。

针对这一关键问题,研究者们提出了一种基于第一性原理的全新方法——POET(Reparameterized Training via Orthogonal Equivalence Transformation),该方法通过重参数化优化策略,旨在从第一性原理出发提升训练效率与稳定性。

图片

Paper:Reparameterized LLM Training via Orthogonal Equivalence Transformation

Project page:https://spherelab.ai/poet/

Arxiv:https://www.arxiv.org/abs/2506.08001

POET:基于第一性原理的大型语言模型全新训练范式

POET 的关键思想是:通过对每个神经元进行结构性重参数化,引入两个可学习的正交矩阵以及一个固定的随机权重矩阵,从而构建一个正交等价的变换结构。该方法在训练过程中严格保持权重的奇异值分布,并天然拥有较低的球面能量,这是 POET 有效性的核心来源。

通过联合建模奇异值不变性与最小超球能量,POET为大模型训练提供了一种兼具物理解释性与泛化能力的新范式。由于该方法严格保持权重矩阵的谱结构,不仅能稳定优化过程,还显著提升了模型的泛化性能。为兼顾计算效率与实用性,研究者还开发了高效的近似算法,使POET可扩展至超大规模神经网络训练。实验结果表明,该方法在大型语言模型训练中表现出卓越的性能与可扩展性。

图片

图 POET 的三个学习阶段:左—示意图;中—角度;右—损失值与验证。

谱性质与泛化

当前训练大型语言模型的事实标准是直接使用Adam优化器对权重矩阵进行更新。尽管这一做法实现简单,但在计算上往往代价高昂,随着模型规模的扩大,其复杂度迅速增长。此外,该方法对超参数极为敏感,需精细调整以保证训练稳定收敛。

更为关键的是,即便训练损失已经被有效最小化,模型的泛化性能仍可能表现不佳。为缓解这一问题,本文提出了多种权重正则化与归一化技术,其核心目标往往可归结为:显式或隐式地改善权重矩阵的谱结构(即奇异值分布)。

从直观角度看,权重矩阵的谱范数(最大奇异值)描述了其对输入向量的放大上界,因此与模型的平滑性和泛化能力密切相关。一般认为,较小的谱范数(意味着更温和的变换)往往有助于提升泛化性能。这一观点促使越来越多研究致力于对谱性质进行精细控制。理论研究亦表明,若能有效约束权重矩阵的谱结构,便可形式化地为模型提供泛化上的保证。

谱保持(Spectrum-preserving)权重更新

为在避免上述局限的同时实现有效的权重量谱控制,文章提出一种基于正交等价变换(OrthogonalEquivalenceTransformation)的重参数化训练算法,用于间接学习权重矩阵。具体而言,POET 将权重矩阵图片 重参数化为图片,其中图片为随机初始化的权重矩阵,图片图片为两个可学习的正交矩阵。在训练过程中,POET 不直接优化权重矩阵,而是保持随机初始化矩阵图片不变,通过学习正交矩阵图片图片来变换图片。这种重参数化策略在允许奇异向量灵活调整的同时,能够严格保持权重矩阵的奇异值谱,从而实现谱保持的权重更新。

图片

图 LLaMA模型中同一权重矩阵奇异值的训练动态。左图为标准训练,严格遵循大型语言模型的常规做法(使用AdamW直接优化);右图为POET,其采用本文提出的近似方法以支持大规模LLM训练。POET的奇异值仅出现轻微(几乎可忽略)的变化,主要归因于数值误差和近似误差。

奇异值谱的训练动态

受 Muon [4]的启发,研究者对 AdamW、Muon与 POET 的奇异值谱进行了谱分析。在训练的不同迭代点,可对训练后的模型计算 SVD 熵 

图片

该指标用于衡量奇异值的多样性;熵值越高,表示谱分布越均匀、越丰富。[4] 将 Muon 相较于 AdamW 的优越性能归因于其权重矩阵更新所带来的更丰富谱分布。正如下图所示,由于采用正交等价变换,POET 在整个训练过程中始终保持较高的谱多样性。

图片

POET方法具备两项核心优势:

  • 高效的谱控制

由于正交变换并不改变权重矩阵的奇异值,POET在训练全程都能保持权重谱与随机初始化矩阵一致——即便采用近似实现,这一点也已得到实证验证。借助恰当的初始化方案,POET可直接约束奇异值分布,避免标准LLM训练后权重出现过大的奇异值。为进一步增强算法效果,研究者们提出了两种新初始化策略:归一化高斯初始化(normalizedGaussianinitialization)和 均匀谱初始化(uniformspectruminitialization),均可确保生成的权重矩阵具有有界奇异值。

  • 高效近似

直接进行POET训练的计算开销较高,但方法本身的灵活性为高效、可扩展训练提供了空间。针对大规模正交矩阵优化这一关键难题,文章提出两级近似方案:

随机基元优化:将大正交矩阵分解为若干参数量更少的基元正交矩阵,并结合“合并再初始化”策略提高效率;

基于Cayley‑Neumann参数化的近似正交性:通过 Neumann 级数近似 Cayley 正交参数化,以较低计算成本保持正交性,同样借助“合并再初始化”策略抑制误差累积。

LLaMA架构的大规模语言模型预训练

本文在多种规模的LLaMATransformer(60M、130M、350M、1.3B 参数)上对POET进行了预训练实验。使用的数据集为C4——从CommonCrawl清洗得到的网页语料,已被广泛用于大型语言模型的预训练。下文汇总了实验结果,报告了验证困惑度(perplexity)及可训练参数量。

图片

图片

图 AdamW和POET在模型规模为350M和1.3B下的可训练参数规模及验证困惑度(perplexity)。

训练加速

为突出POET在性能上的显著改进,文章将AdamW的训练步数(即模型实际看到的token数量)大幅提升至原来的近三倍。即便如此,采用 b=1/2 设置的POET‑FS仍在性能上超越AdamW。

图片

参数与内存复杂度

通过将超参数 b 作为采样预算引入,完全随机 SPO(StochasticPrimitiveOptimization)成功将参数复杂度与权重矩阵规模解耦。当 b 取较小值时,POET 的参数效率显著提升,但收敛速度有所下降,为使用者提供了效率与速度之间的灵活权衡。相比之下,块随机 SPO 的参数复杂度与矩阵尺寸(m+n)成正比,因而较 AdamW(需要 mn 个可训练参数)更具可扩展性。在内存占用方面,只要采样预算 b 设置得当,两种 POET 变体均可显著优于 AdamW。下文给出了参数与内存复杂度的详细对比。

图片

训练算法

步骤1:权重初始化使用归一化高斯初始化为权重矩阵赋值:图片

步骤2:正交矩阵初始化

完全随机SPO(fullystochasticSPO):随机采样索引集合 图片,并使用CNP(Cayley‑NeumannParameterization)对图片图片进行参数化。图片图片二者均以单位矩阵开始。

块随机SPO(block‑stochasticSPO):随机采样置换矩阵图片图片,同样采用CNP对图片与进行参数化,并将它们初始化为单位矩阵。同样,图片图片二者均以单位矩阵开始。

步骤3:高效正交参数化

对于完全随机SPO可得:图片图片

对于块随机SPO可得:图片图片

步骤4:正交矩阵内层训练循环更新

前向传播中的等效权重矩阵为图片

反向传播通过图片图片计算梯度,进而更新完全随机SPO图片的或块随机SPO中的图片;该内循环迭代次数固定。

步骤5:合并并重新初始化(merge‑then‑reinitialize) 

将已学习的正交矩阵合并进权重:图片

若训练未结束,则返回步骤2,重新初始化图片图片,继续下一阶段训练。

POET的优异表现来自于超球能量与谱保持

神经元初始化

鉴于 POET 在训练过程中会保留初始权重矩阵的谱特性,初始化策略显得至关重要。文章运用了归一化高斯初始化:先从零均值、固定方差的高斯分布中抽取神经元权重,再对其进行归一化。下表对多种随机初始化方案进行了实证比较,结果显示归一化高斯初始化取得了最佳最终性能。研究者推测,这一优异表现源于 POET 在该初始化下能够在训练过程中同时保持超球能量与谱特性。

图片

训练中的超球能量

超球能量 HE 用于衡量神经元在单位超球面上的均匀分布程度,可作为刻画各层神经表征的一种度量。文献[2,3]表明,满足正交约束的训练过程可在训练期间保持这一超球能量不变,从而避免表征退化并提升泛化性能。

归一化高斯初始化下的POET 可同时保持能量与奇异值分布

在零均值、各向同性的高斯初始化条件下,POET 能够同时实现谱保持训练与能量保持训练。这一特性为归一化高斯初始化方法的最优性能提供了理论解释(详细证明参见附录 B)。

图片

POET训练机理解析

为深入理解POET的运行机制,我们用向量探测(vector probing)分析正交矩阵的学习动态。具体做法是:固定随机生成的单位向量图片,计算图片,即图片图片之间的余弦相似度,以评估正交矩阵图片图片的演化。

对七个可学习的正交矩阵在训练过程中的余弦相似度进行跟踪后,可以将其学习过程划分为三个阶段(见图1):

  • 锥壳搜索阶段(Conical shell searching)

余弦相似度从1(即图片为单位矩阵)逐渐下降并收敛到[0.60,0.65]。这一现象在所有正交矩阵上均一致,说明图片图片映射到其原始方向附近的狭窄锥壳上。

  • 锥壳上的稳定学习阶段(Stable learning on the conical shell)

余弦相似度保持在该区间内不再显著变化,但模型开始进入稳定学习期;尽管余弦值趋于稳定,验证困惑度仍在线性下降。

  • 最终阶段微调(Final adjusting)

随着学习率逐步衰减至零,学习速度放缓并最终停止。

更为详尽的讨论与实证结果见论文附录。

相关资讯

东京大学 | Adobe 提出InstructMove,可通过观察视频中的动作来实现基于指令的图像编辑

InstructMove是一种基于指令的图像编辑模型,使用多模态 LLM 生成的指令对视频中的帧对进行训练。 该模型擅长非刚性编辑,例如调整主体姿势、表情和改变视点,同时保持内容一致性。 此外,该方法通过集成蒙版、人体姿势和其他控制机制来支持精确的局部编辑。
1/20/2025 10:36:00 AM
AIGC Studio

阿里发布新ID保持项目EcomID, 可从单个ID参考图像生成定制的保ID图像,ComfyUI可使用

阿里妈妈发布了一个新的ID保持项目EcomID,旨在从单个ID参考图像生成定制的保ID图像,优势在于很强的语义一致性,同时受人脸关键点控制。 EcomID 方法结合了 PuLID 和 InstantID 的优点,以获得更好的背景一致性、面部关键点控制、更真实的面部以及更高的相似度。 目前,EcomID 的 Comfyui 原生实现官方插件也已经发布,大家可以从文章中的链接获取。
1/20/2025 11:00:00 AM
AIGC Studio

o1推理框架最新成果:斯坦福&伯克利提出元链式思维,升级模型推理能力

o1背后的推理原理,斯坦福和伯克利帮我们总结好了! 在最新的一篇长达100页的论文中,他们将o1模型背后的推理机制提炼成了一个通用的框架——元链式思维(Meta-CoT)。 这个元链式思维(Meta-CoT)到底是什么意思呢?
1/20/2025 1:08:25 PM
量子位
  • 1