通过最优转移进行词表学习：VOLT

论文：[2012.15671] Vocabulary Learning via Optimal Transport for Machine Translation

Code：Jingjing-NLP/VOLT: Code for paper “Vocabulary Learning via Optimal Transport for Neural Machine Translation”

一句话概述：借鉴边际效用通过最优转移学习词表。

摘要：机器翻译词表影响性能，本文旨在弄清楚什么是好的词表，以及是否可以在不尝试训练的情况下找到最佳词表。为此，作者首先从信息论的视角提供对词表作用的另一种理解。受此启发，将词汇化的探索 —— 找到具有适当大小的最佳词表—— 作为最优转移问题。提出了 VOLT，一个简单有效的无须尝试训练的解决方案。实验结果表明，VOLT 在各种场景中优于广泛使用的词表。此外，与 BPE-search 相比，VOLT 将搜索时间从 384 GPU 小时减少到 30 GPU 小时。

背景

词表的构建是 NMT 任务和其他 NLP 任务的前提条件，最近子词方法（比如 BPE）比较流行，也取得不错的效果。这些方法的核心思想是选择高频的子词（或词片段）作为词表 Token。在信息论中，基于词频的方法是数据压缩（减熵）的简单形式，使得生成的语料库易于学习和预测。不过词表的大小并没有得到充分重视，有些工作表明词表大小也会影响下游任务表现，尤其是在低资源任务上。由于缺乏适当的关于词表尺寸的归纳偏向，通常需要尝试训练（遍历所有可能的尺寸）来搜索最佳尺寸，这需要很高的计算成本。方便起见，大多数已有研究都只采纳了广泛使用的配置，比如 30K-40K。

本文建议通过同时考虑熵和词汇量大小来探索自动词汇化，而无需昂贵的试验训练。这并不容易，主要因为：

  • 难以找到一个合适的目标函数。词表大时语料库熵下降，但更稀疏却不利于模型学习。
  • 即便有评估方法，依然很有挑战，因为优化问题是指数搜索空间。

本文提出 Vocabulary Learning approach via Optimal Transport，简称 VOLT，可以通过考虑语料库熵和词汇量大小在多项式时间内给出合适的词汇量。具体而言：

  • 首先借用了经济学中边际效用的概念，使用 MUV（词汇的边际效用）作为评估方法。形式上，MUV 被定义为熵对词汇量大小的负导数。
  • 然后将目标转向在可处理的时间复杂度中最大化 MUV，将离散优化目标重新表述为最优转移问题，可以通过线性规划在多项式时间内求解。
  • 最后从最佳转移矩阵中生成词表。

直觉上，词表化的过程可以认为是找到从字符分布到词表分布的最佳转移矩阵。

相关工作

  • 词级别
  • Byte 级别
  • 子词方法
  • BPE
  • BPE 变种：BPE-dropout，SentencePiece

MUV

给定两个词表 v(k) 和 v(k+m)，k 和 k+m 表示 token 数量，Hv 表示词表 v 上的语料库熵，等于所有 token 熵的和

Mv(k+m)=−(Hv(k+m)−Hv(k))m,(1)Hv=−1lv∑i∈vP(i)log⁡P(i),(2)\mathcal{M}_{v(k+m)}=\frac{-\left(\mathcal{H}_{v(k+m)}-\mathcal{H}_{v(k)}\right)}{m}, \quad (1)\\ \mathcal{H}_v = - \frac{1}{l_v} \sum_{i \in v} P(i) \log P(i), \quad (2)

P(i) 是 token i 在训练语料中的相对频率，l 是词表 v 中 token 的平均长度。

给定 MUV，有两种选择得到最终的词表：

  • 基于搜索：与广泛使用的词汇化解决方案（比如 BPE）结合起来。简单有效但不是自给自足的方法，而且需要大量时间生成很多词表计算 MUV。
  • 基于学习：VOLT

通过最优转移最大化 MUV

概览

由于词表是离散的，搜索空间太大，根据上面（1）式优化不好处理。本文通过从固定大小的词表中搜索最佳词表来简化原始的离散优化问题。给定 S，每个时间步 t 表示一组数量为 S[t] 的词，对任意的词表，MUV 分数可以基于上一步的词表：

