GRET: Global Representation Enhanced Transformer

We propose to use capsules with dynamic routing to extract the specific and suitable features from the local representations for stronger global representation modeling, which is an effective and strong feature extraction method (?; ?)¹¹1Other details of the Capsule Network are shown in ? (?) .. Features from hidden states of the encoder are summarized into several capsules, and the weights (routes) between hidden states and capsules are updated by dynamic routing algorithm iteratively.

Formally, given an encoder of the Transformer which has $M$ layers and an input sentence $\textbf{x}=\{x_{1},\cdots,x_{i},\cdots,x_{I}\}$ which has $I$ words. The sequence of hidden states $\textbf{H}^{m}=\{\textbf{h}^{m}_{1},\cdots,\textbf{h}^{m}_{i},\cdots,\textbf{h}^{m}_{I}\}$ from the $m^{\text{th}}$ layer of the encoder is computed by

where the $\textbf{Q}^{m}_{e}$, $\textbf{K}^{m-1}_{e}$ and $\textbf{V}^{m-1}_{e}$ are query, key and value vectors which are same as $\textbf{H}^{m-1}$, the hidden states from the $m-1^{\text{th}}$ layer. The $\text{LN}(\cdot)$ and $\text{SAN}(\cdot)$ are layer normalization function (?) and self-attention network (?), respectively. We omit the residual network here.

Then, the capsules $\textbf{U}^{m}$ with size of $K$ are generated by $\textbf{H}^{m}$. Specifically, the $k^{\text{th}}$ capsule $\textbf{u}^{m}_{k}$ is computed by

where $q(\cdot)$ is non-linear squash function (?):

and $\textbf{c}_{k}$ is computed by

where the matrix B is initialized by zero and whose row and column are $K$ and $I$, respectively. This matrix will be updated when all capsules are produced.

The algorithm is shown in Algorithm 1. The sequence of capsules $\textbf{U}^{m}$ could be used to generate the global representation.

Different from the original capsules network which use a concatenation method to generate the final representation, we use an attentive pooling method to generate the global representation²²2Typically, the concatenation and other pooling methods, e.g. mean pooling, could be used here easily, but they will decrease 0.1$\sim$0.2 BLEU in machine translation experiment.. Formally, in the $m^{\text{th}}$ layer, the global representation is computed by

where $\text{FFN}(\cdot)$ is a feed-forward network and the $\hat{\textbf{s}}^{m}$ is computed by

This attentive method can consider the different roles of the capsules and better model the global representation. The overview of the process of generating the global representation are shown in Figure 1.

Traditionally, the Transformer model only fed the last layer’s hidden states $\textbf{H}^{M}$ as representations of input sentence to the decoder to generate the output sentence. Following this, we can feed the last layer’s global representation $\textbf{s}^{M}$ into the decoder directly. However, current global representation only contain the intra-layer information, the other layers’ representations are ignored, which were shown to have different aspects of meaning in previous work (?; ?). Based on this intuition, we propose a layer-wise recurrent structure to aggregate the representations generated by employing the capsule network on all layers of the encoder to model a complete global representation.

The layer-wise recurrent structure aggregates each layer’s intra global state by a gated recurrent unit (?, GRU) which could achieve different aspects of information from the previous layer’s global representation. Formally, we adjust the computing method of $\textbf{s}^{m}$ by

where the $\text{ATP}(\cdot)$ is the attentive pooling function computed by Eq 7-9. The GRU unit can control the information flow by forgetting useless information and capturing suitable information, which can aggregate previous layer’s representations usefully. The layer-wise recurrent structure could achieve a more exquisite and complete representation. Moreover, the proposed structure only need one more step in the encoding stage which is not time-consuming. The overview of the aggregation structure is shown in Figure 2.

We conduct experiments on machine translation and text summarization tasks. In machine translation, we employ our approach on four language pairs: Chinese to English (ZH$\rightarrow$EN), English to German (EN$\rightarrow$DE), German to English (DE$\rightarrow$EN), and Romanian to English (RO$\rightarrow$EN) ³³3http://www.statmt.org/wmt17/translation-task.html. In text summarization, we use LCSTS (?) ⁴⁴4http://icrc.hitsz.edu.cn/Article/show/139.html to evaluate the proposed method. These data-sets are public and widely used in previous work, which will make other researchers replicate our work easily.

