Optimal Corpus Aware Training for Neural Machine Translation

TL;DR

This paper introduces Optimal Corpus Aware Training (OCAT), a lightweight fine-tuning method that improves neural machine translation by effectively leveraging corpus metadata, resulting in higher accuracy and robustness.

Contribution

OCAT fine-tunes pre-trained models by adjusting only a small set of corpus-related parameters, enhancing translation quality with less risk of overfitting and hyperparameter sensitivity.

Findings

+3.6 chrF improvement on WMT23 English-Chinese translation

+1.8 chrF improvement on English-German translation

Comparable or better performance than state-of-the-art fine-tuning methods

Abstract

Corpus Aware Training (CAT) leverages valuable corpus metadata during training by injecting corpus information into each training example, and has been found effective in the literature, commonly known as the "tagging" approach. Models trained with CAT inherently learn the quality, domain and nuance between corpora directly from data, and can easily switch to different inference behavior. To achieve the best evaluation, CAT models pre-define a group of high quality data before training starts which can be error-prone and inefficient. In this work, we propose Optimal Corpus Aware Training (OCAT), which fine-tunes a CAT pre-trained model by freezing most of the model parameters and only tuning small set of corpus-related parameters. We show that OCAT is lightweight, resilient to overfitting, and effective in boosting model accuracy. We use WMT23 English to Chinese and English to German translation tasks as our test ground and show +3.6 and +1.8 chrF improvement, respectively, over vanilla training. Furthermore, our approach is on-par or slightly better than other state-of-the-art fine-tuning techniques while being less sensitive to hyperparameter settings.