\begin{array}{\arg \max } \underset{v(t-1) \in \mathbb{V}_{S[t-1]}, v(t) \in \mathbb{V}_{S[t]}}{\arg \max } \mathcal{M}_{v(t)} =\\ \underset{v(t-1) \in \mathbb{V}_{S[t-1]}, v(t) \in \mathbb{V}_{S[t]}}{\arg \max }-\frac{1}{i}\left[\mathcal{H}_{v(t)}-\mathcal{H}_{v(t-1)}\right] \end{array}

V_S[t-1] 和 V_S[t] 是两组词表大小上界分别为 S[t-1] 和 S[t] 的词表，由于是指数搜索空间，文章建议优化上界：

arg⁡max⁡t1i[max⁡v(t)∈VS[t]Hv(t)−max⁡v(t−1)∈VS[t−1]Hv(t−1)](3)\underset{t}{\arg \max } \frac{1}{i}\left[\max _{v(t) \in \mathbb{V}_{S[t]}} \mathcal{H}_{v(t)}-\max _{v(t-1) \in \mathbb{V}_{S[t-1]}} \mathcal{H}_{v(t-1)}\right] \quad(3)

i 是 t-1 词表和 t 词表 size 的差值。

这里请教过作者，式（3）是上界，具体推导过程如下：

max⁡at,bt−(at−bt)  ⟺  min⁡at,bt(at−bt)  ⟺  min⁡atat−max⁡btbt≤max⁡atat−max⁡btbt\max_{at, bt} - (at - bt) \iff \\ \min_{at,bt} (at-bt) \iff \\ \min_{at} at - \max_{bt} bt \le \max_{at} at - \max_{bt} bt

所以，最大化 MUV 分数可以表达为最大化式（3）。基于该等式，整个解决方案可以分成两步：

  • 在每个时间步 t 搜索具有最大熵的最优词表
  • 枚举所有时间步，输出满足 Eq3 的时间步对应的词表

第一步其实是从 V_S[t] 中搜索具有最大熵的词表 v(t)：

arg⁡max⁡v(t)∈VS[t]−1lv(t)∑i∈v(t)P(i)log⁡P(i)(4)\arg \max_{v(t) \in \mathbb{V}_{S[t]}} - \frac{1}{l_{v(t)}} \sum_{i \in v(t)} P(i) \log P(i) \quad (4)

lv 是 v(t) 中 token 的平均长度，P(i) 是 token i 的概率。但是由于词汇量很大，这个问题通常很难解决。 因此，论文在离散最优转移公式中提出了一个松弛，然后可以通过 Sinkhorn 算法有效地解决。直观地，可以将词汇构建想象成一个转移过程，将字符传输到大小为 S [t] 的 Token 候选中。

每个转移矩阵都可以通过收集带有字符的 Token 来构建词汇表。不同的转移矩阵带来不同的转移成本，最优转移目标是找到一个转移矩阵来最小化转移成本，即我们设置中的负熵。

通过最优转移词表化

给定一组词表 V_S[t]，需要找到具有最大熵的词表，Eq4 可以等价地写为：

min⁡v∈VS[t]1lv∑i∈vP(i)log⁡P(i), s.t. P(i)=Token⁡(i)∑i∈vToken⁡(i),lv=∑i∈vlen⁡(i)∣v∣.\begin{array}{c} \min _{v \in \mathbb{V}_{\boldsymbol{S}[t]}} \frac{1}{l_{v}} \sum_{i \in v} P(i) \log P(i), \\ \text { s.t. } P(i)=\frac{\operatorname{Token}(i)}{\sum_{i \in v} \operatorname{Token}(i)}, l_{v}=\frac{\sum_{i \in v} \operatorname{len}(i)}{|v|} . \end{array}

Token(i) 是 token i 的频率，len(i) 表示 token i 的长度。

目标近似

为了得到一个易处理的熵下界，给出上述目标函数的一个易处理的上界就足够了。论文采用合并规则来分割原始文本，类似于 BPE，如果合并后的 Token 在词汇表中，两个连续的 Token 将合并为一个。令 T ∈ V_S[t] 是包含 top S[t] 个高频 Token 的词表，C 是字符（char）的集合：

