在弯曲权重空间中学习：用于改进优化的指数-线性权重重参数化

# 指数-线性权重重参数化以优化训练
来源：https://arxiv.org/html/2607.09967
###### 摘要

许多神经网络操作本质上是乘性的而非加性的：将范数减半或加倍在相对意义上类似，但采用线性步长时所需的优化距离却不同。自适应优化器（如 Adam）按坐标归一化更新，但更新步长仍然是加性的；幅度差异很大的权重会获得大小相似的绝对变化，从而产生非常不同的相对扰动。我们提出 SymExpLin (SEL)，这是一种神经网络权重重参数化方法，它结合了符号感知的对称指数通路和类似恒等映射的线性通路。对称指数通路对于小的原始权重近似线性，但在幅度较大时弯曲程度增加。对数空间中的加性更新会映射到有效权重空间中与幅度成比例的变化。线性通路通过变换提供了一条直接路径，我们假设这能稳定优化，而可学习的尺度、曲率和偏移参数则控制通路之间的平衡以及指数通路的曲率。这些组件创建了一个弯曲的参数空间几何结构，实证表明，与标准线性参数化相比，它能更快地降低损失。我们还发现了一种有用的*不匹配初始化*：原始权重的选择使得变换的对称版本匹配 Xavier 统计量，但训练使用非对称的前向变换，使得正权重保持完整强度，而负权重的幅度变小；在小模型消融实验中，这改善了早期优化，并可能起到一种对称性破缺的作用。我们在 OpenWebText 上训练了九种宽度×深度的 Transformer 配置，SEL 在 1.32–1.49×更少的训练步数内达到匹配的验证损失，其中宽度最大的配置收益最大。SEL 是实用的：在我们最大的性能分析配置中，1.44×的步数减少和 5.5% 的每步开销对应约 1.37×的墙钟加速，并且训练后参数化可以折叠成标准的线性权重，没有推理成本。

## 1 引言

标准神经网络训练以加性方式优化参数，采用线性更新步长。然而，网络中的许多计算本质上是相对的。将归一化增益、特征幅度或矩阵的奇异值加倍或减半反映了对数空间中的类似变化。但是，它们在线性权重空间中的优化距离却不相等。具体来说，从 1 移动到 0.5 的线性距离是从 1 移动到 2 的一半。这并不是说我们应该将网络的各个方面都视为乘性的。激活函数和偏置认为绝对尺度很重要。尽管如此，训练中使用的优化大多利用了网络的线性/绝对尺度特性。当要通过加性更新实现权重的相对变化时，只要任务需要大的乘法因子，我们就应该预期训练会将一部分参数推离初始化点很远，而将许多其他参数留在其初始尺度附近。

