不同初始化的文本Transformer模型合并

24年3月论文约翰霍普金斯大学和Meta AI的论文“Merging Text Transformer Models from Different Initializations”。

最近关于基于单样本基于组合的模型合并研究表明，来自完全不同初始化的模型之间具有Transformer 最小点学习相似特征的程度，并提出一种模型合并技术来研究损失域中这些最小点之间的关系。该架构的细节，如残差连接、多头注意和离散顺序输入，需要特定的干预才能计算保持在同一函数等价类内的模型排列中。与掩码语言建模任务上训练或在语言理解基准上进行微调的多模型平均方法相比，这些模型与本文方法合并的最小点之间损失屏障更低。这些模型的最小点不那样尖锐和孤立，对未来合并单独训练的 Transformer 模型，本文提供了基础。

根据经验，神经网络权重之间的 线性插值 已成为一种重要的工具。在微调相同的大预训练模型的背景下，平均模型在 ImageNet 上实现了最先进的准确性（Wortsman，2022）。Wortsman（2022）和 Rame（2022）确定，如果微调模型位于单个低误差盆地中，那么权重平均的表现与集成方法类似。然而，不能保证微调模型（从相同的初始化开始）将驻留在相同的损失盆地中。

先前基于线性插值的模型合并工作重点，是改进用于两个网络隐藏单元对齐的算法，减少它们之间插值的障碍。Singh & Jaggi (2020) 开发了一种基于最优传输的方法，允许在一对 ResNet 网络之间进行线性插值 (He et al., 2016)。Entezari (2022) 使用基于模拟退火的方法 (Zhan et al., 2016) 来找到排列，使得在 MNIST 上训练的宽 MLP 可以以接近零的屏障进行线性插值。Ainsworth (2023) 为 MLP 和 ResNet 开发了几种基于排列的算法，并在加宽的 ResNet 上演示了零屏障线性模式连接。Stoica（2024）扩展了这项工作，并允许在每个模型中进行特征合并，消除内部特征冗余并改善模型合并结果。

如图是一个Transformer层的示例，其参数用蓝框概述，特定隐藏状态用圆圈表示，操作流程用箭头表示。- 表示点积运算，+ 表示加法。LN 指LayerNorm 模块。在每个权重矩阵中包括置换矩阵和逆置换矩阵，指示本文方法提出的操作。Pres 指的是残差排列，PMHA 指的是多头注意（MHA）排列，PFF 指的是对齐前馈层。

在寻找多头注意（MHA）参数中相应的神经元，即K（WK）、Q（WQ）、V（WV）和线性层（WO）权重时，提出几种计算潜在排列矩阵的方法。

对于每个K、Q和V权重，完整参数 W 被逻辑划分为 H 个注意头，每个注意头的输出维度 dk = dmodel/H（Vaswani ，2017）。为了对这些全权重矩阵应用置换并保持整个模型的功能等价性，置换必须分别对每个注意头进行操作，而不是置换注意头之间的特征。这是因为来自 MHA 的最终隐向量反映了每个头的结果串联，这些结果是用头 i 的权重 WKi 、 WQi 、 WVi 单独计算的。

由于模型是通过不同的初始化进行训练的，因此除了每个头内特征的对应关系之外，它们的注意头的对应关系也可能不同。这种区别如图的前两个排列所示，由注意头对齐的不同策略所产生的示例排列矩阵。每个 Pi 反映了注意头中特征的排列。

外部排列决定了最终排列中使用的先前计算的内部排列子集，以及连接它们的顺序，从而产生 2 阶段 MHA 排列。如下在算法 1 中展示了这种方法。LinearSumAssignment 指的是方法中遇到的特定分配问题。