min⁡v∈VS[t]1lv∑i∈vP(i)log⁡P(i)≤1lT∑i∈TP(i)log⁡P(i)\min _{v \in \mathbb{V}_{S[t]}} \frac{1}{l_{v}} \sum_{i \in v} P(i) \log P(i) \leq \frac{1}{l_{\mathbb{T}}} \sum_{i \in \mathbb{T}} P(i) \log P(i)

注意，左边式子是 Σ 的最小值。这里，从目标函数的上界开始，从 T 中搜索精确的 Token 集合，这样就把搜索空间降至 T 的子集。令 P(i,j) 是要学习的 Token 和 Char 的联合概率分布：

∑i∈TP(i)log⁡P(i)=∑i∈T∑j∈CP(i,j)log⁡P(i)=∑i∈T∑j∈CP(i,j)log⁡P(i,j)⏟L1+∑i∈T∑j∈CP(i,j)(−log⁡P(j∣i))⏟L2.(6)\begin{aligned} \sum_{i \in \mathbb{T}} P(i) \log P(i) &=\sum_{i \in \mathbb{T}} \sum_{j \in \mathbb{C}} P(i, j) \log P(i) \\ &=\underbrace{\sum_{i \in \mathbb{T}} \sum_{j \in \mathbb{C}} P(i, j) \log P(i, j)}_{\mathcal{L}_{1}} \\ &+\underbrace{\sum_{i \in \mathbb{T}} \sum_{j \in \mathbb{C}} P(i, j)(-\log P(j \mid i))}_{\mathcal{L}_{2}} . \end{aligned} \quad (6)

L1 是联合概率分布 P(i,j) 的负熵，可以记为 -H§。令 D 是 |C| × |T| 矩阵，第 (i,j) 个实体是 -logP(j|i)，令 P 是联合概率矩阵，L2 可以写为：

L2=⟨P,D⟩=∑i∑jP(i,j)D(i,j)\mathcal{L}_{2}=\langle\boldsymbol{P}, \boldsymbol{D}\rangle=\sum_{i} \sum_{j} \boldsymbol{P}(i, j) \boldsymbol{D}(i, j)

Eq6 可以重新表述为如下目标函数，其形式与最优转移中的目标函数相同：

min⁡P∈Rm×n⟨P,D⟩−γH(P)\min _{\boldsymbol{P} \in \mathbb{R}^{m} \times n}\langle\boldsymbol{P}, \boldsymbol{D}\rangle-\gamma H(\boldsymbol{P})

从最优转移的角度来看，P 可以看作是转移矩阵，D 可以看作是距离矩阵。直观地说，最佳转移是用 ⟨P ，D⟩ 定义的最小工作找到从字符分布到目标 Token 分布的最佳转移质量。

OT的设置

