\documentclass[tombow,english]{jnlp_2.3}
\usepackage{jnlpbbl_1.3}
\usepackage[dvipdfm]{graphicx}
\usepackage{gentombow}
\usepackage{amsmath}
\settombowbleed{5mm}
\usepackage{hangcaption_jnlp}
\allowdisplaybreaks
\usepackage{array}
%%%%\usepackage{udline}
%%%%\setulminsep{1.2ex}{0.2ex}
%%%%\let\underline=\ul
%%%%\makeatletter
%%%%\def\@{\spacefactor3000\space}
%%%%\makeatother
\renewcommand{\refname}{References}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% author's definition
%%%%\usepackage[dvipdf]{graphicx}
%%%%\usepackage{amsmath}
%%%%\usepackage{subcaption}
\usepackage{amssymb}
\usepackage{amsthm}
\usepackage{url}
\renewcommand\UrlFont{\rmfamily}
%\usepackage{pgfplots}
\usepackage{multirow}
%%%%\usepackage{tablefootnote}
%%%%\usepackage[multiple]{footmisc}
\newtheorem{axiom}{Axiom}
%%%%\newtheorem{define}{Definition}
%%%%\newtheorem{prop}{Proposition}
\DeclareMathOperator{\softmax}{softmax}
%%%%\DeclareMathOperator{\rand}{rand}
\DeclareMathOperator{\diff}{diff}
\DeclareMathOperator{\pert}{prt}
%%%%\DeclareMathOperator{\IG}{IG}
\DeclareMathOperator{\RNN}{RNN}
\DeclareMathOperator{\LSTM}{LSTM}
\DeclareMathOperator{\MLM}{MLM}
\DeclareMathOperator{\PFN}{PFN}
\DeclareMathOperator{\FFNN}{FFNN}
\DeclareMathOperator{\SA}{SA}
\DeclareMathOperator{\Adj}{adjust}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\Volume{27}
\Number{3}
\Month{September}
\Year{2020}

\received{2019}{11}{26}
\revised{2020}{3}{23}
%%%%\rerevised{2006}{5}{6}
\accepted{2020}{4}{15}
%%%%\Issuetitle{「　　　　」特集号}

\setcounter{page}{531}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\esubject{General Paper}
\etitle{Encoder-Decoder Attention $\neq$ Word Alignment: \\Axiomatic Method of Learning Word Alignments \\for Neural Machine Translation}
\eauthor{Chunpeng Ma\affiref{HIT} \and Akihiro Tamura\affiref{Ehime} \and Masao Utiyama\affiref{NICT} \and Tiejun Zhao\affiref{HIT} \and Eiichiro Sumita\affiref{NICT}} 
\eabstract{%
	The encoder-decoder attention matrix has been regarded as the (soft) alignment model for conventional neural machine translation (NMT) models such as RNN-based models. However, we show empirically that this is not true for the Transformer. By comparing the Transformer with the RNN-based NMT model, we find two inherent differences, and accordingly present two methods of capturing word alignments in the Transformer. Furthermore, instead of focusing on the Transformer, we present three axioms for the attention mechanism that captures word alignments, and propose a new attention mechanism based on these axioms that we have termed the axiomatic attention mechanism (AAM), and which is applicable to any NMT models. The AAM depends on a perturbation function, and we apply several perturbation functions to the AAM, including a novel function based on the masked language model \cite{devlin2018bert}. Using the AAM to guide the training of an NMT model improved both the translation performance and the learning of word alignments of the NMT model. Our research sheds light on the interpretation of sequence-to-sequence models on neural machine translation.
}
\ekeywords{Machine Translation, Deep Learning, Axiomatic Method}

\headauthor{Ma et al.}
\headtitle{Encoder-Decoder Attention $\neq$ Word Alignment}

\affilabel{HIT}{}{Harbin Institute of Technology}
\affilabel{Ehime}{}{Ehime University}
\affilabel{NICT}{}{National Institute of Information and Communications Technology}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{document}

\maketitle

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Introduction}

Different architectures have been proposed for neural machine translation (NMT), ranging from recurrent neural networks (RNNs) \cite{sutskever2014sequence,chen2018best}, convolutional neural networks (CNNs) \cite{gehring2017convolutional,gehring2017convolutional2}, to the Transformer \cite{vaswani2017attention}. The interpretation of neural networks is still an open problem, which is extremely difficult. In this paper, we focus only on one specific important module, the attention mechanism \cite{bahdanau2014neural,luong2015effective}, in the context of machine translation.

For machine translation, the encoder-decoder (hereinafter, enc-dec) attention matrix represents the relationship between the words in the target sentence and the words in the source sentence. Because of its high correlation with word alignments \cite{bahdanau2014neural}, it is thus usually regarded as a soft-alignment model \cite{shankar2018surprisingly,knowles2018context,ghader2017what}. Needless to say, word alignment is extremely important to machine translation \cite{brown1993mathematics,och03systematic,blunsom2006discriminative}. By pinpointing where the word alignment is located in an NMT model, we can improve the performance conveniently. It has been shown to be effective in producing attention matrices that are similar to the gold word alignments for RNN-based NMT models \cite{liu2016neural,mi2016supervised}. Machine translation methods which purely based on word alignments are also proposed \cite{alkhouli2018alignment}.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 1
\begin{figure}[b]
\vspace{-0.5\Cvs}
\begin{center}
\includegraphics{27-3ia2f1.eps}
\end{center}
%%%%\caption{Gold}\label{fig:eggold}
%%%%\caption{RNN model}\label{fig:egrnn}
%%%%\caption{Transformer}\label{fig:egrans}
\caption{Attention matrices of the (a) gold alignment, (b) RNN-based NMT model, and (c) Transformer.}
\label{fig:example}
\vspace{-1\Cvs}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

