Scaling Linear Mode Connectivity and Merging to Billion Parameter Pretrained Transformers¶

会议: NeurIPS 2026
arXiv: 2606.23607
OpenReview: https://openreview.net/forum?id=AT2vWyunMD
代码: https://github.com/VILA-Lab/Dual-Learned-Matching
领域: 预训练
关键词: 线性模式连通性 (LMC), 模型合并, Transformer 对称性, 双向学习匹配, 权重空间优化

一句话总结¶

针对独立预训练 Transformer 在原始权重空间直接线性插值存在巨大损失屏障的问题，本文系统形式化了 Transformer 保持网络功能的完整对称性族，并提出了双向学习匹配框架（LMC-DM），首次在十亿级参数模型上实现了低损失屏障的线性连通与模型融合。

研究背景与动机¶

深度神经网络权重不仅是决定输入输出行为的静态参数，更是承载海量数据、算力与工程经验的可重用构件。在预训练大模型时代，若能直接通过权重空间的算术插值实现模型合并与能力聚合，将极大节省重训练与集成推理的昂贵成本。线性模式连通性（Linear Mode Connectivity, LMC）为这种权重空间合并提供了坚实的理论依托——当两个独立训练的模型可以通过一条处于低损失盆地中的线性线段相连时，简单的凸组合即构成有效的模型融合。然而，由于初始化、数据顺序与训练超参数的随机扰动，两个功能相似的预训练模型往往在原始坐标系中落入几何上相距甚远的高维区域，导致朴素线性插值产生严重的损失屏障与性能崩塌。

导致这种表观隔离的核心根源在于神经网络中广泛存在的“保持功能的对称性”（Functionality-Preserving Symmetries）。在传统 MLP 与小卷积网络中，隐藏神经元的排列对称性已得到较多研究，但在现代 Transformer 架构中，对称性空间呈现出前所未有的丰富度与层级性：残差连接、归一化层、多头注意力（MHA）机制中的 QK 与 OV 子回路以及前馈网络（FFN）共同诱导了复杂的正交旋转、通道排列与可逆缩放变换。现有的权重匹配（Weight Matching）方法仅在参数层面做离散对齐，与任务损失函数脱节；而近期的学习匹配（Learned Matching）虽然引入了插值损失监督，却仅单向变换其中一个端点模型，将优化轨迹锚定在单一固定模型的恶劣局部几何中，且无法稳定扩展至十亿级大参数模型。

本文的切入角度是打破单侧对齐的固有约束，全面释放 Transformer 架构的内在对称自由度：一方面建立包含归一化吸收、残差流正交旋转、头间排列与头内可逆变换的完整代数表征，并通过严格的结构约束（Cayley 变换与极分解）确保梯度更新的数值稳定性；另一方面引入双向协同优化，让两个端点模型同时学习各自的对称变换，共同迁入共享的低损失连通盆地。核心 idea：通过形式化 Transformer 完整的保持功能对称性空间，并引入端点双向联合优化的双向学习匹配（LMC-DM），利用过参数化改善损失曲面几何，首次将线性模式连通性与无损合并扩展至十亿参数级预训练 Transformer。

方法详解¶

整体框架¶

LMC-DM 的整体流程旨在为两个独立预训练但结构相同的端点模型 $\Theta_A$ 与 $\Theta_B$ 寻找功能保持的连续与离散坐标变换 $\Pi_A$ 和 $\Pi_B$，使插值模型在任意混合系数 $\lambda \in [0, 1]$ 下均处于低损失状态。方法首先将模型的归一化层重参数化并吸收入相邻线性层，暴露出残差流上的全局正交自由度；随后利用基于权重的离散匹配初始化头排列与 FFN 通道置换；最后在结构化约束参数化下，双向联合优化两个端点的连续对称变量。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["独立预训练端点模型<br/>Checkpoint A 与 Checkpoint B"] --> B["归一化重参数化与吸收<br/>LN/RMSNorm 转无参形式并中心化残差流"]
    B --> C["离散对称性初始化<br/>基于回路相似度与双线性分配固定排列"]
    C --> D["连续对称性参数化<br/>残差 Cayley 正交变换 + 注意力 Polar 极分解"]
    D --> E["双向联合匹配优化 (LMC-DM)<br/>随机采样混合系数对两端对称变量梯度下降"]
    E --> F["共享低损失连通盆地<br/>平滑线性插值与无损模型合并"]

关键设计¶

