Aligned Cross Entropy for Non-Autoregressive Machine Translation

TL;DR

This paper introduces aligned cross entropy (AXE), a new loss function for non-autoregressive machine translation models that improves alignment and performance by using a differentiable dynamic programming approach.

Contribution

The paper proposes AXE, an innovative loss function that enhances training of non-autoregressive models by better handling word order alignment, leading to state-of-the-art results.

Findings

AXE improves translation quality on WMT benchmarks.

AXE achieves new state-of-the-art performance for non-autoregressive models.

AXE effectively models word order without autoregressive factors.

Abstract

Non-autoregressive machine translation models significantly speed up decoding by allowing for parallel prediction of the entire target sequence. However, modeling word order is more challenging due to the lack of autoregressive factors in the model. This difficultly is compounded during training with cross entropy loss, which can highly penalize small shifts in word order. In this paper, we propose aligned cross entropy (AXE) as an alternative loss function for training of non-autoregressive models. AXE uses a differentiable dynamic program to assign loss based on the best possible monotonic alignment between target tokens and model predictions. AXE-based training of conditional masked language models (CMLMs) substantially improves performance on major WMT benchmarks, while setting a new state of the art for non-autoregressive models.