Surprisingly, we found that for the Transformer,\footnote{Throughout this paper, we focus on Transformer-base \cite{vaswani2017attention}, which is the foundation of larger models and other variants \cite{devlin2018bert}.} 
the enc-dec attention matrix cannot fully capture the word alignment. For example, in Figure \ref{fig:example}, the correspondence between the bold English and Chinese words is not caught by the enc-dec attention matrix of the Transformer, while the attention matrix of the RNN-based model is consistent with the gold word alignments. We found that the incorrespondence with the word alignment is common for the enc-dec attention matrix of the Transformer. Quantitative observations and detailed discussions are provided in Section \ref{sec:compare}.\footnote{Although \cite{tang2018analysis} also observed this, they did not provided statistical experiment results and solutions.}

These observations conflict with previous reports stating that other NMT models learn word alignments by the attention mechanism; it is therefore worthwhile studying the reason why the attention matrices in the Transformer is \textit{not} able to capture the word alignments. We answer this question in Section \ref{sec:explain}. We provided experimental support for our answer, and accordingly added some modules into the Transformer so that it learned better word alignments.

However, these methods are not general, in that they depend on the specific structure of the Transformer. We therefore propose three axioms that the \textit{real} alignments must satisfy, whatever the detailed model structure is (Section \ref{sec:axioms}). Based on these axioms, we propose a new attention mechanism that we have named the axiomatic attention mechanism (AAM). One module of the AAM, the perturbation function, influences the performance significantly, and this is discussed in Section \ref{sec:perturb}. We propose a novel perturbation function based on the masked language model (MLM) proposed recently \cite{devlin2018bert}. The MLM-based perturbation function allows the NMT model to learn better word alignments. Finally, we use the proposed AAM to guide the training process of the NMT model, resulting in better word alignments and a better translation performance (Section \ref{sec:method}).


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Background}

This section gives a brief overview of the RNN-based NMT model and the Transformer, particularly the corresponding attention mechanism.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{RNN-based NMT}

RNN-based models have been the dominant NMT method for several years and have been adopted in large-scale systems \cite{wu2016google,crego2016systran}. With the assistance of an attention mechanism, the translation performance becomes impressive.

The model is composed of two modules: an encoder and a decoder. Given the source sentence $\mathbf{w}_s\in\mathbb{Z}_{+}^{n_s}$, where each element represents the index of the source word, the encoder first generates a sequence of embeddings, denoted by $\mathbf{I}_{s}^{1}\in\mathbb{R}^{n_s\times d}$. Here, $n_s$ is the length of the source sentence, and $d$ is the dimension of the model. Then, the encoder calculates a sequence of hidden vectors $\mathbf{H}_{s}^{1}\in\mathbb{R}^{n_s\times d}$ as 
\pagebreak
follows:\footnote{Throughout this paper, square brackets refer to the columns or elements of a matrix or a vector, and superscripts/subscripts refer to different matrices or vectors.}$^{,}$\footnote{Here we only described the case of uni-direction.}
\begin{equation}
\mathbf{H}_{s}^{1}[i]=\RNN(\mathbf{H}_{s}^{1}[i-1],\mathbf{I}_{s}^{1}[i]).\label{eq:hs}
\end{equation}
The function $\RNN$ can be either the gated recurrent unit (GRU) \cite{cho2014learning} or the long short-term memory (LSTM) 
\linebreak
\cite{hochreiter1997long}. The input to the next layer is
\begin{equation}
\mathbf{I}_{s}^{2} = \mathbf{H}_{s}^{1}.\label{eq:is}
\end{equation}
Repeating this calculation, we obtain the final representation of the source sentence $\mathbf{O}_{s}=\mathbf{I}_{s}^{L_{s}+1}$, where $L_s$ is the number of layers of the encoder.

The decoder generates the target sentence $\mathbf{w}_{t}\in\mathbb{Z}_{+}^{n_t}$ as follows:
\begin{equation}
\mathbf{w}_{t}[j]=\arg\max(\softmax(\FFNN(\mathbf{O}_{t}[j]))),
\end{equation}
where $n_t$ is the length of the target sentence, $\FFNN$ is a feed-forward neural network, $\mathbf{O}_{t}=\mathbf{I}_{t}^{L_{t}+1}$ is the representation of the target sentence, and $L_t$ is the number of decoder layers. For the $l$-th layer, the output $\mathbf{I}_{t}^{l+1}$ is calculated as follows:
\begin{equation}
\mathbf{I}_{t}^{l+1}=f(\mathbf{C}^{l},\mathbf{H}_{t}^{l}),\label{eq:it}
\end{equation}
where $f$ is a non-linear function. $\mathbf{H}_{t}^{l}$ is calculated as follows:
\begin{equation}
\mathbf{H}_{t}^{l}[j]=\RNN(\mathbf{H}_{t}^{l}[j-1],\mathbf{I}_{t}^{l}[j]).\label{eq:ht}
\end{equation}
The vector $\mathbf{C}^{l}[j]$ is the weighted average of the columns of $\mathbf{O}_{s}$:
\begin{equation}
\mathbf{C}^{l}[j]=\sum_{i=1}^{n_s}\softmax((\mathbf{H}_{t}^{l}[j])^{\top}\mathbf{O}_{s})\mathbf{O}_{s}[i].\label{eq:c}
\end{equation}
This is the RNN-based attention (RA).\footnote{Here we only described one kind of RA, i.e., dot-product attention \cite{luong2015effective}, which is widely used.}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Transformer}

