Diformer: Directional Transformer for Neural Machine Translation

The sampling of direction in DLM allows us to train the generation of all directions simultaneously, we introduce the detailed method in this section.

As we all know that the training of Transformer (Vaswani et al., 2017) can be paralleled with teacher forcing, achieved by feeding $y_{1:N-1}$ (context) to the model at once and computing the loss on $y_{2:N}$ (target). The context and target sequence can be easily created by a shifting operation that aligns $y_{i-1}$ to $y_{i}$.

Diformer can also be trained in a similar way, but before that, we have to make a slight change when implementing the computation of the likelihood in Eq 6 due to the difficulty of creating the context sequence $Y_{z_{i}}$ with complicated dependencies. The original equation aims to compute the likelihood on the ground-truth sequence $Y$ where each token is conditioned on a customized context determined by the sampled direction, meaning that the context sequence cannot be shared as Transformer does. Creating specialized context for every token is non-trivial especially when encountered with position changing caused by the shifting when $z_{i}=R\ \text{or}\ L$.

For the convenience of the implementation, we fix the input sequence $y_{1:N}$ and create a new target sequence $Y^{*}$ where tokens are accordingly shifted with the sampled directions:

where $j=i+1$ for $z_{i}=R$, $j=i-1$ for $z_{i}=L$ and $j=i$ for $z_{i}=S$. When training on large corpus with random sampling on directions, we can say that $P(Y^{*}|X)\approx P(Y|X)$ theoretically.

Formally, let the source and target sequence as $X=[x_{1},...,x_{|X|}]$ and $Y=[y_{1},...,y_{N}]$ where $N$ is the target length. Then, we uniformly sample a direction instruction sequence $Z=[z_{1},...,z_{N}]$ with $N$ elements, where $z_{1}$ and $z_{N}$ are fixed to be $R$ and $L$ as they are [BOS] and [EOS], which can only be used to predict tokens inside the sequence for the AR setting, and can never be masked in the NAR setting.

The input sequence $Y_{\text{in}}$ is created by directly copying from ground-truth $Y$, which will be masked accordingly in the decoder to create the disentangled context.

According to the sampled direction sequence $Z$, we can now create the modified target sequence $Y^{*}$ by shifting tokens in $Y$ based on $z_{i}$, which is shown in Figure 1.

To be compatible with the NAR decoding, we also predict the target length $P(N|X)$ with the same way as (Ghazvininejad et al., 2019). Note that the predicted length is only used for NAR decoding, the AR decoding still terminates when [EOS] or [BOS] is generated for L2R and R2L setting.

Finally, the cross-entropy loss is used for both generation ($\mathcal{L}_{\text{DLM}}$) and length prediction ($\mathcal{L}_{\text{LEN}}$) task, the overall loss can be obtained by adding them together:

where $\lambda$ is the factor on which the best performance can be obtained with the value of 0.1, after searched from 0.1 to 1.0 in the experiment.

An embedding matrix is used to map the categorical variable $z_{i}$ into the hidden space, denoted as $\delta$, $\delta(z_{i})\in\mathbb{R}^{d_{\text{model}}}$ where $d_{\text{model}}$ is hidden size of the model. For simplicity, we directly use $z_{i}$ to represent the embedded direction at position $i$ in following sections.

The joint training of L2R and R2L can be problematic with the positional embedding of the original Transformer since the index is counted in a uni-directed order, which can be used for cheating under the bi-directional scenario since future positional index can leak information of the sentence length.

To solve this, we propose to make the positional embedding directed, achieved by encoding the position index counted oppositely based on the direction with separate parameters:

where $\overrightarrow{Pe}$ and $\overleftarrow{Pe}$ are different embedding matrics to encode position indices counting from L2R ($\overrightarrow{i}$) or R2L ($\overleftarrow{i}$) accordingly. More detailed description can be found in Figure 1.

Finally, we add encoded position and direction embeddings together as the initial hidden-state $h_{i}^{l=0}$ for the computation of $Q^{l=0}$ in the first self-attention layer $h_{i}^{0}=p_{z_{i}}+z_{i}$:

In DisCo, to prevent the information leakage from the disentangled context, the input representation for computing $K,V$ is de-contextualized by directly reusing the projection of input embeddings $k_{i},v_{i}=\text{Proj}(w_{i}+p_{i})$. In Diformer, we have to further remove the positional information since the directed positional embedding can still be used for cheating in the computation of self-attention across layers.

