[论文解读] Gradient Vaccine: Investigating and Improving Multi-task Optimization in Massively Multilingual Models

Zi-Rui Wang, Yulia Tsvetkov+2 more|arXiv (Cornell University)|Oct 12, 2020

Multimodal Machine Learning Applications参考文献 40被引用 60

一句话总结

论文分析大规模多语言模型中的损失几何，并提出梯度疫苗 GradVac（Gradient Vaccine），一种基于语言邻近性自适应对齐任务梯度的梯度手术方法，以改善多任务优化。

ABSTRACT

Massively multilingual models subsuming tens or even hundreds of languages pose great challenges to multi-task optimization. While it is a common practice to apply a language-agnostic procedure optimizing a joint multilingual task objective, how to properly characterize and take advantage of its underlying problem structure for improving optimization efficiency remains under-explored. In this paper, we attempt to peek into the black-box of multilingual optimization through the lens of loss function geometry. We find that gradient similarity measured along the optimization trajectory is an important signal, which correlates well with not only language proximity but also the overall model performance. Such observation helps us to identify a critical limitation of existing gradient-based multi-task learning methods, and thus we derive a simple and scalable optimization procedure, named Gradient Vaccine, which encourages more geometrically aligned parameter updates for close tasks. Empirically, our method obtains significant model performance gains on multilingual machine translation and XTREME benchmark tasks for multilingual language models. Our work reveals the importance of properly measuring and utilizing language proximity in multilingual optimization, and has broader implications for multi-task learning beyond multilingual modeling.

研究动机与目标

• 理解在多语言训练中梯度几何如何与语言邻近性相关
• 评估梯度轨迹相似性是否能够预测跨语言迁移与整体模型质量
• 识别现有基于梯度的多任务学习方法在不平衡、嘈杂的多语言数据中的局限性
• 提出并验证 GradVac 作为自适应梯度相似性目标函数以改进优化

提出的方法

• 在 Transformer-Big（375M 参数）上开展 25 种语言（50 对语言对）的 En↔Any 方向和 Any↔En 方向的大规模多语言神经机器翻译实验
• 在检查点计算跨语言对的编码器/解码器梯度的成对余弦相似度以表征损失几何
• 分析梯度相似性在层与训练步骤上的演化，以将几何与语言邻近性及性能联系起来
• 批评 PCGrad 的零相似度目标并推导 GradVac 以将梯度更新自适应为期望的相似性 φ^T_ij
• 引入 GradVac，在 gi 与 gj 的范围内进行梯度更新以匹配目标梯度相似性 φ^T_ij，使用 EMA 在时间/层/任务上自适应 φ^t_ijk
• 在 WMT 翻译任务和 XTREME 命名实体识别基准上比较 GradVac 与统一训练及基线多任务学习方法（GradNorm、MGDA、PCGrad）

实验结果

研究问题

• RQ1 在类型学上相似的语言在多语言训练中是否表现出更相似的损失几何？
• RQ2 任务之间的梯度相似性是否与更好的多语言模型质量和迁移相关？
• RQ3 我们是否可以有意地促进几何对齐的更新以在不平衡、嘈杂的极大规模多语言数据中改善优化？

主要发现

模型	en-fr	en-cs	en-hi	en-tr	平均	fr-en	cs-en	hi-en	tr-en	平均
单一化训练	41.80	24.76	5.77	9.77	20.53	36.38	29.17	8.68	13.87	22.03
(2) 多语言模型	37.24	20.22	13.69	18.77	22.48	34.29	27.66	18.48	22.01	25.61
(3) GradNorm (Chen 等, 2018b)	37.02	18.78	11.57	15.44	20.70	34.58	27.85	18.03	22.37	25.71
(4) MGDA (Sener & Koltun, 2018)	38.22	17.54	12.02	13.69	20.37	35.05	26.87	18.28	22.41	25.65
(5) PCGrad (Yu 等, 2020)	37.72	20.88	13.77	18.23	22.65	34.37	27.82	18.78	22.20	25.79
(6) PCGrad 全层	38.01	21.04	13.95	18.46	22.87	34.57	27.84	18.84	22.48	25.93
(7) GradVac 固定目标$obj	38.41	21.12	13.75	18.68	22.99	34.55	27.97	18.72	22.14	25.85
(8) GradVac 全模型	38.76	21.32	14.22	18.89	23.30	34.84	28.01	18.85	22.24	25.99
(9) GradVac 全层	39.27	21.67	14.88	19.73	23.89	35.28	28.42	19.07	22.58	26.34

• 梯度相似性按语言家族聚类，反映语言距离关系
• 语言对之间的更高梯度相似性与更好的 BLEU 表现和跨语言迁移、以及更高的整体模型质量相关
• 梯度相似性在层级和训练时间上演化，编码器-解码器动态及资源层面依赖（如语言与英语令牌空间）影响相似性
• PCGrad 在正相似性区域常表现不佳；GradVac 通过在任务对、层和训练步骤上利用自适应目标相似性实现泛化
• GradVac（全层变体）在大规模多语言神经机器翻译中相对于统一训练与其他多语言学习基线显示出持续提升；在 Any→En 与 En→Any 方向均有改进