Nowadays, the Transformer has recently become the state-of-the art NMT model in terms of both accuracy and speed. It also follows the encoder-decoder paradigm. Different from RNN-based models, instead of using Eq. \ref{eq:hs}, the encoder calculates the hidden vector as follows:\footnote{For the Transformer, besides word embeddings, $\mathbf{I}_{s}^{1}$ also encodes the word positions (i.e., positional encodings).}
\newpage\vspace*{-17pt}
\begin{equation}
\mathbf{H}_{s}^{1}[i]=\SA(\mathbf{I}_{s}^{1},\mathbf{I}_{s}^{1}[i],\mathbf{I}_{s}^{1}),\label{eq:hs-trans}
\end{equation}
where the self attention (SA) is given by
\begin{equation}
\SA(\mathbf{K},\mathbf{Q},\mathbf{V})=\softmax\Big(\frac{\mathbf{K}^{\top}\mathbf{Q}}{\sqrt{d}}\Big)\mathbf{V}.\label{eq:sa}
\end{equation}
The input to the next layer is calculated as follows:\footnote{Residual connection \cite{he2016deep} has been taken into consideration.}
\begin{equation}
\mathbf{I}_{s}^{2}=\FFNN(\mathbf{H}_{s}^{1})+\mathbf{H}_{s}^{1}.\label{eq:is-trans}
\end{equation}

For the decoder, instead of using Eq. \ref{eq:it}, $\mathbf{I}_{t}^{l+1}$ is calculated as follows:
\begin{align}
\mathbf{I}_{t}^{l+1} & =\FFNN(\mathbf{H}_{t}^{l})+\mathbf{H}_{t}^{l},\\
\mathbf{H}_{t}^{l}[j] & =\SA(\mathbf{O}_{s},\mathbf{\widetilde{H}}_{t}^{l}[j],\mathbf{O}_{s}),\label{eq:ht-trans}\\
\mathbf{\widetilde{H}}_{t}^{l}[j] & =\SA(\mathbf{I}_{t}^{l},\mathbf{I}_{t}^{l}[j],\mathbf{I}_{t}^{l}).
\end{align}
The $\SA$ in Eq. \ref{eq:ht-trans} is the enc-dec attention mechanism (e.g., the attention matrix shown in Figure \ref{fig:example}), which is the focus of this study.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Comparison of RA and SA}
\label{sec:compare}

As illustrated in Figure \ref{fig:example}, we conjecture that for the Transformer, the enc-dec attention matrix is not able to capture the word alignment. To verify our conjecture, we empirically compare the effects of the RNN-based NMT models and the Transformer on learning word alignments.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Datasets and Baselines}