1. 全面 Transformer 保持功能对称性表征：从归一化吸收到多级子空间变换 针对 Transformer 复杂的模块耦合，本文构建了一套完备且数学上严格等价的功能保持重参数化体系。首先，对于带可学习增益 $\boldsymbol{\gamma}$ 和偏置 $\boldsymbol{\beta}$ 的 LayerNorm 或 RMSNorm，本文利用通道中心化投影矩阵 $\mathbf{C} = \mathbf{I} - \frac{1}{d}\mathbf{1}\mathbf{1}^\top$ 将其转化为无参 RMSNorm 形式，并将缩放增益与偏置精准吸收入后接的线性变换权重中。在此无参形式下，残差流获得了全局坐标旋转自由度：对任意正标量 $g > 0$ 与全局正交矩阵 $\mathbf{G} \in \mathbb{R}^{d \times d}$（满足 $\mathbf{G}^\top \mathbf{G} = \mathbf{I}$），残差读操作 $\mathbf{W}_{\text{in}}$ 与写操作 $\mathbf{W}_{\text{out}}$ 分别按如下规则变换： $$\mathbf{W}_{\text{in}}' = \mathbf{G}^\top \mathbf{W}_{\text{in}}, \qquad \mathbf{W}_{\text{out}}' = \mathbf{W}_{\text{out}} g \mathbf{G}$$ 由于无参 RMSNorm 自动消除正标量缩放并保持正交变换，整个残差流等价于统一转动至新坐标系下。其次，对于多头注意力机制，除了常规的头间排列矩阵 $\mathbf{P}_h$，本文精确揭示了每个注意力头内部的双重连续可逆对称性：在 QK 回路中，任意可逆矩阵 $\mathbf{M}_{Q,i}$ 对 Query 的变换可被 Key 的逆转置变换完美抵消（$\mathbf{W}_{Q,i}' = \mathbf{W}_{Q,i}\mathbf{M}_{Q,i}, \mathbf{W}_{K,i}' = \mathbf{W}_{K,i}\mathbf{M}_{Q,i}^{-\top}$），使注意力 Logits 严格不变；在 OV 回路中，任意可逆矩阵 $\mathbf{M}_{V,i}$ 对 Value 的乘法亦可被 Output 投影的求逆操作消除（$\mathbf{W}_{V,i}' = \mathbf{W}_{V,i}\mathbf{M}_{V,i}, \mathbf{W}_{O,i}' = \mathbf{M}_{V,i}^{-1}\mathbf{W}_{O,i}$）。该设计从代数上完全打通了 Transformer 内部各子模块的等价空间。

2. 结构约束下的连续对称性参数化：Cayley 变换与极分解 若在梯度优化中采用无约束矩阵并在前向传播时进行强制正交/可逆投影，极易引发训练不稳定性与梯度发散；而矩阵指数（Matrix Exponential）计算代价高昂且在大规模参数下难以稳定反向传播。为此，本文设计了计算高效且数值稳定的显式结构参数化方案。对于全局残差空间正交矩阵 $\mathbf{G}$，采用 Cayley 变换进行参数化： $$\mathbf{Q} = (\mathbf{I} + \mathbf{S})(\mathbf{I} - \mathbf{S})^{-1}, \quad \mathbf{S}^\top = -\mathbf{S}$$ 仅需维护一个斜对称矩阵 $\mathbf{S}$ 即可天然满足正交约束且便于 GPU 矩阵求逆运算。对于注意力头内的可逆对称矩阵 $\mathbf{M}$，本文引入极分解（Polar Decomposition）风格参数化： $$\mathbf{M} = \mathbf{Q} (\mathbf{L} \mathbf{L}^\top), \quad \text{其中 } \mathbf{Q} = \text{Cayley}(\mathbf{S}), \ \mathbf{L} \text{ 为对角线严格为正的下三角矩阵}$$ 该分解将可逆矩阵拆分为 Cayley 正交旋转分量与基于 Cholesky 分解的对称正定（SPD）缩放分量，既严格保证了矩阵在优化全程的可逆性与正行列式，又赋予了模型比纯正交变换更丰富的表征搜索能力。

