基于任务向量减法实现模型融合,CPU 上几分钟完成,无需训练数据,是 Hugging Face 多数榜单模型的真实制作方式。
一个训练好的 checkpoint 不过是一个超长的数字列表。因此,从同一个基座模型微调出来的两个模型存在于同一个权重空间中,处于相同的坐标位置,你可以对它们做算术运算。这一条观察就是模型合并的全部原理——而 Hugging Face 排行榜上很大一部分开源权重模型实际上就是这样产生的。
另一种方案是在两个数据集的并集上重新训练:需要 GPU、需要数据访问权限、需要数小时到数天。合并不需要训练数据、在 CPU 上运行、只需几分钟。
从微调模型中减去基座模型,基座模型已经知道的一切就会相互抵消。剩余的部分就是微调所改变的内容:
tau_a = {k: ft_a[k] - base[k] for k in base} # the "sentiment" direction
tau_b = {k: ft_b[k] - base[k] for k in base} # the "spam" direction
# reconstruction is exact by construction
assert torch.allclose(base[k] + tau_a[k], ft_a[k])
这个差值就是 task vector,而引人注目的经验事实是:它表现得像一个可以操作的物体。它指向一个一致的方向,它的长度控制着技能被表达的强度,不同任务的向量大致上是正交的。它的绝大多数元素都接近于零——微调只移动了很少的参数。这种近乎正交性正是为什么一个模型可以同时承载两项技能的原因。
从几何性质出发有一个硬性要求:这些模型必须共享一个预训练祖先,且架构和分词器相同。两个从零开始用不同随机种子训练的模型可能计算出相同的函数,但权重却截然不同——神经元被置换、符号被翻转——所以对它们求平均只会产生噪声。对同一个基座模型的微调从同一点出发,并保持在同一个低损失 basin 中,而这正是线性模式连通性(linear mode connectivity)这一特性让这一切成为可能。
然后是一行代码的合并,以及那个旋钮
最粗暴的合并方式是朴素的逐参数求平均,即 model soup:
def soup(models, weights=None):
w = weights or [1 / len(models)] * len(models)
return {k: sum(wi * m[k] for wi, m in zip(w, models)) for k in models[0]}
用 task vector 的形式写出来就是 base + 0.5*tau_a + 0.5*tau_b——系数被强制求和为一的任务算术。出人意料的是,这往往已经是大部分收益了。在下面的演示中,一个情感分类器在其自身任务上得分 100%,在垃圾邮件任务上得分 32%;一个垃圾邮件过滤器则反过来是 100%/42%。它们的 soup 得分是 97% 和 98%——一组权重,两个任务都搞定。
给每个向量分配自己的系数,你就得到了一个混音台:
def task_arithmetic(base, taus, lambdas):
return {k: base[k] + sum(l * t[k] for t, l in zip(taus, lambdas))
for k in base}
Lambda 是最重要的旋钮,而且它不是免费的。在演示中把两个都推到 1.0,准确率就会跌到 89%——低于 0.5 的 soup——因为两个向量的任务外噪声现在是叠加而非平均。扫描曲线在 lambda ≈ 0.4 附近达到峰值 99.2%。每一个真实的合并配方最后都会在某个 held-out eval 集上做超参数搜索。
因为一项技能是一个方向,改变符号就是让模型背离它:
forget = {k: base[k] - 1.0 * tau_toxic[k] for k in base}
在演示中,这把情感准确率从 50% 压低到 20%——低于抛硬币的概率,因为模型不仅仅是遗忘了这个任务,而是被赋予了反向技能,现在在预测相反的标签。另一个任务则纹丝不动。同样的技巧被认真地用于去毒化、剥离记忆化的风格、或从发布的 checkpoint 中移除某种能力。这是有据可查的最便宜的遗忘形式,尽管它是一把钝刀:过大的负 lambda 也会损害通用能力。
合并以两种特定方式失败。
冗余——一个 task vector 中绝大多数元素都是噪声级别的。它们什么都贡献不了,但一旦你做平均,就会稀释那些重要的元素。
符号冲突——在两个技能都关心的参数上,它们往往想要相反的方向。在演示中,单词 "service" 从情感向量得到 +0.42(评论中的赞扬)和从垃圾邮件向量得到 -0.73(收件箱中的Routine)。平均下来得到 -0.15:一个两个模型都不想要的数值,弱到两项技能都在这个参数上受损。
这就是为什么合并越来越多的模型反而越来越差,也是接下来三种方法各自从不同角度攻击的目标。
TIES 通过三遍处理修复了两个问题。
stack = torch.stack([t[k].flatten() for t in taus]) # [n_models, P]
# 1. TRIM - keep top-k% by magnitude, per model
keep = int(k * stack.shape[1])
thresh = stack.abs().kthvalue(stack.shape[1] - keep + 1, dim=1, keepdim=True).values
stack = stack * (stack.abs() >= thresh)
# 2. ELECT - one sign per parameter, by total magnitude
elected = stack.sum(0).sign()
# 3. DISJOINT MERGE - mean over the entries that agree
agree = (stack.sign() == elected) & (stack != 0)
merged = (stack * agree).sum(0) / agree.sum(0).clamp(min=1)
Trim 删除了冗余。Elect 让一个方向彻底胜出,而不是被平均掉。不相交合并只对与 elected 符号一致的条目求平均,因此赢家保留了完整的量级——那个有争议的 service 参数最终是 -0.73,而不是 -0.15。在演示中 TIES 是棋盘上最好的合并方式,达到 99.2%,击败了朴素的 soup,而且与求平均不同,它会随着你添加模型而保持扩展。但如果 trim 过度,你就会开始删除真实的信号:在 k = 20% 时演示跌到 83%。
DARE 把冗余观察推向极端——删除每个 task vector 中随机的一部分,然后重新缩放存活者,使期望向量保持不变:
def dare(tau, p=0.9):
mask = torch.rand_like(tau) >= p # keep each entry w.p. 1-p
return tau * mask / (1 - p) # rescale so E[.] is unchanged
你可以丢弃一个微调 delta 的 90%–99%,而模型仍然能正常运行——这本身就说明微调实际改变的东西少得惊人。对于合并来说,收益是几何级的:两个各自只有 10% 密度的向量在约 1% 的参数上冲突,而不是 100%,所以干扰几乎消失。注意丢弃是随机的,不是基于量级的——那是 TIES trimming 的区别——而且两者可以组合,mergekit 将其暴露为 dare_ties。
两个指向不同方向的权重向量之间的直线会穿过它们所在的球面,所以线性平均比任何一个父模型都短,每个激活都在缩小。SLERP 以恒定角速度走弧线,并单独插值幅度:
def slerp(a, b, t=0.5):
na, nb = a.norm(), b.norm()
ua, ub = a / na, b / nb
omega = torch.acos((ua * ub).sum().clamp(-1, 1))
s = torch.sin(omega)
direction = (torch.sin((1 - t) * omega) * ua + torch.sin(t * omega) * ub) / s
return direction * ((1 - t) * na + t * nb)
范数得以保留,在 t = 0 或 1 时你精确地恢复某个父模型。它精确地处理两个模型;对于三个或更多,使用 TIES 或 DARE。
没人会在 7B 规模上手写这些。mergekit 接收一个 YAML 配方,流式传输张量因此可以在内存适中的 CPU 上运行,然后写出一个可以在任何地方服务的正常 checkpoint:
base_model: mistralai/Mistral-7B-v0.1
merge_method: ties # linear | task_arithmetic | ties | dare_ties | slerp
dtype: bfloat16
models:
- model: org/mistral-sentiment
parameters: {weight: 0.5, density: 0.5} # weight = lambda, density = top-k
- model: org/mistral-spam
parameters: {weight: 0.5, density: 0.5}
合并本身需要几分钟。真正的工作是扫描——lambda、density、t 都是超参数,而且合并无法创造出父模型都没有的技能。生成候选项很便宜,这让评估它们成为实际的工作。一个合并是一个假设;eval 集才是答案。
这个演示真的在你的浏览器里通过梯度下降训练了一个基座和两个微调模型,然后实时合并它们——拖动滑块,观察两项任务在 held-out 准确率上的移动:https://dev48v.infy.uk/ai/days/day59-model-merging.html