Both the RNN-based model and the Transformer are implemented by \texttt{OpenNMT} \cite{klein2017opennmt}.\footnote{\url{http://opennmt.net}} We used the LDC corpus\footnote{LDC2002E18, LDC2003E07, LDC2003E14, Hansards portion of LDC2004T07, LDC2004T08, and LDC2005T06.} (about $1.4$ M sentence pairs) to train the English--Chinese models. Translation performance is evaluated by word-level BLEU \cite{papineni2002bleu}. We use NIST MT 2002 ($878$ sentence pairs) for evaluation.

We used a manually annotated word alignment corpus (the THU corpus)\footnote{\url{http://nlp.csai.tsinghua.edu.cn/~ly/systems/TsinghuaAligner/TsinghuaAligner.html}} to evaluate 
\pagebreak
the performance of learning word alignments. The processed THU corpus contains $854$ sentence pairs. Following \cite{liu2016neural}, when evaluating the alignment error rate (AER), we forced the decoder to generate the words in the reference dataset, and chose the source words with the highest alignment confidence for evaluation.

Table \ref{tab:baseline} lists the baseline results. Although the Transformer yielded a significantly better translation performance, the quality of the learned word alignments given by the enc-dec attention matrices was far worse than the RNN-based model.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 1
\begin{table}[b]
\caption{Baseline results}
\label{tab:baseline}
\input{02table01.tex}
\vspace{4pt}\small
The AER of the Transformer is evaluated based on the enc-dec attention matrix of the last head of the multi-headed attention.
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Effects of Heads of SA}

Previous studies have demonstrated that adjusting the number of heads \cite{tang2018why} or averaging the heads \cite{zenkel2019adding} of a multi-headed SA may affect the model severely. We therefore compared the effects of different heads on learning word alignments.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 2
\begin{table}[b]
\hangcaption{AER of different heads between the last encoder layer and the first decoder layer for the Transformer}
\label{tab:head}
\input{02table02.tex}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

Table \ref{tab:head} gives the AERs of different heads between the last encoder layer and the first decoder layer for the Transformer.\footnote{Intuitively, the connection between these two layers is the ``closest'' connection between the encoder and the decoder, so we use the attention mechanism between them to calculate the AERs. We also ran the experiment on other combinations of encoder-decoder layers (e.g., last encoder layer \& last decoder layer), and found no difference.} We can see that although the AERs were different, all heads failed to learn word alignments well. All the AERs were much higher than the AER of the RNN-based model (AER = 18.9). 

Table \ref{tab:heads} gives the results of adjusting the number of heads for the Transformer, and the results of averaging all heads. Although the translation performance (almost) increased with the number of heads, the AER varied little. The average of all heads did not give a better AER, either. Therefore, for the Transformer, simply adjusting the number of heads cannot make it learn word alignments as good as the RNN-based NMT model.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 3
\begin{table}[b]
\caption{Effect of adjusting the number of heads of the SA}
\label{tab:heads}
\input{02table03.tex}
\vspace{4pt}\small
``\#'' indicates the number of heads. ``$h_{last}$'' indicates using the last head to evaluate the AER, ``aver'' indicates using the average of all heads for evaluation. Note that ``\#=8,($h_{last}$)'' is simply the baseline.
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

In conclusion, there is no significant difference among different heads for learning word alignments, although there is some difference in other tasks \cite{marecek2019balustrades,clark2019bert,vig2019analyzing}. Therefore, through this paper, we simply choose the last head, which occasionally performed the best in Table \ref{tab:head}.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 2
\begin{figure}[b]
\begin{center}
\includegraphics{27-3ia2f2.eps}
\end{center}
\caption{Evolution of the AER (on the THU corpus) and word-level BLEU (on the MT02 corpus)}
\label{fig:bleu}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Evolution during training}

Figure \ref{fig:bleu} shows the evolution of word-level BLEU (on the MT02 corpus) and 
\pagebreak
AER (on the THU corpus) of two models during training. As expected, the BLEUs of both models increased gradually, and the Transformer yielded a better performance. However, the AER of the Transformer was much worse than the AER of the RNN-based model. Furthermore, the AER of the Transformer increased gradually. This indicates that for the Transformer, the enc-dec attention matrix is unable to capture the word alignment.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Effects of Pseudo-supervised attention}
\label{sec:superattn}

Following \cite{mi2016supervised}, we used gold word alignments to guide model training.\footnote{For the Transformer, only the last head is supervised.} Specifically, we first converted the gold word alignments to $0$-$1$ matrices, and smoothed the matrices using a Gaussian filter $\mathcal{N}(0,0.5)$. Then we added the following item to the loss function:
\begin{equation}
d(\mathcal{A},\mathcal{A}^{*})=\Bigg(\sum_{t=1}^{n_t}\sum_{s=1}^{n_s}(\mathcal{A}^{*}[t,s]-\mathcal{A}[t,s])^2\Bigg)^{\frac{1}{2}},\label{eq:superattn}
\end{equation}
where $\mathcal{A}^{*}$ is the smoothed gold word alignment matrix, and $\mathcal{A}$ is the learned attention matrix (i.e., Eq. (6) for RNN-based NMT model and Eq. (11) for the Transformer) learned by the NMT model. Because the golden word alignments are obtained automatically instead of manually, we call the methods in this section ``pseudo-supervised attention.''

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 4
\begin{table}[b]
\caption{Results of pseudo-supervised attention}
\label{tab:sa}
\input{02table04.tex}
\vspace{4pt}\small
The AER of GIZA++ on THU corpus is also reported. Paired bootstrap resampling significance test was carried out. ``$*$'' indicates that the results were significantly better (or worse) than the baseline ($p<0.05$).
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

For our experiments, the gold word alignments of the LDC training corpus were obtained using \texttt{GIZA++} \cite{och03systematic}.\footnote{\url{http://www.fjoch.com/GIZA++.html}} Table \ref{tab:sa} gives the results of pseudo-supervised attention. For the RNN-based NMT model with pseudo-supervised attention, the performance of using both word alignment and machine translation improved. However, for the Transformer, although the word alignment performance increased significantly, the machine translation performance deteriorated. This verified our conjecture: the enc-dec attention matrix learned by the Transformer is not able to capture the word alignments, and therefore gold word alignment misleads the training process.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Why SA is NOT Word Alignment}
\label{sec:explain}

Given the empirical analysis above, a natural question is: why does the SA is not able to capture the word alignment? This is the focus of this section.

Comparing the formulas for RA (Eq. \ref{eq:c}) and SA (Eq. \ref{eq:sa}), we found that except for the constant $\sqrt{d}$, they are almost identical. Therefore, we conjecture that a key point in finding an answer lies in the calculation of the hidden representations in the encoder (i.e., $\mathbf{H}_{s}^{l}$) and the decoder (i.e., $\mathbf{H}_{t}^{l}$). We found that there are two main differences in the calculation of the representations: exposure and dependence. We propose two methods for learning better word alignments based on these two differences and show their effectiveness.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Difference in Exposure}

When calculating the representations of the encoder, besides the apparent difference in the mathematical formulas used (Eqs. \ref{eq:hs} or \ref{eq:hs-trans}), the first significant difference concerns the exposure. 

For the encoder of the RNN-based NMT model, $\mathbf{H}_{s}^{l}[i]$ is calculated recurrently by the RNN, where $\mathbf{H}_{s}^{l}[k]$ is used for calculation. For the forward direction, $k\in[1,i-1]$, while for the backward direction, $k\in[i+1,n_s]$. Although the representations of both directions are then concatenated, for one specific element of the final representation, only several $\mathbf{H}_{s}^{l}[k]$ are exposed.

Conversely, for the encoder of the Transformer, for any specific element of the final representation, all $\mathbf{H}_{s}^{l}[k] (k\in[1,n_s])$ are exposed. Although the additional information used by the Transformer may improve the translation performance, it may mislead the learning of word alignments.

Because of this, we conjecture that adding an encoding mask to the encoder of the Transformer could allow the Transformer to learn better word alignments. Specifically, when calculating $\mathbf{H}_{s}^{1}[i]$, instead of using Eq. \ref{eq:hs-trans}, we calculated as follows:
\begin{equation}
\mathbf{H}_{s}^{1}[i]=\SA(\mathbf{I}_{s}^{1}[\leq i],\mathbf{I}_{s}^{1}[i],\mathbf{I}_{s}^{1}[\leq i]).
\end{equation}
This is called the forward encoding mask. We can also calculate $\mathbf{H}_{s}^{1}[i]$ as follows:
\begin{equation}
\mathbf{H}_{s}^{1}[i]=\SA(\mathbf{I}_{s}^{1}[\geq i],\mathbf{I}_{s}^{1}[i],\mathbf{I}_{s}^{1}[\geq i]),
\end{equation}
which is called the backward encoding mask. Other encoder layers are calculated similarly.

Table \ref{tab:encmask} gives the result of adding the encoding mask. We can see although the AER was still worse than for the RNN-based NMT model, adding the encoding mask indeed made the Transformer learn better word alignments. Furthermore, although the encoding mask reduced the information used, the deterioration in the translation performance on the THU corpus was not as severe. For MT 2002 corpus, the BLEU even increased to some extent, which is surprising. Thus, the difference in exposure is indeed one of the causes of the failure on learning word alignments for the Transformer.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 5
\begin{table}[t]
\caption{Effects of encoding mask}
\label{tab:encmask}
\input{02table05.tex}
\vspace{4pt}\small
``F'' and ``B'' are forward and backward encoding masks, respectively. ``$*$'': significantly better than Baseline ($p<0.05$).
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Difference in Dependence}

In the calculation of the decoder representation (Eqs. \ref{eq:ht} and \ref{eq:ht-trans}), we found that for the RNN-based NMT model, $\mathbf{H}_{t}^{l}[j]$ depends on $\mathbf{H}_{t}^{l}[j-1]$ (i.e., the history of the \textit{same} layer). However, for the Transformer, $\mathbf{H}_{t}^{l}[j-1]$ is not involved. The value of $\mathbf{H}_{t}^{l}[j]$ depends on $\mathbf{I}_{t}^{l}$, which further depends on $\mathbf{H}_{t}^{l-1}$ (i.e., the history of the \textit{previous} layer). This is also true for the encoder (Eqs. \ref{eq:hs} and \ref{eq:hs-trans}). We argue that the history of the same layer is important for learning word alignments.

Incorporating historical information on the same layer into the Transformer is not straightforward, due to the parallel nature of the Transformer. We offer a stopgap solution: adding recurrent layers above the final layer of the encoder and/or decoder of the Transformer. Specifically, the final encoder representation $\widetilde{\mathbf{O}}_{s}$ is calculated as follows:
\begin{equation}
\widetilde{\mathbf{O}}_{s}=\PFN(\LSTM(\FFNN(\mathbf{O}_{s})+\mathbf{O}_{s})),
\end{equation}
where $\PFN$ is a position-wise feed-forward network \cite{vaswani2017attention}. The final decoder representation $\widetilde{\mathbf{O}}_{t}$ can be calculated similarly.

Intuitively, $\FFNN(\mathbf{O}_{s})+\mathbf{O}_{s}$ in Eq. (16) is a kind of the residual connection \cite{he2016deep}, which has been proved very useful for representation learning. Using an $\LSTM$ makes different representations in different time stamps connect with each other. $\PFN$ is proved useful in the original Transformer, so we also used it in our network.

Table \ref{tab:encpfn} shows the results of adding bi-directional LSTMs above the final layers of the encoder/decoder. First, we can see that adding recurrent layers to the encoder or the decoder or both did not make much difference. Second, we can see that using more recurrent layers made the model learn better word alignments while using one recurrent layer made the performance worse, because of the lack of model capacity of the added layer. Last, in most cases, the translation performance improved. This verified empirically that the difference in dependence is also one of the causes of the failure on learning word alignments for the Transformer.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 6
\begin{table}[t]
\caption{Results of adding bi-directional recurrent layers above the final layers of the encoder/decoder}
\label{tab:encpfn}
\input{02table06.tex}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Axiomatic Attention Mechanism}
\label{sec:axioms}