如图 1 (https://arxiv.org/html/2607.09967#S3.F1) 所示，我们可以在预训练语言模型中观察到这种效应，它们会形成重尾权重分布。许多参数保持在初始化尺度附近，而一小部分则移动到尾部很远。在我们对 Qwen3-4B-Instruct 的分析中（图 1 (https://arxiv.org/html/2607.09967#S3.F1)），最大权重约为 Xavier 初始化尺度的 30 倍，尽管标准差仅增加了 1.7 倍。这创建了一个异构的优化问题，其中许多坐标可能从其原始位置移动很少，而一些坐标则必须在权重空间中遍历大得多的距离。

自适应优化器（如 Adam）[Kingma and Ba, 2014 (https://arxiv.org/html/2607.09967#bib.bib21)] 按坐标归一化更新幅度，使得各维度的更新幅度相似，从而使更新对参数幅度不变。预条件将步长与梯度尺度解耦以保持优化良好条件，但参数 0.004 和参数 0.4 可以接收到相似的绝对更新，同时导致非常不同的相对扰动：相同的步长对于一个小参数可能是显著的，但对于一个大参数在相对意义上却几乎不移动。全局学习率可能被选择为尊重最敏感的小尺度坐标，从而可能无法充分服务于较大幅度的方向，如特征增益、矩阵模式或其他可能受益于相对移动的坐标。

一种解决方案可能是在优化器层面通过参数幅度缩放来纠正这一点，但这可能与预条件的益处相冲突。因此，我们提出改变权重本身的参数化（从而改变几何结构）。纳入乘法相对性的自然方法是利用指数映射，它将原始空间中的加性更新转换为输出空间中的乘性更新。然而，当我们的权重大致以 0 为中心且对称时，标准指数函数仅输出正值。为了解决这个问题，我们使用了 symexp 函数，它是保号的、奇函数、单调的，并且具有指数行为。为了不过度承诺指数结构，我们引入了一个线性通路以及可学习的权重，允许网络缩放指数和线性分量的效果。我们还假设线性通路可能提供一条直接的梯度路径，这可能有助于优化。

每个有效权重被分解为两个互补的通路：非对称指数通路和线性通路。在我们推荐的配置中，有效权重为：

weff = sign(w) [Eθ(|w|) + Lθ(w)] + Δφ   (1)

其中 Eθ(|w|) 是可学习控制的 symexp 幅度，Lθ(w) 是类似恒等映射的线性通路，Δφ 是可选的加性残差，通常实现为低秩的 LoRA 风格修正。符号包装将弯曲几何应用于 |w|，同时将符号保留为可分离的自由度。我们进一步发现，初始化方式与前向传播组合这些通路的方式之间的刻意*不对称*会带来额外的收益，我们在第 4 节 (https://arxiv.org/html/2607.09967#S4) 中对此进行了分析，并将其归因于早期的对称性破缺效应。

贡献：

- • 我们提出了 SymExpLin (SEL)，一种非线性权重重参数化方法，可应用于任何可学习的权重矩阵或偏置向量，将有效权重分解为对称指数和线性分量，并具有独立可学习的结构化尺度。
- • 我们识别了*不匹配初始化*，即为权重反演使用与前向传播不同的组合规则，作为我们小规模消融实验中有益的非对称偏差来源。
- • 我们引入了具有各种分解模式的结构化尺度参数，以及用于管理冗余指数控制的目标学习率退火调度。
- • SEL 实现友好：训练变换是逐元素的，在编译实现中开销适中，而有效权重在训练后可以预先计算，降级为标准线性层，*零推理成本*。
- • 我们注意到，SEL 可以在训练后折叠到标准权重中，反之亦然，标准权重也可以反演回 SEL 参数空间，用于未来的持续训练或微调实验。
- • 我们在模型尺度网格 {1024,2048,3072}×{12,24,36} 上评估自回归语言建模，并进行针对初始化不对称性、线性通路和尺度学习率退火的目标消融实验。

## 2 相关工作

##### 权重重参数化。

权重归一化 [Salimans and Kingma, 2016 (https://arxiv.org/html/2607.09967#bib.bib3)] 将权重幅度与方向分离，改善了优化条件。谱归一化 [Miyato et al., 2018 (https://arxiv.org/html/2607.09967#bib.bib4)] 约束谱范数以保证训练稳定性。两者都应用了相对简单的变换。SEL 引入了一种更丰富的非线性重参数化，混合了对称指数和线性分量以及可学习的混合系数，作用于权重矩阵本身而不是其范数。

##### Symlog 变换。

对称对数 sign(x) ln(|x|+1) 及其逆函数已用于强化学习中的预测目标 [Hafner et al., 2023 (https://arxiv.org/html/2607.09967#bib.bib5)]，处理跨越多个数量级的值。SEL 将类似的操作（参数化的对称指数）应用于*权重空间*而非预测目标，其可学习的尺度参数控制对称指数和线性分量。

##### 初始化方案。

Xavier [Glorot and Bengio, 2010 (https://arxiv.org/html/2607.09967#bib.bib6)] 和 Kaiming [He et al., 2015 (https://arxiv.org/html/2607.09967#bib.bib7)] 初始化在层间保持方差。数据依赖初始化 [Mishkin and Matas, 2016 (https://arxiv.org/html/2607.09967#bib.bib10)] 逐层匹配目标统计量。我们通过牛顿法反演进行的初始化确保*尽管*存在非线性重参数化，初始化时的有效权重仍与标准 Xavier 统计量匹配。我们的不匹配初始化不同于这些方案：它在一个变换规则下保持预期的 Xavier 尺度，然后故意在略微不对称的规则下评估相同的原始权重。

##### 学习率参数化。

最大更新参数化 (μP) [Yang et al., 2022 (https://arxiv.org/html/2607.09967#bib.bib11)] 提供了跨模型宽度的有原则的每张量学习率缩放，解决了我们在第 3 节 (https://arxiv.org/html/2607.09967#S3) 中讨论的初始化/优化器不匹配的一个方面。SEL 瞄准了一个互补的方面，在训练过程中而非初始化时按张量调整每参数的有效步长；两者并非相互排斥。

##### 乘法交互。

Jayakumar et al. [2020 (https://arxiv.org/html/2607.09967#bib.bib13)] 研究了乘法交互在神经网络中出现的位置及其益处。我们的结构化尺度参数采用 `row_col_mult` 模式，在权重缩放上实现了秩 1 乘法交互，是一种轻量级的结构化乘法调制形式。

乘法权重更新 (MWU) 方法 [Arora et al., 2012 (https://arxiv.org/html/2607.09967#bib.bib32)] 使相对尺度更新显式化：权重按 (1±ε)^{payoff} 缩放而不是加性移动，因此调整自然与当前幅度成比例。这一原理是许多算法（从 boosting 到在线凸优化）的基础，但在神经网络训练中采用有限，其中加性 SGD 风格更新仍是标准。SEL 并非以乘法方式更新。优化器仍然在原始参数空间中采取加性步长，但对称指数通路在有效权重层面诱导出相关效应：原始空间敏感性按 exp(β|w|) 缩放，因此相同的原始步长在较大幅度处会产生更大的有效权重位移。线性通路在原点附近保持加性行为，产生一个混合机制而非完全的乘性机制。这使得 SEL 成为一种在保持与标准梯度优化器和预条件兼容的同时，获得 MWU 部分相对尺度行为的方法。

## 3 动机

训练将参数从采样的简单起始分布移动到更广泛的最终分布。初始化时，方差保持方案根据矩阵宽度确定的典型幅度附近初始化参数。到训练结束时，我们观察到权重的分布趋向于重尾形状，一小部分参数存活在初始典型范围之外很远的地方（图 1 (https://arxiv.org/html/2607.09967#S3.F1)）。不同的方法系列响应这个轨迹上的不同点：μP 和相关学习率参数化从初始化侧典型幅度设置步长，这是一个依赖于模型维度的每张量汇总统计量，而 SEL 则根据每个参数的*当前*幅度动态调整有效步长，因此随着训练过程中分布的扩展，已经移动得更远的参数会获得相应更大的有效步长。本节讨论预条件、学习率调优以及相对于参数幅度的更新幅度。我们将 SEL 定位为沿第二个动态轴与 μP 互补的方法。

##### Adam 归一化消除了尺度的影响。

考虑 Adam 的更新规则。一阶和二阶矩估计为 mt = β1 mt−1 + (1−β1) gt 和 vt = β2 vt−1 + (1−β2) gt^2，更新为 Δw ∝ mt / √vt。在稳态条件下，梯度幅度近似恒定，vt ≈ E[g^2]，因此归一化后的更新 mt / √vt 的预期绝对值接近 1，因为除以了估计梯度的标准差。在 β1 = β2 的特殊情况下，两个移动平均跟踪相同的时间尺度，柯西-施瓦茨不等式将其收紧为一个硬边界：|mt / √vt| ≤ 1 在稳态下。除了稳态条件外，动量在梯度符号翻转时会受到破坏性干扰，从而降低分子。同时，vt 对较大值敏感，可能增加分母。关键点是不变性：将每个梯度统一缩放一个常数会使得 |mt / √vt| 不变，被 √vt 的归一化抵消，并且更新中没有项依赖于 |w|。每维度的变化反映了梯度分布的*形状*（噪声模式、异常值频率）而不是其整体尺度，因此更新集中在一个由学习率设定的常见幅度周围，对预期梯度幅度和参数幅度都不变。因此，参数 0.004 和参数 0.4 接收到的更新绝对值大小相似，对应着 100× 的相对扰动差异；全局学习率调优必须适应最敏感的小尺度坐标，可能无法充分服务于较大幅度的坐标。

##### 方差保持初始化

方差保持初始化在训练开始时提供了一个典型的参数幅度，也可能指导训练结束状态的常见幅度。Xavier 初始化使用 σ ∝ 1/√((fan_in+fan_out)/2)，因此较小的矩阵接收较大的初始化值，并