[Paper Review] 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 Applications40 references60 citations

TL;DR

The paper analyzes loss geometry in massively multilingual models and introduces Gradient Vaccine (GradVac), a gradient-surgery method that adaptively aligns task gradients based on language proximity to improve multi-task optimization.

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.

Motivation & Objective

• Understand how gradient geometry correlates with language proximity during multilingual training.
• Evaluate whether gradient trajectory similarity predicts cross-lingual transfer and overall model quality.
• Identify limitations of existing gradient-based MTL methods in unbalanced, noisy multilingual data.
• Propose and validate GradVac as an adaptive gradient similarity objective to improve optimization.

Proposed method

• Run large-scale multilingual NMT experiments with 25 languages (50 language pairs) in En↔Any directions and Any↔En directions using Transformer-Big (375M params).
• Compute pairwise cosine similarities of encoder/decoder gradients across language pairs at checkpoints to characterize loss geometry.
• Analyze gradient similarity evolution across layers and training steps to relate geometry to language proximity and performance.
• Critique PCGrad's zero-similarity objective and derive GradVac to adapt gradient updates to desired similarity φ^T_ij.
• Introduce GradVac with a gradient update in the span of gi and gj to match a target gradient similarity φ^T_ij, using EMA to adapt φ^t_ijk across time/layers/tasks.
• Evaluate GradVac versus monolithic training and baseline MTL methods (GradNorm, MGDA, PCGrad) on WMT translation tasks and XTREME NER benchmarks.

Experimental results

Research questions

• RQ1 Do typologically similar languages exhibit more similar loss geometries during multilingual training?
• RQ2 Do higher gradient similarities between tasks correlate with better multilingual model quality and transfer?
• RQ3 Can we deliberately encourage geometrically aligned updates to improve optimization in unbalanced, noisy massively multilingual data?

Key findings

• Gradient similarities cluster by language family, reflecting linguistic proximity.
• Higher gradient similarity between language pairs correlates with better BLEU performance and cross-lingual transfer, and with greater overall model quality.
• Gradient similarities evolve across layers and training time, with encoder-decoder dynamics and resource-level dependencies (e.g., language vs English token spaces) influencing similarity.
• PCGrad often underperforms in positive-similarity regimes; GradVac generalizes by leveraging adaptive target similarities across task pairs, layers, and training steps.
• GradVac (all_layer variant) yields consistent gains over monolithic training and other MTL baselines on large-scale multilingual NMT; improvements seen in both Any→En and En→Any directions.