Completely removing the the positional information on $K,V$ and only using the word-embedding $w_{i}$ can be harmful to the performance. Therefore, we propose an alternative solution by replacing the removed absolute positional embedding with the relative positional embedding proposed in (Shaw et al., 2018) for two reasons: 1) The relative position is computed in a 2 dimensional space, meaning that $p_{ij}$ and $p_{kj}$ for token $y_{j}$ is not shared between $y_{i}$ and $y_{k}$, which satisfies our requirements that each token in the context should have the position information only used for $y_{i}$ but not shared for $y_{k}$. 2) The position information is only injected during the computation of self-attention without affecting the original word embedding used in $K,V$.

Formally, we directly use the method in (Shaw et al., 2018) but replace the hidden representation for computing $K,V$ with word embeddings:

where $h_{i}^{l^{\prime}}$ is the output of the self-attention in current layer, $w_{j}$ is the word embedding, $p_{ij}^{V},p_{ij}^{K}$ are embedded relative positions, $W^{Q},W^{K},W^{V}$ are parameters for $Q,K$ and $V$, $h_{i}^{l-1}$ is the last layer’s hidden state, $d_{\text{head}}$ is the hidden size of a single head. Two parameter matrics are used as embeddings — $Re^{K}$ and $Re^{V}$, with the shape of $[2k+1,d_{\text{head}}]$, where $k$ is the max length. $p_{ij}$ is obtained by embedding the distance between $i$ and $j$ clipped by the maximum length $k$.

Finally, a customized attention mask (see Figure 1) is created during training to simulate specific dependencies based on the sampled direction sequence $Z$ with following rules:

• •

  If $z_{i}=R$, all tokens for $j>i$ will be masked.

• •

  If $z_{i}=L$, all tokens for $j<i$ will be masked.

• •

  If $z_{i}=S$, $y_{i}$ and a subset of randomly selected tokens will be masked following the method in (Ghazvininejad et al., 2020), excluding [BOS] and [EOS].

Diformer can generate a sequence with 4 modes including L2R and R2L for AR decoding, mask-predict and parallel easy-first for NAR decoding.

For the AR decoding, the model works exactly same as the conventional Transformer, except that for each step, a fixed direction $z_{i}=R$ or $L$ should also be also be given, together with previously generated tokens, making it a pure-linear autoregressive generation. Beam search can be directly used in both L2R and R2L decoding. For the NAR decoding, the model uses mask-predict or easy-first by applying specific masking operation during each iteration, where all tokens are assigned with $z=S$. Length beam can be used to further improve the performance. Detailed examples are shown in the Appendix.

More importantly, we find that the multi-directional property of Diformer can be used for reranking, which is quite beneficial for the NAR decoding. Specifically, compared to other NAR models that uses an external AR model for reranking, Diformer can do it all by its own without introducing additional computational costs. For example, it first refines 5 candidates with 8 iterations and performs reranking with the rest of 2 iterations by re-using the encoder states and scoring candidates with L2R and R2L modes, which is equivalent to the computational cost of a 10-stepped refinement reported in CMLM. The scores computed in two directions are averaged to obtain the final rank. Experimental results show that 8 steps of refinement + 2 steps of reranking obtains significant performance improvements compared to 10 steps of refinement without re-ranking. It can also be used for AR decoding, where all tokens are scored under the reversed direction, e.g. generating with L2R and scoring with R2L. We name this method as self-reranking.

We evaluate Diformer on 4 benchmarks including WMT14 En$\leftrightarrow$De (4.5M sentence pairs) and WMT16 En$\leftrightarrow$Ro (610k sentence pairs). The data is preprocessed in the same way with (Vaswani et al., 2017; Lee et al., 2018), where each sentence is tokenized with Moses toolkit (Koehn et al., 2007) and encoded into subwords using BPE (Sennrich et al., 2016). We follow Gu et al. (2018); Ghazvininejad et al. (2019); Zhou et al. (2020) to create the knowledge distilled (KD) data with L2R Transformer-big and Transformer-base for En$\leftrightarrow$De and En$\leftrightarrow$Ro, the reported performance in the overall results are all obtained by training on the KD data.