3. 双向对齐学习（Dual Learned Matching）：解耦坐标约束的过参数化优化 现有单向匹配将一个模型视为固定不变的坐标系基准，强迫另一个模型“跨越鸿沟”迁就前者，这种单侧优化极易陷入高度非凸的局部极小。本文提出双向对齐学习（LMC-DM），同时为端点 $\Theta_A$ 和 $\Theta_B$ 赋予独立的连续对称性参数 $\Pi_A$ 与 $\Pi_B$，构成联合插值模型： $$\Theta_{\text{interp}}(\lambda) = \lambda \pi_{\Pi_A}(\Theta_A) + (1-\lambda) \pi_{\Pi_B}(\Theta_B), \quad \lambda \in [0, 1]$$ 尽管从代数角度看，任意双向对称性均可严格折叠吸收为一个相对单向变换加上一个全局规范（Gauge）（即存在相对变换 $\mathbf{R}_G = \mathbf{G}_B \mathbf{G}_A^\top$ 使得两端插值完全等价于单侧插值），因此双向参数化在理论容量上并未增加任何额外自由度；但从优化动力学角度，双向参数化构成了关键的“过参数化”。它彻底解耦了端点之间的硬性几何锚定，将原本复杂的非凸相对映射分解为两个更易于梯度寻优的平滑子路径，促使两个检查点同时向曲率更平缓、各向同性更好的共享低损失盆地迁移。

损失函数 / 训练策略¶

在训练阶段，算法保持离散的头排列 $\mathbf{P}_h$ 与 FFN 神经元置换 $\mathbf{P}_f$ 冻结（由初始化阶段的回路相似度匹配确定，实验证实允许其更新并不会改变离散分配），仅对连续对称参数 $\Pi_A, \Pi_B$（包含标量 $g_m$、残差正交参数 $\mathbf{S}_m$、以及各头 QK/OV 极分解参数）进行梯度更新。在每个训练步骤中，从均匀分布中随机采样插值系数 $\lambda \sim \mathcal{U}(0, 1)$ 与训练批次 $\mathcal{B} \sim \mathcal{D}$，以插值模型的任务交叉熵损失作为目标函数： $$\mathcal{J}(\Pi_A, \Pi_B) = \frac{1}{|\mathcal{B}|} \sum_{(x, y) \in \mathcal{B}} \mathcal{L}(\Theta_{\text{interp}}(\lambda); x, y)$$ 优化器采用 AdamW 搭配余弦退火学习率调度，以极小算力预算（例如 ViT-B 与 Pythia-160M 仅需单张 A6000 训练数小时至一天，Pythia-6.9B 与 OLMo-7B 仅需 4 张 A6000 训练约 1 天）即可实现收敛。

实验关键数据¶

主实验¶

实验评估涵盖视觉 Transformer（ViT-Small/Base/Large）与自回归语言模型（Pythia 14M~1B），度量插值路径上的最大损失屏障（Loss Barrier，定义为插值路径最高损失与两端点平均损失之差，越低越好）。

模型类别 / 规模	评估数据集	原始插值 (Raw)	权重匹配 (WM)	学习匹配 (LM)	增强学习匹配 (单侧)	双向学习匹配 (LMC-DM, 本文)	屏障消除幅度 (相对 Raw)
ViT-Small	ImageNet-1K	6.58	6.58	2.18	1.20	1.11	-5.47 (-83.1%)
ViT-Base	ImageNet-1K	6.57	6.50	2.00	1.05	0.82	-5.75 (-87.5%)
ViT-Large	ImageNet-1K	6.47	6.48	1.23	0.71	0.66	-5.81 (-89.8%)
Pythia-14M	WikiText-103	9.89	8.17	1.26	1.00	0.86	-9.03 (-91.3%)
Pythia-160M	WikiText-103	6.28	6.04	0.39	0.27	0.18	-6.10 (-97.1%)
Pythia-410M	WikiText-103	4.99	5.65	0.57	0.44	0.37	-4.62 (-92.6%)
Pythia-1B	WikiText-103	4.60	4.75	0.42	0.28	0.23	-4.37 (-95.0%)

在端点精度保持方面，ViT-Large 在整个线性插值路径 $\lambda \in [0, 1]$ 上的 ImageNet Top-1 准确率始终保持在 69.40% 以上（两端分别为 85.82% 与 82.97%），而未经对齐的原始插值在中间点准确率直接暴跌至 0.06%~0.46%。

消融实验¶

消融实验在 ViT-Base 和 Pythia-160M 上深入检验了对称模块组成、连续参数化方式以及单向/双向设置的核心贡献。