In machine translation, on the ZH$\rightarrow$EN task, we use WMT17 as training set which consists of about 7.5M sentence pairs. We use newsdev2017 as validation set and newstest2017 as test set which have 2002 and 2001 sentence pairs, respectively. On the EN$\rightarrow$DE and DE$\rightarrow$EN tasks, we use WMT14 as training set which consists of about 4.5M sentence pairs. We use newstest2013 as validation set and newstest2014 as test set which have 2169 and 3000 sentence pairs, respectively. On the RO$\rightarrow$EN task, we use WMT16 as training set which consists of about 0.6M sentence pairs. We use newstest2015 as validation set and newstest2016 as test set which has 3000 and 3002 sentence pairs, respectively.

In text summarization, following in ? (?) , we use PART I as training set which consists of 2M sentence pairs. We use the subsets of PART II and PART III scored from 3 to 5 as validation and test sets which consists of 8685 and 725 sentence pairs, respectively.

In machine translation, we apply byte pair encoding (BPE) (?) to all language pairs and limit the vocabulary size to 32K. In text summarization, we limit the vocabulary size to 3500 based on the character level. Out-of-vocabulary words and chars are replaced by the special token UNK.

For the Transformer, we set the dimension of the input and output of all layers as 512, and that of the feed-forward layer to 2048. We employ 8 parallel attention heads. The number of layers for the encoder and decoder are 6. Sentence pairs are batched together by approximate sentence length. Each batch has 50 sentence and the maximum length of a sentence is limited to 100. We set the value of dropout rate to 0.1. We use the Adam (?) to update the parameters, and the learning rate was varied under a warm-up strategy with 4000 steps (?). Other details are shown in ? (?) . The number of capsules is set 32 and the default time of iteration is set 3. The training time of the Transformer is about 6 days on the DE$\rightarrow$EN task. And the training time of the GRET model is about 12 hours when using the parameters of baseline as initialization.

After the training stage, we use beam search for heuristic decoding, and the beam size is set to 4. We measure translation quality with the NIST-BLEU (?) and summarization quality with the ROUGE (?).

We employ the proposed GRET model on four machine translation tasks. All results are summarized in Table 1. For fair comparison, we reported several Transformer baselines with same settings reported by previous work (?; ?; ?) and researches about enhancing local word level representations (?; ?; ?; ?).

The results on the WMT17 ZH$\rightarrow$EN task are shown in the second column of Table 1. The improvement of our GRET model could be up to 1.22 based on a strong baseline system, which outperforms all previous work we reported. To our best knowledge, our approach attains the state-of-the-art in relevant researches.

Then, the results on the WMT14 EN$\rightarrow$DE and DE$\rightarrow$EN tasks, which is the most widely used data-set recently, are shown in the third and fourth columns. The GRET model could attain 28.46 BLEU (+1.26) on the EN$\rightarrow$DE and 33.79 BLEU (+1.45) on the DE$\rightarrow$EN, which are competitive results compared with previous studies.

To verify the generality of our approach, we also experiment it on low resource language pair of the WMT16 RO$\rightarrow$EN task. Results are shown in the last column. The improvement of the GRET is 0.89 BLEU, which is a material improvement in low resource language pair. And it shows that proposed methods could improve translation quality in low resource scenario.

Experimental results on four machine translation tasks show that modeling global representation in the current Transformer network is a general approach, which is not limited by the language or size of training data, for improving translation quality.

Besides machine translation, we also employ proposed methods in text summarization, a monolingual generation task, which is an important and typical task in natural language generation.

The results are shown in Table 2, we also reports several popular methods in this data-set as a comparison. Our approach achieves considerable improvements in ROUGE-1/2/L (+1.63/+1.70/+1.49) and outperforms other work with same settings. The improvement on text summarization is even more than machine translation. Compared with machine translation, text summarization focuses more on extracting suitable information from the input sentence, which is an advantage of the GRET model.

Experiments on the two tasks also show that our approach could work on different types of language generation task and may improve the performance of other text generation tasks.