We follow the same configurations with previous works (Vaswani et al., 2017; Ghazvininejad et al., 2019, 2020) on hyperparameters: $n_{(encoder+decoder)\_layers}$ = 6 + 6, $n_{heads}$ = 8, $d_{hidden}$ = 512, $d_{FFN}$ = 2048. For customized components in Diformer, we tune the max relative distance $k$ in [1,8,16,256] and find that $k=256$ obtains best performance. Adam (Kingma and Ba, 2015) is used for optimization with 128k tokens per batch on 8 V100 GPUs. The learning rate warms up for 10k steps to 5e-4 and decays with inversed-sqrt. Models for En$\leftrightarrow$De and En$\leftrightarrow$Ro are trained for 300k and 100k steps, last 5 checkpoints are averaged for final evaluation. We set beam size as 4 and 5 for AR and NAR decoding. When decoding in NAR mode, we set the max iteration for mask-predict and easy-fist decoding as 10 without using any early-stopping strategy. For fair comparison, we reduce the max iteration to 8 when decoding with self-reranking in NAR model. Our model is implemented with pytorch (Paszke et al., 2019) and fairseq (Ott et al., 2019). BLEU (Papineni et al., 2002) is used for evaluation.

For the comparison to unified-models (Mansimov et al., 2020; Tian et al., 2020), Diformer outperforms others in all generation directions, by obtaining more than 1.5 BLEU.

As discussed in the section 1, their support on multiple generation strategies is achieved by applying certain position selection strategy on the masked language model or generating with certain permutation with the permutation language model. This creates the gap between the training and inference since a specific decoding strategy might not be fully trained with the generalized objective as analyzed in (Mansimov et al., 2020). So, compared to both, we use the task-specific modelling in exchange for better performance by abandon certain flexibility, thus makes the learned distribution to be same with the one used in decoding, which answers why Diformer performs better.

For the En$\leftrightarrow$De dataset, since we use a larger teacher model (Transformer-big), therefore, we only compare Diformer with same sized Transformer-base trained on the raw and distilled data. The Diformer outperforms Transformer trained on the raw data with a large margin and reaches the same level to the one trained on distilled data. Interesting, the best performance of Diformer are usually obtained by the R2L decoding and the reranked results on L2R, the reason of it will be further discussed in ablation study sections. For the En$\leftrightarrow$Ro dataset, Diformer can also obtain similar performance compared to the same sized Transformer trained on the distilled data produced by a same sized teacher.

Diformer is also competitive to a series of NAR models including iterative-refinement based and fully NAR models. We speculate the strong performance of Diformer comes from the joint training of AR and NAR, since it is similar to the multi-task scenario, where tasks are closely correlated but not same. This could be beneficial for the task that is more difficult i.e. NAR, because the learned common knowledge on AR tasks could be directly used in it. By applying the self-reranking method, Diformer could obtain additional 0.5 BLEU over the strong baseline.

We train Diformer not only with distilled data but also with raw data as shown in table 2. The degradation of NAR decoding when training on raw data is not surprising which is a common problem faced by all NAR models. However, the performance of AR decoding also degrades. We speculate that on the raw data, the difficulty of learning to generate from straight and R2L increased significantly, making the model to allocate more capacity to fit them, resulting in the negative influence on the performance of L2R. We verify this by fixing $z_{i}=R$ for all tokens and train the model on raw data. The result confirms it because the performance recovers to its original level. On the contrary, the knowledge distilled data is cleaner and more monotonous (Zhou et al., 2020), making it easier to learn for all directions, and allows the model to allocate balanced capacity on each direction. As for the better performance obtained by R2L decoding, we consider the reason is that, the R2L is able to learn the distilled data generated by the L2R teacher in a complementary manner, making it more efficient to learn the knowledge that cannot be learned by L2R due to the same modeling method.

We also demonstrate the importance of the relative positional embedding by evaluating the model with different maximum relative distance $k$ and obtain the same conclusion (Shaw et al., 2018) — the distance should be at least 8. Meanwhile, we observe that NAR is more sensitive to the positional information, which is reasonable, since the decoding of NAR is conditioned on the bi-directional context, where the positional information contains both distance and direction thereby is more important compared to that in AR.

As shown in the overall results, self-reranking is a useful method to improve the performance especially for NAR decoding. For the AR decoding, the improvements is not that significant since the outputs are already good enough for L2R or R2L, the tiny gap between reranking and generation direction cannot provide enough help, which indicates that using self-reranking for AR is not that profitable compared to NAR.

We further investigate its ability on NAR decoding (mask-predict) given different max iteration number and length beam size, as shown in Figure 2. It clearly shows that without reranking, the incorrect selection on beam candidates may even reduce the performance with larger beam size. The use of self-reranking actually lets the performance and beam-size positively correlated, meaning that exchanging 2 steps of iteration with self-reranking can be profitable with larger beam size. In practical usage of self-reranking, it is critical to find the optimal combination by balancing the beam size and max iteration number so that both high performance and low latency can be obtained.