今天介绍一篇名为 EMO 的工作,它基于最优传输思想提出了新的改进损失函数,声称能大幅提高 LLM 的微调效果。
论文标题:
EMO: Earth Mover Distance Optimization for Auto-Regressive Language Modeling
论文地址:
https://arxiv.org/abs/2310.04691
概率散度
假设
p
i
p_i
pi 是模型预测的第
i
i
i个类别的概率,
i
=
1
,
2
,
.
.
.
,
n
i=1,2,...,n
i=1,2,...,n,
t
t
t则是目标类别,那么交叉熵损失为
L
=
−
l
o
g
p
t
\mathcal{L}=-logp_t
L=−logpt
如果将标签
t
t
t 用one hot形式的分布
τ
\tau
τ 表示出来(即
τ
t
=
1
,
τ
t
=
0
∣
i
≠
t
,
i
∈
[
1
,
n
]
\tau_t = 1,\tau_t = 0|i\neq t,i\in[1,n]
τt=1,τt=0∣i=t,i∈[1,n]) ,那么它可以重写成
L
=
−
∑
i
τ
i
l
o
g
p
i
\mathcal{L}=-\sum_{i}\tau_ilogp_i
L=−i∑τilogpi:
这个形式同时适用于非 one hot 的标签
τ
\tau
τ(即软标签),它等价于优化
τ
,
p
\tau,p
τ,p的 KL 散度:
K
L
(
τ
∣
∣
p
)
=
∑
i
τ
i
l
o
g
τ
i
p
i
=
∑
i
τ
i
l
o
g
τ
i
−
∑
i
τ
i
l
o
g
p
i
KL(\tau||p)=\sum_{i}\tau_ilog\frac{\tau_i}{p_i}=\sum_{i}\tau_ilog{\tau_i}-\sum_{i}\tau_ilog{p_i}
KL(τ∣∣p)=i∑τilogpiτi=i∑τilogτi−i∑τilogpi
当
τ
\tau
τ 给定时,最右端第一项就是一个常数,所以它跟交叉熵目标是等价的。
这个结果表明,我们在做 MLE,或者说以交叉熵为损失时,实则就是在最小化目标分布和预测分布的 KL 散度。由于 KL 散度的一般推广是 f 散度:
D
f
(
P
∣
∣
Q
)
=
∫
q
(
x
)
f
(
p
(
x
)
q
(
x
)
)
d
x
\mathcal{D}_f(P||Q)=\int q(x)f(\frac{p(x)}{q(x)})dx
Df(P∣∣Q)=∫q(x)f(q(x)p(x))dx
不过,每种 f 散度或多或少有些问题,要说概率分布之间的理想度量,当属基于最优传输思想的“推土机距离(Earth Mover’s Distance,EMD)
简单来说,推土机距离定义为两个分布之间的最优传输成本:
C
[
p
,
τ
]
=
inf
γ
∈
Π
[
p
,
τ
]
∑
i
,
j
γ
i
,
j
c
i
,
j
\mathcal{C}[p,\tau]=\inf_{\gamma \in \Pi[p,\tau]}\sum_{i,j}\gamma_{i,j}c_{i,j}
C[p,τ]=γ∈Π[p,τ]infi,j∑γi,jci,j
这里的
γ
∈
[
p
,
τ
]
\gamma \in[p,\tau]
γ∈[p,τ]说的是
γ
\gamma
γ 是任意以
p
,
τ
p,\tau
p,τ 为边缘分布的联合分布,
c
i
,
j
c_{i,j}
ci,j是实现给定的成本函数,代表“从
i
i
i 搬运到
j
j
j 的成本”, inf是下确界,意思就是说将最低的运输成本作为
p
,
τ
p,\tau
p,τ 之间的差异度量。本文期望如果将分类的损失函数换成两个分布的 W 距离,也能收敛到更好的结果。
当
τ
\tau
τ 是 one hot 分布时,目标分布就是一个点
t
t
t,那么就无所谓最不最优了,传输方案就只有一个,即把
p
p
p 的所有东西都搬到同一个点 ,所以此时就有
C
[
p
,
τ
]
=
∑
i
p
i
c
i
,
t
\mathcal{C}[p,\tau]=\sum_{i}p_{i}c_{i,t}
C[p,τ]=i∑pici,t
如果
τ
\tau
τ 是一般的软标签分布,那么
C
[
p
,
τ
]
\mathcal{C}[p,\tau]
C[p,τ] 的计算是一个线性规划问题,求解起来比较复杂,由于
p
i
τ
j
p_i\tau_j
piτj 所定义的分布也属于
Π
[
p
,
τ
]
\Pi[p,\tau]
Π[p,τ],那么我们有
C
[
p
,
τ
]
=
inf
γ
∈
Π
[
p
,
τ
]
∑
i
,
j
γ
i
,
j
c
i
,
j
≤
∑
i
,
j
p
i
τ
j
c
i
,
j
\mathcal{C}[p,\tau]=\inf_{\gamma \in \Pi[p,\tau]}\sum_{i,j}\gamma_{i,j}c_{i,j} \leq\sum_{i,j}p_{i}\tau_{j}c_{i,j}
C[p,τ]=γ∈Π[p,τ]infi,j∑γi,jci,j≤i,j∑piτjci,j
这是一个容易计算的上界,也可以作为优化目标。
成本函数
现在回到原论文所关心的场景——LLM 的微调,包括二次预训练和微调到下游任务等。正如本文开头所述,LLM 的训练可以理解为逐 token 的分类任务(类别即所有 token),每个标签是 one hot 的,所以适用于上面式子。
上面式子还差成本函数
c
i
,
t
c_{i,t}
ci,t 还没定下来。如果简单地认为只要
i
≠
t
i\neq t
i=t,那么成本都是 1,即
c
i
,
t
=
1
−
δ
i
,
t
c_{i,t}=1-\delta_{i,t}
ci,t=1−δi,t,那么
C
[
p
,
τ
]
=
∑
i
p
i
c
i
,
t
=
∑
i
p
i
−
p
i
δ
i
,
t
=
1
−
p
t
\mathcal{C}[p,\tau]=\sum_{i}p_{i}c_{i,t}=\sum_{i}p_i-p_{i}\delta_{i,t}=1-p_t
C[p,τ]=i∑pici,t=i∑pi−piδi,t=1−pt
这其实就是在最大化准确率的光滑近似,但直觉上,所有
i
≠
t
i\neq t
i=t 都给予同样程度的惩罚似乎过于简单了,理想情况下应该根据相似度来给每个不同的
i
i
i 设计不同的成本,即相似度越大,传输成本越低,那么我们可以将传输成本设计为
c
i
,
t
=
1
−
c
o
s
(
e
i
,
e
t
)
=
1
−
⟨
e
i
∣
∣
e
i
∣
∣
,
e
t
∣
∣
e
t
∣
∣
⟩
c_{i,t} = 1-cos(e_i,e_t)=1-\langle \frac{e_i}{||e_i||},\frac{e_t}{||e_t||} \rangle
ci,t=1−cos(ei,et)=1−⟨∣∣ei∣∣ei,∣∣et∣∣et⟩
这里的
e
i
,
e
t
e_i,e_t
ei,et 是事先获取到 Token Embedding,原论文是将预训练模型的 LM Head 作为 Token Embedding 的,并且根据最优传输的定义成本函数是要实现给定的,因此计算相似度的 Token Embedding 要在训练过程中固定不变。
有了成本函数后,我们就可以计算
C
[
p
,
τ
]
=
∑
i
p
i
c
i
,
t
=
∑
i
(
p
i
−
p
i
⟨
e
i
∣
∣
e
i
∣
∣
,
e
t
∣
∣
e
t
∣
∣
⟩
)
=
1
−
⟨
∑
i
p
i
e
i
∣
∣
e
i
∣
∣
,
e
t
∣
∣
e
t
∣
∣
⟩
\mathcal{C}[p,\tau]=\sum_{i}p_{i}c_{i,t}=\sum_{i}(p_i-p_i\langle \frac{e_i}{||e_i||},\frac{e_t}{||e_t||} \rangle)=1-\langle \sum_{i}p_i\frac{e_i}{||e_i||},\frac{e_t}{||e_t||} \rangle
C[p,τ]=i∑pici,t=i∑(pi−pi⟨∣∣ei∣∣ei,∣∣et∣∣et⟩)=1−⟨i∑pi∣∣ei∣∣ei,∣∣et∣∣et⟩
这就是 EMO(Earth Mover Distance Optimization)最终的训练损失。由于 embedding_size 通常远小于 vocab_size,所以先算
∑
i
p
i
e
i
∣
∣
e
i
∣
∣
\sum_{i}p_i\frac{e_i}{||e_i||}
∑ipi∣∣ei∣∣ei 能明显降低计算量。
实验效果
一起看看原论文的实验。不得不说,原论文的实验结果还是比较惊艳的。
首先,是小模型上的继续预训练实验,相比交叉熵(MLE)的提升最多的有 10 个点,并且是全面 SOTA:
值得一提的是,这里的评价指标是 MAUVE,越大越好,它提出自《MAUVE: Measuring the Gap Between Neural Text and Human Text using Divergence Frontiers》,是跟人工评价最相关的自动评测指标之一。
那么 EMO 更好是不是单纯因为评价指标选得好?并不是,让人意外的是,EMO 训练的模型,甚至 PPL 都更好(PPL 跟 MLE 更相关):
然后是将 LLAMA-7B/13B 微调到下游任务做 Few Shot 的效果,同样很出色:
最后对比了不同模型规模和数据规模的效果,显示出 EMO 在不同模型和数据规模上都有不错的表现:
个人思考
总的来说,原论文的“成绩单”还是非常漂亮的,值得一试。唯一的疑虑可能是原论文的实验数据量其实都不算大,不清楚进一步增大数据量后是否会缩小 EMO 和 MLE 的差距。
就笔者看来,EMO 之所以能取得更好的结果,是因为它通过 Embedding 算相似度,来为“近义词”分配了更合理的损失,从而使得模型的学习更加合理。因为虽然形式上 LLM 也是分类任务,但它并不是一个简单的对与错问题,并不是说下一个预测的 token 跟标签 token 不一致,句子就不合理了,因此引入语义上的相似度来设计损失对 LLM 的训练是有帮助的。可以进一步猜测的是,vocab_size 越大、token 颗粒度越大的情况下,EMO 的效果应该越好,因为 vocab_size 大了“近义词”就可能越多。
当然,引入语义相似度也导致了 EMO 不适用于从零训练,因为它需要一个训练好的 LM Head 作为 Token Embedding。当然,一个可能的解决方案是考虑用其他方式,比如经典的 Word2Vec 来事先训练好 Token Embedding,但这可能会有一个风险,即经典方式训练的 Token Embedding 是否会降低 LLM 能力的天花板(毕竟存在不一致性)。
此外,即便 Token Embedding 没问题,从零预训练时单纯用 EMO 可能还存在收敛过慢的问题,从损失函数视角来看:
首先寻找评测指标的一个光滑近似,最好能表达成每个样本的期望形式,然后将错误方向的误差逐渐拉到无穷大(保证模型能更关注错误样本),但同时在正确方向保证与原始形式是一阶近似。也就是说,为了保证(从零训练的)收敛速度,错误方向的损失最好能拉到无穷大,而 EMO 显然不满足这一点,因此将 EMO 用于从零训练的时候,大概率是 EMO 与 MLE 的某个加权组合,才能平衡收敛速度和最终效果。