Despite the strategies proposed above make the Transformer successfully learn word alignments, the strategies are not general and are restricted to the Transformer, making it difficult to extend them to other NMT models.

We propose a universally applicable model: the axiomatic attention mechanism (AAM). Although axiomatic methods have been used for various tasks \cite{sundararajan2017axiomatic,yoshikawa2018combining}, as far as we know, the proposed method is the first axiomatic method in MT.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Axioms}

A large number of previous studies have shown the important role of word alignment in machine translation \cite{brown1993mathematics,tu2016modeling}. We argue that an attention mechanism that can capture word alignments should satisfy the following three axioms.

\begin{axiom}
\textbf{Sensitivity.} The attention mechanism should catch the corresponding changes in the source and target sentences. 
\end{axiom}

For example, in Figure \ref{fig:example}, when the English word \textit{``show''} is changed to \textit{``news,''} the corresponding Chinese word is also changed. The attention mechanism should catch this information and assign high scores to this correspondence (i.e., the second English word should be aligned with the third Chinese word).

\begin{axiom}
\textbf{Model Agnostic.} Whatever the detailed structure of the model is, the correspondence between the source and target words should be invariant.
\end{axiom}

For example, in Figure \ref{fig:example}, when generating the bold Chinese word, the bold English word should always be paid attention to because of the inherent semantic correspondence between the Chinese and English words, no matter whether the model is the Transformer or the RNN-based model.

\begin{axiom}
\textbf{Context Dependency.} In different contexts (or sentences), correspondence between two words may change.
\end{axiom}

For example, in Figure \ref{fig:example}, the English word \textit{``was''} does not correspond to any of the Chinese words, which is correct for the given sentence. However, for other English sentences (e.g., \textit{``It was a book''}), the English word \textit{``was''} should correspond to a Chinese word.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Definition}

