Source-linked AI summary
Gradient Vaccine: Investigating and Improving Multi-task Optimization in Massively Multilingual Models
Zirui Wang, Yulia Tsvetkov, Orhan Firat, Yuan Cao
TL;DR
Multilingual models are commonly optimized with language-agnostic objectives that overlook task relatedness and may cause gradient interference. This paper analyzes gradient geometry, proposes Gradient Vaccine to align updates according to task relatedness, and reports superior multilingual performance.
Problem
Multilingual optimization commonly uses a language-agnostic joint objective that overlooks task relatedness and leaves cross-lingual interactions under-explored.
Method
Gradient Vaccine analyzes gradient similarity and adaptively aligns task gradients toward similarity objectives determined by task relatedness.
Results
Gradient Vaccine significantly outperforms baselines, improving average BLEU and individual performance on all tasks over naive joint training.
Takeaways & Limitations
Gradient similarity reflects language proximity and model quality, supporting geometrically consistent updates aligned with task closeness for multilingual optimization.
Takeaways & Limitations
PCGrad implicitly assigns every task pair a zero gradient-similarity objective and is ineffective when gradients already have positive similarity.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Multilingual models are efficient but are optimized as monolithic multi-task systems that overlook task relatedness and obscure cross-lingual interactions. The paper studies loss geometry to understand these dynamics and develop more aligned updates for multilingual optimization.
- 1 INTRODUCTION: Multilingual models process tens or hundreds of languages in one model, reducing resources relative to maintaining separate models.They are also motivated by positive cross-lingual transfer.
- 1 INTRODUCTION: Existing multilingual training commonly uses one language-agnostic objective over concatenated multilingual data, ignoring task relatedness and risking negative interference.This monolithic procedure also leaves interactions among languages and cross-lingual transfer mechanisms poorly understood.
- 1 INTRODUCTION: The paper asks whether typologically similar languages have similar loss geometries, whether aligned gradients relate to lower interference and better quality, and whether alignment can be encouraged.The study frames these questions around gradient behavior during multilingual optimization.
- 1 INTRODUCTION: Experiments analyze gradient similarity along optimization trajectories in massively multilingual neural machine translation, relating it to language proximity, cross-lingual transfer, and joint performance.Each language pair is treated as a separate task.
- 1 INTRODUCTION: The proposed Gradient Vaccine uses task relatedness to set gradient-similarity objectives and adaptively align task gradients during training.It is designed for massively multilingual settings with noisy and unbalanced training data.
2 INVESTIGATING MULTI-TASK OPTIMIZATION IN MASSIVELY MULTILINGUAL MODELS
Control experiments examine how gradients interact across languages, layers, and training steps in massively multilingual translation. They find that gradient similarity tracks linguistic proximity and model quality, motivating systematic optimization improvements.
- 2 INVESTIGATING MULTI-TASK OPTIMIZATION IN MASSIVELY MULTILINGUAL MODELS: The study uses multilingual NMT models to investigate gradient interactions and their effects on model performance.The authors state that these optimization dynamics are not well understood despite prior work on data, architecture, and scale.
- 2.1 EXPERIMENTAL SETUP: The experiments jointly train 50 English-pivoted language pairs selected from over 8 billion sentence pairs across 25 languages and 10 language families.The corpus covers 102 languages overall and multiple data-size levels.
- 2.2 OBSERVATIONS: Gradient similarities cluster by linguistic-family proximity, with Western European similarities declining across West Slavic, South Slavic, and East Slavic languages.Figure 1 averages encoder-gradient cosine similarities across training checkpoints.
- 2.2 OBSERVATIONS: En→Any tasks have less similar gradients than Any→En tasks, while paired-task BLEU comparisons show that gradient similarity is associated with model performance.The experiments compare directions in a joint model and evaluate En→Fr after pairing it with different En→Any languages.
- 2.2 OBSERVATIONS: Gradient similarity changes differently across layers and training steps: En→Fr with En→Hi becomes less similar in one decoder layer but more similar in the corresponding encoder layer.En→Fr and En→Es remain more similar than En→Fr and En→Hi in the same layer.
- 2.3 IMPLICATIONS: These observations motivate improving multilingual multi-task optimization by accounting for loss geometry rather than treating training as an undifferentiated joint objective.The paper transitions from analyzing gradient behavior to designing a systematic optimization method.
3 PROPOSED METHOD
The method identifies a limitation in PCGrad’s fixed zero-similarity objective and proposes Gradient Vaccine, which adapts gradient similarity objectives across tasks, layers, and training steps.
- 3.1 GRADIENT SURGERY: PCGrad detects negative gradient cosine similarity and projects conflicting gradients to deconflict task updates.It processes task gradients pairwise, replacing a gradient after projection and repeating across tasks in random order.
- 3.1 GRADIENT SURGERY: PCGrad implicitly targets zero gradient similarity, assuming all tasks should have identical gradient interactions.After projection, the resulting gradient cosine similarity is always zero.
- 3.1 GRADIENT SURGERY: PCGrad is ineffective for positive gradient similarities, which occur frequently among language-pair tasks.Its negative-similarity precondition leaves it inactive for positive similarities and sparse during training.
- 3.2 GRADIENT VACCINE: Gradient Vaccine sets adaptive similarity objectives and alters gradient magnitude and direction within the space spanned by two task gradients.The target can be arbitrary within [-1, 1], allowing the method to accommodate differing inter-task relationships.
- 3.2 GRADIENT VACCINE: Gradient Vaccine uses an exponential moving average to incorporate task-pair, parameter-group, and training-step variation in gradient interactions.The method uses computed gradient similarities and a hyper-parameter to update its moving-average objective.
- 3.2 GRADIENT VACCINE: Gradient surgery is a special case of Gradient Vaccine with a target similarity of zero, whereas Gradient Vaccine updates gradients proactively for both positive and negative similarities.Figure 5 contrasts PCGrad’s inactivity for positive similarity with GradVac’s activity in both cases.
4 EXPERIMENTS
The experiments compare scalable multi-task optimization methods with monolithic multilingual training across multilingual translation and language-model benchmarks. GradVac improves multilingual performance, including average and individual WMT task scores, while ablations support geometry-aware, task-related objectives.
- Experimental Setup: The study compares monolithic multilingual training with GradNorm, MGDA, PCGrad, and the proposed method on multilingual NMT and language-model tasks.Language-specific gradients are computed per batch, with sampling temperature fixed at T=5 unless stated otherwise.
- WMT Machine Translation: The naive multilingual baseline helps low-resource languages but underperforms bilingual models on high-resource languages because of negative interference and constrained capacity.The result is reported for the WMT experiments.
- WMT Machine Translation: GradNorm and MGDA obtain marginal or no improvement, indicating that reweighting task weights without considering loss geometry has limited efficacy.MGDA can assign a larger weight to en-fr while worsening performance on other languages.
- WMT Machine Translation: GradVac significantly outperforms all baselines, improving both average BLEU and every individual task relative to naive joint training.The gain is larger for En→Any than Any→En, consistent with less similar and more conflicting gradients in En→Any.
- Ablation Studies: The all-layer GradVac variant outperforms whole-model application, constant objectives underperform EMA variables, and GradVac variants outperform PCGrad counterparts.These ablations support fine-grained parameter objectives, task-aware objectives based on relatedness, and preemptive gradient-similarity objectives.
- Massively Multilingual Tasks: On the massive dataset and XTREME, GradVac improves multilingual models across resource levels and consistently outperforms naive joint finetuning and other multi-task baselines.The massive-dataset comparison reports consistent gains across tasks, while XTREME results demonstrate practicality for general multilingual tasks.
5 RELATED WORK
Related work frames multilingual models as both language-processing systems and multi-task learners. Prior approaches address cross-lingual ability and optimization through task relatedness, task-specific routing, gradient modification, and Pareto optimization.
- Multilingual Modeling: Prior multilingual research studies cross-lingual ability, training limitations, and the effects of multilingual modeling across many languages.This work includes multilingual language models and neural machine translation systems.
- Multi-task Optimization: Multi-task optimization research improves training through task relatedness, task-specific routing, direct gradient alteration, or Pareto-solution search.These approaches treat multilingual models as multi-task learning methods.
6 CONCLUSION
The paper studies multilingual loss geometry through gradient similarity and proposes GradVac to align updates according to task closeness. Experiments report superior performance on massive multilingual tasks and argue that the approach may extend beyond multilingual settings.
- Conclusion: The paper systematically studies loss geometry through gradient similarity and proposes GradVac based on the resulting findings.The approach leverages linguistic proximity among multilingual tasks.
- Conclusion: More geometrically consistent updates aligned with task closeness improve performance on massive multilingual tasks.The authors also state that gradient conflicts can hurt model performance and that the approach is generic beyond multilingual settings.
A DATA STATISTICS
The dataset spans diverse languages, linguistic families, scripts, and highly uneven training-data sizes. The selected tasks are grouped into high-, medium-, low-, and extremely low-resource categories.
- Dataset Composition: The selected study set contains 25 languages and 50 language pairs spanning different language families and training-data sizes.The languages are chosen for typological diversity and varied data availability.
- Language Diversity: The dataset contains languages from diverse scripts and linguistic families, with the selected 25 languages covering 10 families or branches.Language-family membership is used to define language proximity because related languages tend to share vocabulary and grammatical constructs.
- Data Sizes: Training data range from 10^5 to 10^9 sentence pairs, divided into High (7), Medium (8), Low (5), and Extremely Low (5) resource levels.High-resource tasks exceed 10^8 pairs; medium-resource tasks contain 10^7–10^8 pairs; tasks below 5 million pairs are extremely low-resource.
B TRAINING DETAILS
The models use shared encoder-decoder Transformer parameters across language pairs, trained with Adam, learning-rate scheduling, and temperature-based multilingual sampling.
- All language pairs share the encoder-decoder Transformer parameters, including word embeddings and the output softmax layer.
- Training uses Adam with a 30K-step linear warm-up to 1e-3, followed by inverse-square-root decay.
- Language pairs are sampled at each training step using temperature-based sampling, with T=5 used in most experiments.
C.1 DATA
The WMT analysis selects eight languages from four families and evaluates gradient similarities across training checkpoints, finding clustering by linguistic proximity and stronger patterns in smaller models.
- Data: The WMT analysis uses eight languages forming 16 language pairs across four language families, drawing data from WMT 13 through WMT 19.
- Analysis: Gradient similarities are computed on validation sets at all checkpoints and averaged across checkpoints for visualization.
- Results: Gradient similarities cluster by language proximity, with languages from the same family sharing the most similar gradients.
- Results: The observed gradient-similarity patterns are more evident in smaller Transformer models.
D ADDITIONAL RESULTS ON OUR DATASET
Additional experiments show that gradient similarities depend on language direction, batch size, and task relatedness, motivating the proposed Gradient Vaccine method.
- Gradient geometry: En→Any models show smaller and more negative gradient cosine similarities than Any→En models, making them harder to train.
- Gradient geometry: Nepali has mostly conflicting gradients with other languages except those in the same language family.
- Batch size: Smaller 125k-token batches produce lower cosine similarities than larger batches, indicating less similar gradients.
- Task relatedness: Gradient similarities primarily reflect task relatedness, since semantically identical sentence pairs can still have negative cosine similarity across languages.
- Method: Gradient Vaccine is introduced as the method studied in the subsequent algorithmic description.
E.1 METHOD DERIVATION
Gradient Vaccine adjusts task gradients toward pair-specific similarity targets, generalizing PCGrad and supporting convergence under stated convexity and smoothness assumptions.
- Method derivation: Gradient Vaccine replaces each task gradient with a vector in the span of two task gradients to achieve a preset cosine similarity.
- Method derivation: The update rule permits arbitrary pair-specific similarity objectives, enabling task relatedness to determine gradient alignment targets.
- Relation to prior work: PCGrad is a special case of Gradient Vaccine when the target gradient cosine similarity is zero.
- Implementation: GradVac uses exponential moving averages to set pair-wise objectives and can operate independently across model components or layers.
- Theory: When task-gradient cosine similarity is negative, the method is largely equivalent to PCGrad and inherits PCGrad’s convergence analysis.
- Theory: Under convex differentiable losses and Lipschitz-continuous gradients, sufficiently small steps make the objective strictly decrease toward its optimum.
- Hyper-parameters: The best task-selection setting applies GradVac to low-resource languages, while an effective EMA window of 100 training steps works best.
- Evaluation: The experiments evaluate multilingual NMT and multilingual language-model transfer tasks, including XTREME POS and NER settings.
G.2 POS RESULT
Gradient Vaccine outperforms other multi-task methods on average in the POS tasks evaluated in Table 9.
- Gradient Vaccine outperforms other multi-task methods on average on the POS tasks in Table 9.