为了保证方案有效性，论文添加了以下约束：

  • 阻止从 Char j 到 Token i 的无效的转移，如果 i 不包括 j，则将距离设置为 +∞；否则使用 1/len(i) 评估 P(j|i)，其中 len(i) 是 Token i 的长度

    D(i,j)={−log⁡P(j∣i)=+∞, if j∉i−log⁡P(j∣i)=−log⁡1len⁡(i), otherwise \boldsymbol{D}(i, j)=\left\{\begin{array}{ll} -\log P(j \mid i)=+\infty, & \text { if } j \notin i \\ -\log P(j \mid i)=-\log \frac{1}{\operatorname{len}(i)}, & \text { otherwise } \end{array}\right.

  • 字符（Char）的数量是固定的，设置转移矩阵每一行的和为 Char j 的概率；每个 Token 需要的字符（Char）的上限是固定的，将转移矩阵中每列的和设置为 Token i 的概率。

    ∑iP(i,j)=P(j)∣∑jP(i,j)−P(i)∣≤ϵ\sum_i P(i,j) = P(j) \\ | \sum_j P(i,j) - P(i) | \le \epsilon

给定转移矩阵 P 和距离矩阵 D，最终目标函数为：

arg⁡min⁡P∈R∣C∣×∣∣−H(P)+⟨P,D⟩ s.t. ∑iP(i,j)=P(j),∣∑jP(i,j)−P(i)∣≤ϵ\begin{array}{l} \quad \underset{\boldsymbol{P} \in \mathbb{R}^{|\mathbb{C}| \times|\mathbb{}|}}{\arg \min }-H(\boldsymbol{P})+\langle\boldsymbol{P}, \boldsymbol{D}\rangle \\ \text { s.t. } \quad \sum_{i} \boldsymbol{P}(i, j)=P(j), \quad\left|\sum_{j} \boldsymbol{P}(i, j)-P(i)\right| \leq \epsilon \end{array}

具体可以参考下图：

严格来说，这是一个非平衡熵正则化最优转移问题。尽管如此，仍然可以使用广义 Sinkhorn 算法来有效地找到目标词表。算法细节如下所示。在每个时间步长 t，可以根据转移矩阵 P 生成与熵分数相关的新词表。最后，收集这些与熵分数相关的词表，并输出满足 Eq3 的词表。

实现

  • 对所有 Token 按频率从大到小排序，简单起见，采用 BPE 生成的 Token 作为候选 Token。实际上任何分割算法都可以初始化 Token 候选，不同的初始化方法结果接近。
  • 所有的 Token 候选以及它们的概率用来初始化算法中的 L。
  • 递增整数序列 S 的大小是超参数，双语翻译设置为（1K,…,10K），多语翻译设置为（40K,…,160K）
  • 每一个时间步 t，可以基于转移矩阵使用最大熵得到词表。由于限制放宽，处理非法转移情况在所难免。因此，删除了那些频率低于 0.001 的 Token。
  • 最后，枚举出所有时间步，选择出满足 Eq3 的词表作为最终词表。

生成词表后，VOLT 和 BPE 类似，使用贪婪策略对文本进行编码。即首先将句子切分成字符级，然后合并连续的两个 Token（如果合并后在词表中），直到没有 Token 可以被合并为止。

实验

三个数据集：WMT-14 英德翻译，TED 双语翻译和 TED 多语翻译。

主要结果

  1. VOLT 搜索的词表效果优于之前广泛使用的方法，可以找到具有更高 BLEU 和更小尺寸的性能良好的词汇表。

  1. VOLT 搜索的词汇表与低资源数据集上启发式搜索的词汇表相当。

  1. VOLT 在多语言翻译上运行良好。

  1. VOLT 是一种绿色词汇解决方案，与需要数百 GPU 小时的 BPE-Search 相比，可以在单个 CPU 上 0.5 小时内找到具有竞争力的词汇。

讨论

  1. 带有 VOLT 生成的词汇表的简单基线达到 SOTA 结果，换句话说，简单的基线可以通过定义明确的词汇表获得良好的结果。（Table5）
  2. 胜过 SentencePiece 和 WordPiece。（Table5）
  3. 适用于各种架构。

  1. 可以在训练期间带来轻微的加速。
    • 2 Tesla-V 100-GPUs + 1 Gold-6130-CPU
    • VOLT-searched（11.6k tokens）：133 句子/秒 VS BPE-30K（33.6k tokens）：101句子/秒
    • 加速主要来自更大的批量和减少的 Embedding 参数
  2. VOLT 词汇表和 BPE 词汇表高度重叠。从经验的角度来看，具有 VOLT 大小的 BPE 也是一个不错的选择。

结论

  • 提出了无须尝试训练的新的词汇搜索方法。
  • 框架从信息论开始，借用经济学中边际效用的概念，使用 MUV（词汇的边际效用）作为评估方法。
  • 将词汇化制定为一个两步离散优化目标，并将其表述为最优转移问题。

感想

在刚看到这篇 Paper 介绍时就被其中借鉴的边际效用吸引，一方面是因为我之前经济学专业，比较敏感，另一方面是之前做过一些将 “法则” 或 “规则” 与机器学习融合的尝试，比如最省力法则、齐夫定律等。因此，虽然目前工作不是做这一块的，依然迫不及待地第一时间去拜读了。读完后感觉确实不错，不只是指结果，更多地是「如何将某一个规则从想法应用到实际，并产生还不错结果」这一过程。其中碰到几个问题，给一作许晶晶同学发了邮件请教，人非常耐心、友好，很愉快的交流。希望她们未来能够带来更多出色的作品。