Given the above axioms, we formally define the AAM as follows:
\begin{equation}
\mathbf{a}_{i}(\mathbf{S},\mathbf{T})=\softmax\Big(\frac{\diff(\mathbf{\widetilde{T}}^{(i)},\mathbf{T})}{\sum_{j}\diff(\mathbf{\widetilde{S}}^{(i)},\mathbf{S})[j]}\Big).\label{eq:am}
\end{equation}
This is a vector of length $n_t$. $\mathbf{S}$ is the representation matrix of the source sentence (e.g., $\mathbf{I}_{s}^{1}$ or $\mathbf{O}_{s}$), and $\mathbf{T}$ is the representation of the target sentence. The $j$-th element of $\mathbf{a}_{i}$ indicates the attention score of the $j$-th target word to the $i$-th source word. There are $n_s$ different $\mathbf{a}_i$ in total.

Matrix $\mathbf{\widetilde{S}}^{(i)}$ is a matrix identical to $\mathbf{S}$, except for the $i$-th column. It is calculated as follows:
\begin{equation}
\mathbf{\widetilde{S}}^{(i)}=[\mathbf{S}[1],\ldots,\mathbf{S}[i-1],\pert(\mathbf{S}[i]),\mathbf{S}[i+1],\ldots,\mathbf{S}[n_s]].\label{eq:pert}
\end{equation}
The function $\pert$ denotes the perturbation function. On changing the source representation to $\mathbf{\widetilde{S}}^{(i)}$, the target representation will correspondingly change, which is denoted by $\mathbf{\widetilde{T}}^{(i)}$. 

The function $\diff$ calculates the semantic difference between two matrices. Given matrices $\mathbf{X}$ and $\mathbf{X}^\prime$ of the same size, $n\times d$, the value of $\diff$ is a vector of length $n$ whose $j$-th element is calculated as follows:
\pagebreak
\begin{equation}
\diff(\mathbf{X},\mathbf{X}^{\prime})[j]=\frac{1}{\pi}\arccos\Big(\frac{\mathbf{X}[j]\cdot\mathbf{X}^{\prime}[j]}{\Vert\mathbf{X}[j]\Vert\Vert\mathbf{X}^{\prime}[j]\Vert}\Big),
\end{equation}
which is proportional to the \textit{angular} difference of two vectors.
The final attention matrix is calculated as follows:
\begin{equation}
\mathbf{A}(\mathbf{S},\mathbf{T})=\Adj\Big([\mathbf{a}_{1}(\mathbf{S},\mathbf{T}),\ldots,\mathbf{a}_{n_s}(\mathbf{S},\mathbf{T})]\Big).\label{eq:ast}
\end{equation}
The size of matrix $\mathbf{A}$ is $n_{t}\times n_s$. The function ``$\Adj$'' is the adjustment function and is defined as follows:
\begin{equation}
\Adj(\mathbf{X})[p,q]=\bigg(1-\frac{1}{p-1}\sum_{k=1}^{p-1}\mathbf{X}[k,q]\bigg)\mathbf{X}[p,q].
\end{equation}
This function is used to prevent all target words from paying attention to the same source word. Intuitively, if the $q$-th source word has been assigned a large attention score by some target words previous to $p-1$, then the $p$-th target word should decrease its attention to this word.

\textbf{Satisfaction of Axiom 1.} It is obvious that for $\diff(\mathbf{\widetilde{S}}^{(i)},\mathbf{S})$, all elements except the $i$-th element are equal to $0$. Assume that the $p$-th and $q$-th columns of $\mathbf{\widetilde{T}}^{(i)}$ are different from those of $\mathbf{T}$, then for $\mathbf{a}_{i}(\mathbf{S},\mathbf{T})$, all elements except the $p$-th and $q$-th elements are equal to $0$. Therefore, Axiom 1 is satisfied.

\textbf{Satisfaction of Axiom 2.} We can see that our definition (Eq. \ref{eq:am}) does not depend on any characteristics of the network structures. As long as the representations of source and target sentences can be calculated, which is true for any NMT model, the axiomatic attention mechanism can be obtained. Therefore, Axiom 2 is satisfied.

\textbf{Satisfaction of Axiom 3.} Because the representations of the words depend on other words in the sentence, which is true for all NMT models, it is obvious that for different sentences, $\mathbf{A}(\mathbf{S},\mathbf{T})$ is different. Therefore, Axiom 3 is satisfied. Note that the satisfaction of this axiom needs the following restriction: word representations depend on other words, therefore AAM is applicable to hidden representations, instead of word embeddings. This restriction can be satisfied easily for NMT. 


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Intuitive Interpretation}

Axiom 1 is the main consideration of our definition of AAM (Eq. \ref{eq:am}). First, we want to represent the corresponding changes on two sides, therefore we change the words in the source sentence one by one, and use $\widetilde{\mathbf{S}}^{(i)}$ to represent the change of the $i$-th source word, and use $\widetilde{\mathbf{T}}^{(i)}$ to represent the corresponding change of the target sentence.

Then, because the change of representations can be either large or small, so we use ``$\diff$'' to calculate the scale of change. Intuitively, if a small perturbation on the source side causes a large perturbation on the target side, then a large attention should be assigned. Hence, we use the division.

Last, by considering Axiom 2, we do not import any features depending on specific network structures. Therefore, in Eq. \ref{eq:am}, only representations of two sides appeared, while functions like $\RNN$ or $\SA$ did not appear.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Perturbation Function}
\label{sec:perturb}

