Robust Optimization for Multilingual Translation with Imbalanced Data

TL;DR

This paper introduces CATS, a novel optimization algorithm that improves multilingual translation models by addressing data imbalance issues, leading to better performance on low-resource languages without compromising high-resource ones.

Contribution

The paper proposes Curvature Aware Task Scaling (CATS), a new optimization method that adaptively rescales gradients to improve training for imbalanced multilingual translation models.

Findings

CATS improves low-resource language BLEU scores by 0.8 to 2.2 points.

CATS enhances multilingual training robustness across different model sizes and batch configurations.

The method maintains high-resource language performance while boosting low-resource language translation quality.

Abstract

Multilingual models are parameter-efficient and especially effective in improving low-resource languages by leveraging crosslingual transfer. Despite recent advance in massive multilingual translation with ever-growing model and data, how to effectively train multilingual models has not been well understood. In this paper, we show that a common situation in multilingual training, data imbalance among languages, poses optimization tension between high resource and low resource languages where the found multilingual solution is often sub-optimal for low resources. We show that common training method which upsamples low resources can not robustly optimize population loss with risks of either underfitting high resource languages or overfitting low resource ones. Drawing on recent findings on the geometry of loss landscape and its effect on generalization, we propose a principled optimization algorithm, Curvature Aware Task Scaling (CATS), which adaptively rescales gradients from different tasks with a meta objective of guiding multilingual training to low-curvature neighborhoods with uniformly low loss for all languages. We ran experiments on common benchmarks (TED, WMT and OPUS-100) with varying degrees of data imbalance. CATS effectively improved multilingual optimization and as a result demonstrated consistent gains on low resources ($+0.8$ to $+2.2$ BLEU) without hurting high resources. In addition, CATS is robust to overparameterization and large batch size training, making it a promising training method for massive multilingual models that truly improve low resource languages.