EMO: Earth Mover Distance Optimization for Auto-Regressive Language Modeling

今天介绍一篇名为 EMO 的工作，它基于最优传输思想提出了新的改进损失函数，声称能大幅提高 LLM 的微调效果。
论文标题：
EMO: Earth Mover Distance Optimization for Auto-Regressive Language Modeling
论文地址：
https://arxiv.org/abs/2310.04691

概率散度

假设 $p_i$ 是模型预测的第 $i$个类别的概率，$i=1,2,...,n$，$t$则是目标类别，那么交叉熵损失为
$$\mathcal{L}=-log{p}_t$$
如果将标签 $t$ 用one hot形式的分布 $\tau$ 表示出来（即 ${\tau }_t=1,{\tau }_t=0|i\ne t，i\in [1,n]$） ，那么它可以重写成
$$\mathcal{L}=-\sum _i{\tau }_{i}log{p}_i$$：
这个形式同时适用于非 one hot 的标签 $\tau$（即软标签），它等价于优化 $\tau ,p$的 KL 散度：
$$KL(\tau ||p)=\sum _i{\tau }_{i}log\frac{{\tau }_i}{p_i}=\sum _i{\tau }_{i}log{\tau }_i-\sum _i{\tau }_{i}log{p}_i$$
当 $\tau$ 给定时，最右端第一项就是一个常数，所以它跟交叉熵目标是等价的。
这个结果表明，我们在做 MLE，或者说以交叉熵为损失时，实则就是在最小化目标分布和预测分布的 KL 散度。由于 KL 散度的一般推广是 f 散度：
$${\mathcal{D}}_f(P||Q)=\int q(x)f(\frac{p(x)}{q(x)})dx$$

不过，每种 f 散度或多或少有些问题，要说概率分布之间的理想度量，当属基于最优传输思想的“推土机距离（Earth Mover’s Distance，EMD）
简单来说，推土机距离定义为两个分布之间的最优传输成本：
$$\mathcal{C}[p,\tau ]=\underset{\gamma \in \Pi [p,\tau ]}{\inf }\sum _{i,j}{\gamma }_{i,j}{c}_{i,j}$$
这里的$\gamma \in [p,\tau ]$说的是 $\gamma$ 是任意以 $p,\tau$ 为边缘分布的联合分布，$c_{i,j}$是实现给定的成本函数，代表“从 $i$ 搬运到 $j$ 的成本”， inf是下确界，意思就是说将最低的运输成本作为 $p,\tau$ 之间的差异度量。本文期望如果将分类的损失函数换成两个分布的 W 距离，也能收敛到更好的结果。
当 $\tau$ 是 one hot 分布时，目标分布就是一个点 $t$，那么就无所谓最不最优了，传输方案就只有一个，即把 $p$ 的所有东西都搬到同一个点 ，所以此时就有
$$\mathcal{C}[p,\tau ]=\sum _i{p}_i{c}_{i,t}$$
如果 $\tau$ 是一般的软标签分布，那么 $\mathcal{C}[p,\tau ]$ 的计算是一个线性规划问题，求解起来比较复杂，由于 $p_i{\tau }_j$ 所定义的分布也属于 $\Pi [p,\tau ]$，那么我们有
$$\mathcal{C}[p,\tau ]=\underset{\gamma \in \Pi [p,\tau ]}{\inf }\sum _{i,j}{\gamma }_{i,j}{c}_{i,j}\le \sum _{i,j}{p}_i{\tau }_j{c}_{i,j}$$
这是一个容易计算的上界，也可以作为优化目标。

成本函数

现在回到原论文所关心的场景——LLM 的微调，包括二次预训练和微调到下游任务等。正如本文开头所述，LLM 的训练可以理解为逐 token 的分类任务（类别即所有 token），每个标签是 one hot 的，所以适用于上面式子。
上面式子还差成本函数$c_{i,t}$ 还没定下来。如果简单地认为只要 $i\ne t$，那么成本都是 1，即$c_{i,t}=1-{\delta }_{i,t}$，那么
$$\mathcal{C}[p,\tau ]=\sum _i{p}_i{c}_{i,t}=\sum _i{p}_i-p_i{\delta }_{i,t}=1-p_t$$
这其实就是在最大化准确率的光滑近似，但直觉上，所有 $i\ne t$ 都给予同样程度的惩罚似乎过于简单了，理想情况下应该根据相似度来给每个不同的 $i$ 设计不同的成本，即相似度越大，传输成本越低，那么我们可以将传输成本设计为
$$c_{i,t}=1-cos(e_i,e_t)=1-⟨\frac{e_i}{||e_i||},\frac{e_t}{||e_t||}⟩$$
这里的 $e_i,e_t$ 是事先获取到 Token Embedding，原论文是将预训练模型的 LM Head 作为 Token Embedding 的，并且根据最优传输的定义成本函数是要实现给定的，因此计算相似度的 Token Embedding 要在训练过程中固定不变。
有了成本函数后，我们就可以计算
$$\mathcal{C}[p,\tau ]=\sum _i{p}_i{c}_{i,t}=\sum _i(p_i-p_i⟨\frac{e_i}{||e_i||},\frac{e_t}{||e_t||}⟩)=1-⟨\sum _i{p}_i\frac{e_i}{||e_i||},\frac{e_t}{||e_t||}⟩$$
这就是 EMO（Earth Mover Distance Optimization）最终的训练损失。由于 embedding_size 通常远小于 vocab_size，所以先算 ${\sum }_i{p}_i\frac{e_i}{||e_i||}$ 能明显降低计算量。

实验效果

一起看看原论文的实验。不得不说，原论文的实验结果还是比较惊艳的。

首先，是小模型上的继续预训练实验，相比交叉熵（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 的某个加权组合，才能平衡收敛速度和最终效果。