We found that the choice of the perturbation function (i.e., $\pert$ in Eq. \ref{eq:pert}) affected the AAM significantly, so we discuss this topic in detail.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Perturbation based on rules}

A straightforward way to define the perturbation function is to manually establish some rules. A simple rule is to let the values of the perturbation function always be zero vectors, i.e.,
\begin{equation}
\pert(\mathbf{x})=\mathbf{0}.\label{eq:pert-0}
\end{equation}
Another possible definition is as follows:
\begin{equation}
\pert(\mathbf{x})=-\mathbf{x}.
\end{equation}
We may also rotate the vector by a fixed angle $\theta$:
\begin{equation}
\begin{aligned}
\pert(\mathbf{x})[2k-1]&=\mathbf{x}[2k-1]\cos\theta-\mathbf{x}[2k]\sin\theta,\\
\pert(\mathbf{x})[2k]&=\mathbf{x}[2k-1]\sin\theta+\mathbf{x}[2k]\cos\theta,\label{eq:pert-theta}
\end{aligned}
\end{equation}
where $k\in[1,d/2]$, and $\theta$ is a hyper-parameter.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Perturbation based on language model}

Although the definitions of rule-based perturbation functions are succinct, these functions may hurt the smoothness of the sentence. For example, in Figure \ref{fig:example}, if we change the second English word to a verb (e.g., \textit{``The \textbf{eat} today ...''}), which may correspond to a $\theta$ in Eq. \ref{eq:pert-theta}, the sentence becomes so weird that it is impossible for the NMT model to generate the Chinese word corresponding to \textit{``eat''} at the third position. 

It is therefore important to perturb the source sentence while keeping its smoothness. A simple solution is to resort to the language model. However, standard language models generally predict the words left-to-right, making it inappropriate when perturbing the first several words.

Recently, Devlin et al. (2019) proposed the masked language model (MLM) which is able to predict a masked word given all other words in the sentence. Based on the MLM, we propose a perturbation function. Intuitively, we mask the words in the source sentence one by one, use the MLM to predict the masked word, and assign the embedding of the predicted word to the perturbation function. Formally, it is defined as follows:
\begin{equation}
\pert(\mathbf{x},i)= \mathbf{x}^{(t)},~~~~~~~~t= \MLM(\mathbf{X},i),
\end{equation}
where $\mathbf{x}^{(t)}$ is the word embedding of the $t$-th word in the vocabulary. $\MLM(\mathbf{X},i)$ predicts the index of the masked word given sentence $\mathbf{X}$ with the $i$-th word masked. If the predicted masked word is identical to the original word, the word with the second-highest score is returned. 

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 7
\begin{table}[b]
\caption{AER of AAM with different perturbation functions on two NMT models}
\label{tab:pert-rule}
\input{02table07.tex}
\vspace{4pt}\small
Two rows are the results of letting $\mathbf{S}$ in Eq. \ref{eq:am} be $\mathbf{I}_{s}^{1}$ and $\mathbf{O}_{s}$, respectively. For the ``MLM'' columns, the first letter indicates the scale of the model is base (B) or large (L), and the second letter indicates the training corpus is cased (C) or uncased (U). Bold: the best results for rule/MLM-based perturbation functions.
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{Comparison of perturbation functions}