配置变体	ViT-Base 损失屏障 ↓	ViT-Base 最差准确率 (%) ↑	ViT-Base 中点准确率 (%) ↑	Pythia-160M 损失屏障 ↓	Pythia-160M 最差困惑度 ↓	Pythia-160M 中点困惑度 ↓
LMC-DM 完整方案 (双向 + Cayley + Polar)	0.82	63.79	65.48	0.18	82.47	73.33
移除 QK 回路对称性 (w/o QK)	0.97 (+0.15)	60.68 (-3.11)	61.80 (-3.68)	0.18 (+0.00)	81.81 (-0.66)	76.85 (+3.52)
移除 OV 回路对称性 (w/o OV)	1.17 (+0.35)	56.85 (-6.94)	61.04 (-4.44)	0.37 (+0.19)	97.03 (+14.56)	82.97 (+9.64)
同时移除 QK 与 OV 对称性	1.48 (+0.66)	50.98 (-12.81)	54.96 (-10.52)	0.37 (+0.19)	97.06 (+14.59)	82.11 (+8.78)
移除全局残差缩放 (w/o global scale)	0.82 (+0.00)	63.50 (-0.29)	65.39 (-0.09)	0.23 (+0.05)	86.01 (+3.54)	77.53 (+4.20)
注意力参数化：无约束直接优化	0.97 (+0.15)	60.72 (-3.07)	63.68 (-1.80)	0.29 (+0.11)	90.20 (+7.73)	80.31 (+6.98)
注意力参数化：矩阵指数 (MatExp)	0.98 (+0.16)	60.26 (-3.53)	63.57 (-1.91)	0.30 (+0.12)	90.59 (+8.12)	78.06 (+4.73)
注意力参数化：纯正交 Cayley 变换	0.88 (+0.06)	62.42 (-1.37)	63.62 (-1.86)	0.24 (+0.06)	85.18 (+2.71)	71.86 (-1.47)
单侧匹配 (Single w. Absorb)	1.05 (+0.23)	58.20 (-5.59)	61.34 (-4.14)	0.27 (+0.09)	89.45 (+6.98)	78.12 (+4.79)

关键发现¶

OV 与 QK 回路对称性的主导地位：去掉注意力的 OV 或 QK 变换会导致屏障显著抬升。特别是在 ViT-Base 上，完全移除注意力回路对称性后损失屏障激增 0.66，中点分类准确率下挫 10.52%，表明注意力机制内部的连续对称性是打通 Transformer 模式连通性的最核心开关。
Cayley 与极分解的优越稳定性：在连续参数化对比中，直角坐标系下的无约束投影和矩阵指数均落后于 Cayley + Polar 组合。矩阵指数虽然覆盖全连续群，但在深层反向传播时梯度病态；而极分解不仅收敛稳定，且比纯正交参数化（Cayley）提供了必要的正定缩放裕度。
双向优化的优化红利：双向联合优化在没有引入额外有效自由度的情况下，系统性超越所有单向优化方案（ViT-Base 屏障从 1.05 降至 0.82，Pythia-160M 从 0.27 降至 0.18），强力印证了过参数化对平滑损失景观的几何增益。

亮点与洞察¶

代数等价与优化非等价的深刻洞见：论文从理论上严格证明了双向对称变换可以无损折叠为单向相对变换，但在梯度下降的世界里，双向优化却表现出截然不同的收敛动力学，打破了单侧锚定带来的恶劣几何阻碍。
结构化约束对可微对齐至关重要：避免粗暴的梯度后投影，利用 Cayley 变换和 Polar 极分解将流形约束内化为无约束自由参数，为超大模型权重对齐提供了兼顾严格约束与数值稳定的范本。
拓展至大模型权重复用的普适范式：此前学界普遍怀疑数十亿参数的大模型是否仍具备线性连通性，本文首次将 LMC 推进至 6.9B 参数规模，为大模型无训练/微算力模型融合开辟了全新路径。

局限与展望¶

下游零样本常识能力未达完全无损：在常识推理基准（PIQA、SciQ、ARC-E）上，尽管 LMC-DM 的插值模型远优于原始插值，但相比原始单模型仍有性能损耗，说明预训练自监督损失的低屏障尚未完全等价于下游复杂推理能力的无损保留。
依赖严格同构架构假设：当前方法依赖层数、隐藏维度、注意力头数与分词器完全一致的孪生架构，无法直接应用于结构异构、通道剪枝或不同分词空间的跨模型融合。
极端异质模型对齐仍具挑战：实验显示在不同训练语料体系下的 OLMo-7B 模型上，损失屏障仍相对较高（1.53），说明初始化与预训练分布的巨大偏离仍会限制对称性对齐的效果。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出双向联合学习匹配框架与 Polar 连续对称性参数化，拓展了 Transformer 完备对称性代数表征。
实验充分度: ⭐⭐⭐⭐⭐ 跨越 ViT 与 Pythia/OLMo 系列，覆盖 14M 到 7B 多尺度模型，提供严谨的插值全轨迹、多维消融与常识评测。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述自然流畅，数学推导完备自洽，图表信息充实且分析深入。
价值: ⭐⭐⭐⭐⭐ 为大模型权重空间几何研究与实用化模型无损合并提供了关键基石。