SlowGuess/NAACL_2022_OCR huggingface dataset:
Yukun Feng $^{1}$ , Feng Li $^{2}$ , Ziang Song $^{1}$ , Boyuan Zheng $^{1}$ , and Philipp Koehn $^{1}$
$^{1}$ Department of Computer Science, Johns Hopkins University
$^{2}$ Department of Computer Science, University of Illinois Urbana-Champaign
{yfeng55, zsong17, bzheng12, phi} @jhu.edu, {fengl3} @illinois.edu
The Transformer architecture has led to significant gains in machine translation. However, most studies focus on only sentence-level translation without considering the context dependency within documents, leading to the inadequacy of document-level coherence. Some recent research tried to mitigate this issue by introducing an additional context encoder or translating with multiple sentences or even the entire document. Such methods may lose the information on the target side or have an increasing computational complexity as documents get longer. To address such problems, we introduce a recurrent memory unit to the vanilla Transformer, which supports the information exchange between the sentence and previous context. The memory unit is recurrently updated by acquiring information from sentences, and passing the aggregated knowledge back to subsequent sentence states. We follow a two-stage training strategy, in which the model is first trained at the sentence level and then finetuned for document-level translation. We conduct experiments on three popular datasets for document-level machine translation and our model has an average improvement of 0.91 s-BLEU over the sentence-level baseline. We also achieve state-of-the-art results on TED and News, outperforming the previous work by 0.36 s-BLEU and 1.49 d-BLEU on average.
Most previous machine translation methods are designed for sentence-level translation. Recent studies have shown that the effective use of contextual information between sentences can achieve better performance in document-level machine translation (Garcia et al., 2015; Maruf and Haffari, 2018; Miculicich et al., 2018; Zhang et al., 2020; Bao et al., 2021). Built on the Transformer model (Vaswani et al., 2017), a general approach is to incorporate neighboring sentence states (Tiedemann and Scher
rer, 2017; Zheng et al., 2020) into the attention mechanism, which has also been widely used in many long sequence modeling methods (Dai et al., 2019; Rae et al., 2020; Yang et al., 2019; Beltagy et al., 2020). Zhang et al. (2018); Maruf et al. (2019) have introduced an additional context encoder to solve the limitation of sentence-level translation, which, however, is separated from the original translation model and context states is only applied on the source side. Other works (Junczys-Dowmunt, 2019; Scherrer et al., 2019; Zhang et al., 2020; Bao et al., 2021) concatenated sentences or the entire document and feed into the attention module of the Transformer. Since more extended contexts may confound attention on meaningful portions of the current sentence, the model is difficult to select valuable inputs from extra contexts to navigate the redundancy of information. Such methods also suffer from the quadratically increasing complexity when documents get longer.
We solve such problems by introducing a memory mechanism to recurrently integrate contextualized knowledge from intermediate state in Transformer layers. As recurrent memory has been widely researched since RNN (Rumelhart et al., 1986), which has been incorporated with Transformer by Transformer-XL (Dai et al., 2019) and further extended by Rae et al. (2020) who compress previous states into a two-layer hidden memory. In our approach, we update the memory through an attention module to select practical information from sentences and reduce the context space into multiple dense vectors in the memory. Besides, we use another attention module to pass the knowledge retained in the memory back to the sentence state in the next step. Such information exchange is expected to convey contextualized dependency between sentences. This memory mechanism can be applied in each layer for both the source and target documents, and our study shows that incorporating memory only in the last layer achieves the
Findings of the Association for Computational Linguistics: NAACL 2022, pages 1409 - 1420 July 10-15, 2022 ©2022 Association for Computational Linguistics
best performance. Also, as sentences are ordered in documents, our model reads one sentence pair at each step, keeping the computational cost as same as the sentence-level translation.
We experiment across three widely used datasets for document-level translation: TED, NEWS, and Europarl, and evaluate our model with s-BLEU and d-BLEU. We first train a vanilla Transformer on sentence-level translation as the baseline and finetune the model for the documents by initializing the memory mechanism to the Transformer. Our model outperforms previous SOTA work by 0.5 s-BLEU and 2.30 d-BLEU on TED, and 0.21 s-BLEU and 0.57 d-BLEU on News. We do not achieve the SOTA result on Europarl, which might be caused by the different results between the baselines for sentence-level translation. However, we further evaluate the improvement of previous works from their reported baseline Transformer, and we achieve the most relative gain on all three datasets. We also analyze our model from the memory usage, long-range effect, context dependency, and computational complexity, and demonstrate the effectiveness and efficiency of our approach in the general understanding of the document machine translation.
Overall, this paper makes several contributions: (i) Our work reduces the contextualized knowledge space of sentences states to multiple dense vectors, and considers the sentence dependency for both source and target documents, while keeping computational complexity in sentence-level. (ii) Our model significantly improves the sentence level baseline by 0.91 s-BLEU average and achieved the SOTA results on TED and News. (iii) Our model shows the effective use of memory, long-range influence, context-dependency across sentences, and decoding efficiency through convincing analysis.
Recurrent Sequence Modeling RNN (Rumelhart et al., 1986) was the first class of models that introduced hidden states as the memory in neural models. Although improved on sequential-oriented tasks, RNN has unsatisfactory learning of long-term information due to gradient vanishing and explosion. LSTM (Hochreiter and Schmidhuber, 1997) improved RNN by introducing gate mechanisms to selectively retain knowledge at each step.
This RNN variant dominated NLP models until the Transformer replaced the memory unit with a
self-attention mechanism and achieved great success in a wide range of NLP applications. Although we cannot deny the robustness and effectiveness of the Transformer model, the quadratically increased computational cost as the increase of token numbers makes Transformer unable to fit the long-range sequence. Some studies (Parmar et al., 2018; Child et al., 2019; Beltagy et al., 2020; Ainslie et al., 2020; Qiu et al., 2020; Zaheer et al., 2020; Martins et al., 2021) try to mitigate this issue by reducing the complexity of the attention module. However, such work still suffers from the problems by unlimited the document length and the document modeling is hard to solve.
Transformer-XL (Dai et al., 2019) broke this dilemma by introducing the recurrent memory into Transformer-based models. It cached previous hidden sentences computation and mapped such states to subsequent sentences states. Theoretically, Transformer-XL could handle infinite length text but storing uncompressed hidden state requires tremendous memory space, which impeded Transformer-XL from good performance on dealing with practical long-sequence tasks. The Compressive Transformer (Rae et al., 2020) further addressed this problem by mapping the evicted hidden state from cached memory to a more compressed representation. However, two-layer caching still requires a huge memory space and may be improved with trainable memories.
Document Machine Translation Machine Translation has been a widely researched area for decades. A series of models have addressed various translation problems (Koehn et al., 2003; Kalchbrenner and Blunsom, 2013; Bahdanau et al., 2015; Luong et al., 2015). As most of them target translation at the sentence level, document-level translation poses a fundamental challenge requiring models to pass intra-sentential information throughout consecutive sequences of sentences, and it has been addressed by Gong et al. (2011); Hardmeier et al. (2013); Pouget-Abadie et al. (2014); Garcia et al. (2015); Koehn and Knowles (2017); Läubli et al. (2018); Agrawal et al. (2018) among others.
Recent studies have attempted to incorporate additional contextual information into the Transformer structure to improve the performance of neural machine translation models further. The intuitive way is to leverage neighboring sentences from paragraphs or the documents (Tiedemann and
Figure 1: An overview of the model architecture, where E and D refers to Encoder and Decoder respectively.
Scherrer, 2017; Maruf and Haffari, 2018; Zheng et al., 2020), demonstrating the effectiveness of the additional contexts. Specifically, in the first class of methodologies for document-level translation, independent from the architecture of vanilla Transformer processing current sentences, some studies (Miculicich et al., 2018; Zhang et al., 2018; Maruf et al., 2019; Voita et al., 2019a,b; Ma et al., 2020; Donato et al., 2021) introduces context-aware components only attend to source or target contexts and usually jointly train with the rest of the network from scratch. The second class of models follows the pattern of concatenating multiple sentences for translation (Agrawal et al., 2018; Scherrer et al., 2019; Junczys-Dowmunt, 2019; Zhang et al., 2020). Such a method is expected to capture the contextual correlations between sentences. However, one of its drawbacks is the quadratically increased computational complexity in the face of longer contexts sequences. Also, longer sequences usually confound document-level attention and sometimes even overlook key information on the current sentences. Bao et al. (2021) uses group masks to introduce locality constraints to reinforce sentence information in multi-head attention to resolve the confounding issue in long contexts.
Our work incorporates the idea of the recurrent memory to document-level machine translation. It follows the locality assumptions by reducing the context space into multiple memory vectors and passes dependencies between sentences. The mechanism to update and output memory is similar to models which store cached bilingual sentence pairs
in the memory to enhance the sentence-level translation (Feng et al., 2017; He et al., 2021; Jiang et al., 2021). We believe our approach is intuitive to efficiently store sentence states and transfer context information across sentences.
Our model is shown in Figure 1. Additional to the vanilla Transformer, we introduce a contextual memory unit and two attention modules to manipulate the memory defined as Update Attention and Output Attention. These modules can be applied at each layer in both the encoder and decoder.
As input sentences are ordered from left to right in the document, our model only reads one sentence every time. The memory is expected to store contextualized information from the input sentence states and convey such knowledge to the next sentence. At each step, the Update Attention step maps the contextual information from the sentence state to the memory, and updates the memory to the next step. Meanwhile, the Output Attention step fuses the information from the current sentence and the contextual memory, and outputs the aggregated knowledge to the remaining modules of the layer.
Formally, we define  $\mathbf{h}^{\mathrm{t}}$  as the sentence state from self-attention module in Transformer layer, and  $\mathbf{M}^{\mathrm{t}}$  refers to the contextual memory M at step t, where t refers to the index of  $t^{\mathrm{th}}$  sentence in the document.  $\mathbf{M}^{\mathrm{t}}$  and  $\mathbf{h}^{\mathrm{t}}$  are updated and outputted as:
$$
\mathrm {M} ^ {\mathrm {t} + 1} = \text {U p d a t e A t t e n t i o n} \left(\mathrm {M} ^ {\mathrm {t}}, \mathrm {h} ^ {\mathrm {t}}\right) \tag {1}
$$
$$
\widetilde {\mathrm {h} ^ {\mathrm {t}}} = \mathrm {O u t p u t A t t e n t i o n} (\mathrm {M} ^ {\mathrm {t}}, \mathrm {h} ^ {\mathrm {t}})
$$
3.1 Contextual Memory
Memory  $\mathrm{M} \in \mathbb{R}^{\mathrm{d}_{\mathrm{M}} \times \mathrm{d}_{\mathrm{model}}}$  where  $\mathrm{d}_{\mathrm{model}}$  refers to the hidden dimension and  $\mathrm{d}_{\mathrm{M}}$  is a hyper-parameter, indicating how many vectors will be allocated for memory. To avoid the redundancy of memory space, we set  $\mathrm{d}_{\mathrm{M}}$  to 16. Detailed analysis is discussed later.
3.2 Update Attention
We update contextual memory through an attention module (Vaswani et al., 2017). Attention is a mapping function between input vectors of query (Q) and key-value (K-V) pairs. The output is the weighted sum of values with corresponding scores.
$$
\mathrm {A t t e n t i o n (Q , K , V)} = \mathrm {S o f t m a x} (\frac {\mathrm {Q K ^ {T}}}{\sqrt {\mathrm {d _ {k}}}}) \mathrm {V}
$$
Multi-Head attention extends the vanilla attention by projecting input vectors  $(\mathrm{Q},\mathrm{K},\mathrm{V})$  to different representation subspaces, and attention is performed in parallel in each head. Attention outputs from multiple heads will be concatenated and projected to the expected space.
$$
\mathrm {M H A} (\mathrm {Q}, \mathrm {K}, \mathrm {V}) = \operatorname {C o n c a t} \left(\text {h e a d} _ {1},.., \text {h e a d} _ {\mathrm {n}}\right) \mathrm {W} ^ {\mathrm {o}}
$$
$$
\mathrm {h e a d} _ {\mathrm {i}} = \text {A t t e n t i o n} \left(\mathrm {Q W} _ {\mathrm {i}} ^ {\mathrm {q}}, \mathrm {K W} _ {\mathrm {i}} ^ {\mathrm {k}}, \mathrm {V W} _ {\mathrm {i}} ^ {\mathrm {v}}\right)
$$
where  $\mathrm{d_k}$  is the hidden dimension of the K,  $\mathrm{W^q}$ ,  $\mathrm{W^k}$ ,  $\mathrm{W^v} \in \mathbb{R}^{d_{\mathrm{model}} \times d_{\mathrm{h}}}$ , and  $\mathrm{W^o} \in \mathbb{R}^{n \times d_{\mathrm{h}} \times d_{\mathrm{model}}}$  are learnable parameters.  $d_{\mathrm{model}}$  and  $d_{\mathrm{h}}$  refer to the hidden dimension of the model and each head.
To update the contextual memory  $\mathbf{M}^{\mathrm{t}}$  to next step, sentence state  $\mathrm{h}^{\mathrm{t}}$  is mapped to  $\mathbf{M}^{\mathrm{t}}$  through the Multi-Head Attention. Both the memory and context state are projected into different sub-spaces and contextualized knowledge is expected to be mapped to each memory vector from different perspectives. The memory at step t is updated as:
$$
\widetilde {\mathrm {M} ^ {\mathrm {t}}} = \operatorname {A d d N o r m} \left(\mathrm {M H A} \left(\mathrm {M} ^ {\mathrm {t}}, \mathrm {h} ^ {\mathrm {t}}, \mathrm {h} ^ {\mathrm {t}}\right)\right) \tag {2}
$$
A Feed-Forward Network is then used to further enhance the memory representation from the attention output.
$$
\mathrm {M} ^ {\mathrm {t} + 1} = \operatorname {A d d N o r m} (\text {F e e d F o r w a r d} (\widetilde {\mathrm {M} ^ {\mathrm {t}}})) \tag {3}
$$
In the memory matrix M, each vector is expected to select contextualized information from different perspectives. However, it is hard to distinguish such vectors since they do not have actual positional meanings, and the same key-value pairs are
mapped to these vectors in the attention phase resulting in the same representation in each memory vector. To solve such a problem, we use the positional encoding  $\mathrm{PE}()$  as introduced in Vaswani et al. (2017) to differentiate multiple memory vectors. M is added by such position-level bias in each update phase.
$$
\mathrm {M} ^ {\mathrm {t}} = \mathrm {M} ^ {\mathrm {t}} + \mathrm {P E} (\mathrm {M} ^ {\mathrm {t}}) \tag {4}
$$
3.3 Output Attention
To map the contextualized knowledge from  $\mathbf{M}^{\mathrm{t}}$  to the sentence state  $\mathrm{h}^{\mathrm{t}}$ , multi-head attention is used to take the representation of  $\mathrm{h}^{\mathrm{t}}$  and  $\mathbf{M}^{\mathrm{t}}$  as query and key-value, respectively.
$$
\widetilde {\mathrm {h} ^ {\mathrm {t}}} = \mathrm {M H A} \left(\mathrm {h} ^ {\mathrm {t}}, \mathrm {M} ^ {\mathrm {t}}, \mathrm {M} ^ {\mathrm {t}}\right) \tag {5}
$$
$\widetilde{\mathrm{h}}_{\mathrm{t}}$  will be passed to the subsequent modules in the Transformer layer.
Similar approaches have been discussed in previous works. Simply increasing the context space does not help but introduces a lot of noise. Instead of incorporating multiple sentences to the context attention, we compress contextualized information into multiple memorized vectors and map such vectors back to the sentence state at the next step. We find that both the BLEU score and the information gained from the context attention space do not increase when the memory length increases from 64 to 128. Therefore, a large context space in M seems redundant for the model to learn, and we find  $\mathrm{d}_{\mathrm{M}} = 16$  for the most effectiveness and efficiency.
3.4 Document Neural Machine Translation
In the task of document-level machine translation, the source and target documents are represented as sequences of sentences  $\mathrm{X} = \{\mathrm{x_t}|1\leq t\leq n\}$ , and  $\mathrm{Y} = \{\mathrm{y_t}|1\leq t\leq n\}$  respectively, where  $\mathbf{t}$  refers to the sentence index. Given a vanilla Transformer and its parameters  $\theta$ , the objective is to maximize the target document probability conditioned on the source document.
$$
\operatorname * {a r g m a x} _ {\theta} \mathrm {P} (\mathrm {Y} | \mathrm {X}, \theta)
$$
Our approach recurrently translates an ordered document sentence by sentence, and the objective is:
$$
\underset {\widetilde {\theta}} {\operatorname {a r g m a x}} \prod_ {t} P \left(Y _ {t} | X _ {\leqslant t}, Y _ {<   t}, \widetilde {\theta}\right)
$$
where  $\widetilde{\theta}$  refers to Transformer parameters including Memory, Update Attention and Output Attention.
Model	TED	News	Europarl
s-BLEU	d-BLEU	s-BLEU	d-BLEU	s-BLEU	d-BLEU
Vaswani et al. (2017)	23.10	-	22.40	-	29.40	-
Miculicich et al. (2018)	24.58	-	25.03	-	28.60	-
Maruf et al. (2019)	24.42	-	24.84	-	29.75	-
Ma et al. (2020)	24.87	-	23.55	-	30.09	-
Zheng et al. (2020)	25.10	-	24.91	-	30.40	-
Bao et al. (2021)	25.12 (+0.30)	27.17	25.52 (+0.33)	27.11	32.39 (+1.02)	34.08
Sentence Baseline	24.73	-	25.18	-	30.13	-
Finetune on Sentence	25.62 (+0.89)	29.47	25.73 (+0.55)	27.78	31.41 (+1.28)	33.50
Table 1: Experiments results of BLEU scores on three datasets. The improvement from the Transformer baseline for previous models are also reported as in "()". It indicates the score improved from sentence-level translation provided by their implementations. Results are averaged from two runs.
Data	# of Docs	# of Sents/Doc
TED	1.7K/93/23	123/98/105
News	6.1K/71/155	40/25/20
Europarl	118K/240/360	14/15/14
Table 2: Dataset Statistics for Train/Valid/Test
As suggested by Beltagy et al. (2020); Bao et al. (2021), context would be better applied in higher layers and keep only local information in lower layers. Therefore we only apply the memory unit M in the top layer of encoder and decoder, and in lower layers, we keep using the original Transformer structure. Analysis regarding the location of memory is discussed in Section 5.1.
Training During training, our model takes an input of  $\mathbf{x}_{\mathrm{t}}$  and  $\mathbf{y}_{\mathrm{t}}$ , which refer to sequences of tokens of the  $\mathfrak{t}^{\mathrm{th}}$  sentence in source and target documents. Memory unit M is initialized trainable parameters before the first input of each document, and it will be updated after each input sentence pair, which are batched as the sentence order in the document. For computational convenience, the gradients are only back-propagated to the current sentence and the most recent sentence in each update step, and we stop the gradient for M before it is passed to the next step.
Inference In the decoding phase, our model translates the source document sentence by sentence. In the generation of each sentence, tokens are decoded in an auto-regressive order until the stop sign or exceeds the max length. The memory M will not be updated until the complete sentence is generated since the update of M depends on all tokens in the current sentence. If M is updated
after each token generation, the attention space in the output attention does not represent the complete contextualized information of the expected sentence. The computational complexity keeps in sentence-level since we only feed one sentence every time, and there is no cache vector besides M.
We experiment across three widely used datasets for English  $\rightarrow$  German document translation.
TED Training data for TED comes from IWSLT'17. We use tst2016-2017 as test set and a held-out set from training as valid.
News The corpus comes from News Commentary v11. We use tst2016-2017 as test set and a held-out set from training as valid.
Europarl Train, valid and test sets are extracted from the corpus Europarl v7, as mentioned in (Maruf et al., 2019).
Detailed statistics for the datasets is in Table 2. Moses (Koehn et al., 2007) is used for data processing and BPE (Sennrich et al., 2016) is used with vocab-size of  $30\mathrm{K}$  for all datasets.
We adopt Transformer model with the transformer-base configurations as the baseline, which has six layers with a hidden size of 512 and an intermediate size of 2048. Token embedding is shared for source and target languages, and token indexes are encoded with a learnable embedding matrix. We first train a baseline model with vanilla Transformer architecture for sentence-level translation and finetune our model based on the sentence-level baseline. We use the AdamW optimizer with an ini
Figure 2: Evaluation on TED with different memory sizes.
tial learning rate of  $5 \times 10^{-4}$  and warm-up steps of 4000 for training sentence-level baseline. The dropout rate is set to 0.3 for TED and News and 0.1 for the Europarl. As for finetuning after sentence-level Transformer, the learning rate is set as  $3 \times 10^{-4}$  for newly initialized parameters and  $6 \times 10^{-5}$  for pretrained parameters, and warm-up steps of 1000 are set for TED and 2000 for News and Europarl. The drop-out rate is set to 0.1 for the Europarl and 0.2 for the TED and News during finetuning. We also apply gradient accumulation, and detailed studies are discussed in the section 5.3. Models are trained with a patience of 5 for both sentence-level and document-level. We use the beam size of 5 during inference and compute the BLEU score in a max order of 4 after removing BPE-tokens. s-BLEU and d-BLEU are used as evaluation metrics, where s-BLEU refers to the BLEU score for sentences, and d-BLEU is the score for documents.
Experiment results are shown in Table 1. Our method shows consistent improvements over three datasets from sentence-level Transformer. We achieve the state-of-the-art results of s-BLEU of 25.62 and d-BLEU of 29.47 on TED and s-BLEU of 25.73 and d-BLEU of 27.78 on News. Though our results do not outperform G-Transformer (Bao et al., 2021) on Europarl, we think the difference mostly comes from the gap between sentence-level baselines. Such difference may be caused by the implementation framework and computing resources, which they use Fairseq library and multiple GPUs, while we adopt the code from HuggingFace and only a single 1080-Ti GPU is used for our training. We further report the score of works gained from their reported baseline, and our model makes the greatest improvement on all three
Side	Index	s-BLEU	d-BLEU
Source+Target	0-1	25.31	29.13
Source+Target	2-3	25.30	29.23
Source+Target	4-5	25.43	29.22
Source+Target	5	25.62	29.47
Source Only	5	25.42	29.33
Target Only	5	25.43	29.25
Table 3: Evaluation on TED with memory on different sides and layers. We adopt the 6-layer Transformer model finetued on the sentence-level baseline, and 0 refers to the first layer, and 5 refers to the last layer.
datasets. Overall, the results could demonstrate the advantages of our method in the general understanding of the document machine translation.