Table \ref{tab:pert-rule} gives the AERs of AAM using different perturbation functions. For Eq.~\ref{eq:pert-theta}, we only reported the result of $\theta=\frac{\pi}{2}$, which performs the best among all possible $\theta$. For MLM, we used the pre-trained models provided by \texttt{BERT}.\footnote{\url{https://github.com/google-research/bert}} Note that for the MLM-based perturbation function, because the embedding matrix is needed for calculation, $\mathbf{X}$ can only be $\mathbf{I}_{s}^{1}$. We compare different perturbation functions as follows.

\textbf{Rule-based vs.\ MLM-based.} For both models, we found that MLM-based functions performed consistently better than rule-based ones. This verified the importance of keeping the smoothness of the sentence for perturbation. Table \ref{tab:mlm-example} provides the qualitative evidence. It gives the top-$2$ words predicted by ``L+U'' model by masking each word respectively, and the words selected by the perturbation function. We can see that the perturbed sentences are all smooth and correct sentences.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 8
\begin{table}[t]
\caption{Results of pseudo-supervised attention}
\label{tab:sa-aam}
\input{02table08.tex}
\vspace{4pt}\small
AAM with the MLM-based perturbation function was guided by word alignments learned by \texttt{GIZA++}. ``AER (Mat.)'' are the results given by the attention matrices. ``$*$'': significantly better ($p<0.05$). For the case of ``+ Pseudo-supervision'', we used AAM, because the AERs are lower.
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 9
\begin{table}[t]
\caption{Top-$2$ words predicted by MLM}
\label{tab:mlm-example}
\input{02table09.tex}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\textbf{Choice of MLM.} Four different MLMs were evaluated: cased or uncased, base or large. We found that uncased models performed better. This is because our training corpus is uncased. Table \ref{tab:pert-rule} also shows that large models performed better although the difference was small in some cases. 

\textbf{RNN vs.\ Transformer.} For the Transformer, the MLM-based perturbation functions gave significantly better AERs than the baseline, while for the RNN-based model, the differences in performance were moderate. This is because the attention matrix of the RNN model has been able to learn word alignments so well that they are difficult to beat. This is consistent with previous studies that have proven the effectiveness of RNNs on learning word alignments \cite{tamura2014recurrent}.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{Improving NMT with AAM}
\label{sec:method}

In this section, we present a method of improving NMT with AAM by incorporating AAM into the training of an NMT model. The method is based on the pseudo-supervised attention technique introduced in Section \ref{sec:superattn}. In particular, we replace the matrix $\mathcal{A}$ in Eq. \ref{eq:superattn} by the AAM (i.e., the matrix $\mathbf{A}(\mathbf{S},\mathbf{T})$ in Eq. \ref{eq:ast}). The MLM-based perturbation function is used. Note that the guidance word alignments (i.e., $\mathcal{A}^{*}$) were obtained using \texttt{GIZA++}, which indicates that the proposed AAM can be used when annotated word alignments are not given. We used the ``L+U'' model to predict perturbed words because of the low AERs for both NMT models.


Table \ref{tab:sa-aam} gives the results. Note that for both NMT models, the ``enc-dec attention matrix'' module always exist, whether the AAM is used or not. We found that for the RNN-based NMT model, both the attention matrix and the AAM yield better word alignments, while for the Transformer, only the AAM gives a better AER. This again verifies the difference in the roles played by attention matrices of the two models. The BLEU of the RNN-based model only increased slightly, because the performance of the baseline RNN-based model on learning word alignments is quite good. However, the BLEU of the Transformer increased considerably. The AAM caught the word alignments learned by the model, and the guidance of the AAM made the model perform better. 

We can see that we may improve the performance of NMT with AAM. By building the AAM into the model, we can manipulate the model conveniently. Therefore, the proposed AAM is not only an explanation, but also an improvement of the model.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Related Work}

For statistical machine translation model, word alignment is an important part. There are a lot of researches to improve the performance of word alignment \cite{brown1993mathematics,och03systematic,blunsom2006discriminative}, so that the performance of machine translation can be improved.

For neural machine translation, there are many researches capturing word alignments by attention mechanisms, to some extent. For RNN-based models \cite{bahdanau2014neural,luong2015effective} and CNN-based models \cite{gehring2017convolutional,gehring2017convolutional2}, word alignments and attention mechanisms are strongly correlated, therefore guiding the attention mechanisms with correct word alignments improved machine translation \cite{liu2016neural,mi2016supervised}. 

However, the behavior of the attention mechanism in the Transformer is quite different. The interpretation of the attention mechanism in the Transformer and its relationship to the word alignment are the points of focus of this paper.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{Conclusion}

We found that the attention matrix of the Transformer is unable to capture the word alignments. We presented reasons for this, and proposed two methods to make the attention matrix learn better word alignments. Instead of focusing on the Transformer, we proposed three axioms that should be satisfied by all kinds of attention mechanisms in NMT models. Based on these axioms, we proposed our axiomatic attention mechanism. There are various choices for the perturbation function that forms the core of the proposed AAM. We found that the perturbation function based on the MLM performed the best. By incorporating the AAM into the training of NMT models with the pseudo-supervised attention, both the translation performance and word alignment performance of the Transformer improved. 


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\bibliographystyle{jnlpbbl_1.5}
\bibliography{02refs}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{biography}

\bioauthor[:]{Chunpeng Ma}{%
is currently a researcher in Fujitsu Laboratory. He received his B.E., M.E., Ph.D. from Harbin Institute of Technology in 2010, 2014 and 2020, respectively. His advisor was Prof. Tiejun Zhao. From July 2013 to January 2014, he interned at Baidu Inc. with instruction from Xianchao Wu. From October 2015 to September 2016, he interned at National Institute of Information and Communications Technology, with the instruction from Akihiro Tamura, cooperating with Lemao Liu. His research interests include syntactic parsing, machine translation and deep learning.
}

\bioauthor[:]{Akihiro Tamura}{%
received his B.E., M.E., and Dr. Eng. from the Tokyo Institute of Technology in 2005, 2007, and 2013, respectively. After serving as a researcher at NEC and NICT, he worked as an assistant professor at Ehime University from 2017. Since 2020, he has been working as an associate professor at Doshisha University. His research interests include natural language processing and machine learning. He is a member of the Association for Computational Linguistics, the Information Processing Society of Japan, the Japanese Society for Artificial Intelligence, and the Association for Natural
Language Processing.
}

\bioauthor[:]{Masao Utiyama}{%
received the Ph.D. degree in engineering from the University of Tsukuba, Tsukuba, Japan, in 1997. He is currently a Research Manager in the National Institute of Information and Communications Technology, Tokyo, Japan. He is a member of the Association for Natural Language Processing, and Asia-Pacific Association for Machine Translation.
}

\bioauthor[:]{Tiejun Zhao}{%
Doctor, professor of Research Center of Language Technology, School of Computer Science and Technology, Harbin Institute of Technology. He is associate dean of SCST and director of Ministry of Education-Microsoft Key Laboratory of NLP \& Speech in HIT. He is associate director of Machine Translation Subject Committee of Chinese Information Processing Society, the member of editorial board of Journal of Chinese Information Processing, Journal of Automation, the senior member and associate director of Chinese Computing Committee of China Computer Federation. His research fields include: natural language understanding, content-based web information processing, applied artificial intelligence.
}

\bioauthor[:]{Eiichiro Sumita}{%
received the M.S. degree in computer science from the University of Electro-Communications in 1982 and the Ph.D. degree in engineering from Kyoto University in 1999. He is the Associate Director General, ASTREC (Advanced Speech Translation Research and Development Promotion Center). He is currently the fellow of NICT. He is a member of the Association for Natural Language Processing, director of Japan Translation Federation, and the president of Asia-Pacific Association for Machine Translation.
}

\end{biography}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\biodate



\end{document}

