SlowGuess/NAACL_2022_OCR huggingface dataset:
Haozhe Ji, Rongsheng Zhang, Zhenyu Yang, Zhipeng Hu, Minlie Huang
The CoAI group, DCST, Institute for Artificial Intelligence, State Key Lab of Intelligent Technology and Systems, Beijing National Research Center for Information Science and Technology, Tsinghua University, Beijing 100084, China
Fuxi AI Lab, NetEase Inc., China, Guangdong OPPO Mobile Telecommunications Corp., Ltd.
jhz20@mails.tsinghua.edu.cn, aihuang@tsinghua.edu.cn
{zhangrongsheng, zphu}@corp.netease.com, yangzhenyu@oppo.com
Although Transformers with fully connected self-attention are powerful to model long-term dependencies, they are struggling to scale to long texts with thousands of words in language modeling. One of the solutions is to equip the model with a recurrence memory. However, existing approaches directly reuse hidden states from the previous segment that encodes contexts in a uni-directional way. As a result, this prohibits the memory to dynamically interact with the current context that provides up-to-date information for token prediction. To remedy this issue, we propose Look-Ahead Memory (LaMemo)<sup>1</sup> that enhances the recurrence memory by incrementally attending to the right-side tokens, and interpolating with the old memory states to maintain long-term information in the history. LaMemo embraces bi-directional attention and segment recurrence with an additional computation overhead only linearly proportional to the memory length. Experiments on widely used language modeling benchmarks demonstrate its superiority over the baselines equipped with different types of memory.<sup>2</sup>
Language modeling is an important task that tests the ability of modeling long-term dependencies by predicting the current token based on the previous context (Mikolov and Zweig, 2012; Merity et al., 2017). Recently, Transformer-based language models achieved remarkable performance by enabling direct interaction between long-distance word pairs. However, as the computation overhead grows with the length of the input sequence, Transformers can only process a fixed length segment at a time. To allow long-term information flow across individual segments, existing approaches augment the model
Figure 1: Attention weights on the context (in log-scale) in the final layer of Transformer-XL and LaMemo averaged on 15K tokens. Transformer-XL quickly loses attention to older contexts, while LaMemo maintains awareness to the history with the grow of the context length.
with a recurrence memory that stores hidden states computed in previous time steps (Dai et al., 2019) and their compressions (Rae et al., 2020; Martins et al., 2021) for the target tokens to attend to.
One limitation of this approach is that the recurrence memory is only aware of older contexts since they are previously computed to predict the next word from left to right. As a result, distant memory states become outdated and less activated by the current context, as illustrated in Figure 1. When humans read or write a document, they maintain a memory that records important information from the past and often refresh them under the current context to keep it up-to-date.
In this paper, we propose Look-Ahead Memory (LaMemo) where memory states "look ahead" to future time steps by attending to the token representations on their right side to provide up-to-date contextualization. To maintain information from the long-term history, we propose memory interpolation to take both past and future tokens into consideration, which mimics the bi-directional attention. Note that, directly applying bi-directional attention to update the memory representations brings an additional complexity of  $\mathcal{O}(M^2)$  ( $M$  is the mem
<sup>3</sup>Note that the look-ahead attention does not exceed the current step of the autoregressive model to prevent information leakage.
<sup>1</sup>We are also inspired by the French word “La Mémoire”, meaning “the memory”.
Source code available at https://github.com/ thu-coai/LaMemo.
Human Language Technologies, pages 5747 - 5762
ory length). This is expensive when the memory is very long. LaMemo incrementally attends to the right and accumulate the weighted attention sum from previous segments to simulate the full attention in only  $\mathcal{O}(M\times N)$  complexity ( $N$  is the target sequence length), which does not increase the attention complexity of Transformer-XL, namely  $\mathcal{O}(N^2 +M\times N)$ . We provide an illustration of this mechanism in Figure 3.
Another technique proved to be effective in language modeling is the relative positional encoding (Shaw et al., 2018; Huang et al., 2018; Dai et al., 2019), which biases the pair-wise attention score purely based on the relative distance of the two tokens. However its ability to generalize to the attention of the future tokens remains unknown, since both the distance and the direction need to be taken into consideration. In preliminary experiments, we observed the unstability of directly applying the relative positional encoding of Dai et al. (2019) to this setting. We propose a simple yet effective modification based on Dai et al. (2019) that disentangles the bias of the relative distance and the attention direction which facilitates the training of LaMemo. We give both theoretical and empirical analysis to the unstability issue and demonstrate the effectiveness of the proposed disentangled relative positional encoding method.
(1) We propose LaMemo, a memory mechanism that incrementally attends to the right-side tokens, and interpolates with the old memory, which enables bi-directional interaction with a complexity linear in memory length.
(2) We propose disentangled relative positional encoding, a simple yet effective solution that disentangles the relative distance and the attention direction that can better generalize to the attention of the future tokens.
(3) We conduct experiments on standard language modeling benchmarks and demonstrate LaMemo's superiority over various baselines equipped with different types of memory mechanisms, despite some having an access to longer contexts. Comprehensive comparisons show the benefits of learning memory representations contextualized with up-to-date information.
2.1 Transformer for Language Modeling
A Transformer (Vaswani et al., 2017) is composed of multiple layers of identical blocks, including a multi-head self-attention (Bahdanau et al., 2015) that calculates pair-wise token interaction and a feed-forward layer for position-wise projection with a non-linear activation. Both two modules are followed by residual connections (He et al., 2016) and layer normalization (Ba et al., 2016) to facilitate optimization.
Given the input sequence representations of the current  $\tau$ -th segment  $X_{\tau} = [x_{\tau +1},\dots ,x_{\tau +N}]\in$ $\mathbb{R}^{N\times d}$  where  $N$  is the target sequence length and  $d$  is the hidden state size, they are first mapped into queries  $Q$ , keys  $K$  and values  $V$  by learned weight matrix to compute self-attention:
$$
\boldsymbol {Q} _ {\tau} = \boldsymbol {X} _ {\tau} \boldsymbol {W} _ {q}, \boldsymbol {K} _ {\tau} = \boldsymbol {X} _ {\tau} \boldsymbol {W} _ {k}, \boldsymbol {V} _ {\tau} = \boldsymbol {X} _ {\tau} \boldsymbol {W} _ {v}, \tag {1}
$$
where  $W_{q}, W_{k}, W_{v} \in \mathbb{R}^{d \times d}$  are learnable projection matrices. To perform multi-head self-attention,  $Q, K, V$  are further split into  $H$  heads. For simplicity, we only consider the case of a single head. In language modeling, the attention map is always added by a causal mask to avoid information leakage from the future when predicting the next token:
$$
\begin{array}{l} \boldsymbol {C} _ {\tau} ^ {\rightarrow} = \operatorname {C a u s a l - A t t n} \left(\boldsymbol {Q} _ {\tau}, \boldsymbol {K} _ {\tau}, \boldsymbol {V} _ {\tau}\right) \\ = \operatorname {s o f t m a x} _ {\triangle} \left(\frac {\boldsymbol {Q} _ {\tau} \boldsymbol {K} _ {\tau} ^ {\top}}{\sqrt {d}}\right) \boldsymbol {V} _ {\tau}, \tag {2} \\ \end{array}
$$
where  $\mathrm{softmax}_{\triangleq}(\cdot)$  masks position  $j > i$  for the  $i$ -th row of the input matrix with  $-\infty$  before taking the softmax. The resulted context representations are concatenated and then projected to the final outputs  $O_{\tau} \in \mathbb{R}^{N \times d}$  with a learnable projection matrix  $W_{o} \in \mathbb{R}^{d \times d}$ . Finally, the self-attention outputs  $O_{\tau}$  are added by the input representations  $X_{\tau}$  and fed to the following point-wise non-linear transformation, denoted as  $f(\cdot)$ :
$$
f (\boldsymbol {x}) = \operatorname {L N} \left(\operatorname {F F N} \left(\operatorname {L N} (\boldsymbol {x})\right) + \operatorname {L N} (\boldsymbol {x})\right), \tag {3}
$$
where  $\mathrm{LN}(\cdot)$  is the layer normalization and  $\mathrm{FFN}(\cdot)$  is the feed-forward layer, both of which are applied to each row vector individually. The final output of this Transformer layer is  $f(O_{\tau} + X_{\tau})$
Outputs of the final layer are projected to the vocabulary to predict  $\operatorname{Pr}(w_t | w_1, \dots, w_{t-1})$ . The joint probability of predicting the whole segment
Figure 2: The architecture of Transformer-XL augmenting with a recurrence memory.
is the product of these conditional factors. The final objective is to maximize the following log-likelihood:
$$
\log \Pr (\boldsymbol {w}) = \prod_ {t = 1} ^ {N} \log \Pr \left(w _ {t} \mid w _ {1}, \dots , w _ {t - 1}\right). \tag {4}
$$
2.2 Recurrence Memory Mechanism
To enable the Transformer to consider more contextual information from previous segments, Dai et al. (2019) proposed to augment the Transformer with a recurrence memory which stores the hidden states of previous time steps as extended keys and values, as shown in Figure 2. Concretely, let us consider a memory length of  $M$  and memory representations  $\mathbf{X}_{\tau -1} = [\pmb{x}_{\tau -M + 1},\dots ,\pmb{x}_{\tau}]\in \mathbb{R}^{M\times d}$ . The extended key and value matrices are obtained by preprocess  $\mathbf{X}_{\tau -1}$  to  $\mathbf{X}_{\tau}$  before projection:
$$
\tilde {\boldsymbol {X}} _ {\tau} ^ {\mathrm {s g}} = \left[ \operatorname {s g} \left(\boldsymbol {X} _ {\tau - 1}\right) \circ \boldsymbol {X} _ {\tau} \right] \in \mathbb {R} ^ {(M + N) \times d}, \tag {5}
$$
where  $\mathrm{sg}(\cdot)$  stands for stop-gradient which disables gradient propagation to previous segments, and  $[\cdot \circ \cdot ]$  indicates concatenation of hidden states along the length dimension. Extended by the recurrence memory, each query vector can consider contexts even beyond the total context length of the attention  $M + N$ . As illustrated by Dai et al. (2019), the effective context length grows linearly to the number of layers and the attention context length due to layer-wise reusing.
Another technique necessary to the recurrence memory is the relative positional encodings. By considering only the relative distance between two tokens when computing the attention score, it avoids temporal confusion caused by indexing the same position across segments and injects useful relative bias. Transformer-XL uses the fixed sinusoidal encoding matrix (Vaswani et al., 2017) to provide relative distance bias and learns global
Figure 3: Illustration of LaMemo with a memory length  $M = 2$  and a target sequence length  $N = 1$  for clarity. Solid lines stand for the attention connections computed at this iteration while dashed lines represent the previously computed attention.
bias terms shared across different layers, which can extrapolate to longer contexts with a great reduction of parameters compared to Shaw et al. (2018):
$$
\begin{array}{l} \boldsymbol {A} _ {i, j} ^ {x l} = \boldsymbol {X} _ {i} ^ {\top} \boldsymbol {W} _ {q} ^ {\top} \boldsymbol {W} _ {k} ^ {E} \boldsymbol {X} _ {j} + \boldsymbol {X} _ {i} ^ {\top} \boldsymbol {W} _ {q} ^ {\top} \boldsymbol {W} _ {k} ^ {R} \boldsymbol {R} _ {i - j} \\ + \boldsymbol {u} ^ {\top} \boldsymbol {W} _ {k} ^ {E} \boldsymbol {X} _ {j} + \boldsymbol {v} ^ {\top} \boldsymbol {W} _ {k} ^ {R} \boldsymbol {R} _ {i - j}, \tag {6} \\ \end{array}
$$
where  $\pmb{R}$  is the sinusoid encoding matrix,  $\pmb{u},\pmb{v}$  are learnable weight vectors governing the global content and position bias, and  $\pmb{W}_k^E,\pmb{W}_k^R$  are separate key projection matrices for the content and position respectively.
In this section, we describe our method in detail with our motivation to learn better representations for the memory.
3.1 Look-Ahead Attention
Human language is sequential with one word following another, but humans process information usually in a non-sequential way and recontextualize certain contents for several times. For example, when countering complicated contents during reading, humans usually first store them temporarily in the memory and continue to scan for relevant information if any, and revisit those old contents to refresh their meaning quite often. This dynamic memory refreshing mechanism enables us to thoroughly understand the passage under current contexts.
Existing recurrence memory however, lacks this dynamic contextualization ability. As the representations in the recurrence memory are previously computed conditioned on their past, they are not aware of the current contexts which provide more
relevant information for the current token prediction.
To address this limitation, we propose a look-ahead attention that allows the memory to attend to the contexts on their right. Formally, we reuse the notation  $\mathbf{X}_{\tau} = [\pmb{x}_{\tau + 1}, \dots, \pmb{x}_{\tau + N}] \in \mathbb{R}^{N \times d}$  for the representations of the current target sequence and  $\mathbf{X}_{\tau - 1} = [\pmb{x}_{\tau - M + 1}, \dots, \pmb{x}_{\tau}] \in \mathbb{R}^{M \times d}$  for the representations of the memory.
Let us consider the  $i$ -th position of the memory  $\mathbf{X}_{\tau -1}$ ,  $\pmb{x}_i$  can attend to position  $\pmb{x}_j$  on its right  $(j > i)$  without causing information leakage as long as  $j\leq \tau +1$ . Though appealing, this naive approach requires to calculate an  $M$  by  $M$  attention map, which would become inefficient and redundant when  $M$  is significantly greater than  $N$ . Actually, since the target segment moves forward  $N$  positions at each iteration, we devise an incremental manner of look-ahead attention computation that only requires the newest  $N$  positions on the right as key-value pairs.
$$
\tilde {\boldsymbol {X}} _ {\tau - 1} = [ \boldsymbol {x} _ {\tau - N + 2}, \dots , \boldsymbol {x} _ {\tau + 1} ] \in \mathbb {R} ^ {N \times d}. (7)
$$
Then the look-ahead attention results computed previously can be effectively reused and interpolated with the current ones (§3.2). Concretely, we formalize the look-ahead attention as follows:
$$
\tilde {\boldsymbol {K}} _ {\tau - 1} = \tilde {\boldsymbol {X}} _ {\tau - 1} \boldsymbol {W} _ {k}, \tilde {\boldsymbol {V}} _ {\tau - 1} = \tilde {\boldsymbol {X}} _ {\tau - 1} \boldsymbol {W} _ {v}, \tag {8}
$$
$$
\begin{array}{l} \boldsymbol {C} _ {\tau - 1} ^ {\leftarrow} = \mathrm {L o o k A h e a d - A t t n} (\boldsymbol {Q} _ {\tau - 1}, \tilde {\boldsymbol {K}} _ {\tau - 1}, \tilde {\boldsymbol {V}} _ {\tau - 1}) \\ = \operatorname {s o f t m a x} _ {\triangledown} \left(\frac {\boldsymbol {Q} _ {\tau - 1} \tilde {\boldsymbol {K}} _ {\tau - 1} ^ {\top}}{\sqrt {d}}\right) \tilde {\boldsymbol {V}} _ {\tau - 1}, \tag {9} \\ \end{array}
$$
where  $\mathrm{softmax}_{\nabla}(\cdot)$  masks position  $j \leq i$  for the  $i$ -th row of the input matrix with  $-\infty$  before softmax.  $Q_{\tau - 1}$  is obtained by Eq. (1), and the projection matrices of query, key and value are all shared with the causal attention. We illustrate this in Figure 3 where the look-ahead attention (yello paths) increases the attention window of each memory state to  $M$  tokens on its right.
3.2 Memory Interpolation
To save computations for looking-ahead and effectively reuse the attention results of the past, we propose memory interpolation that smoothly interpolates attention results from both the future and the past to provide bi-directional contextualization.
Recall that in the previous iteration, we have calculated the causal context representations  $C_{\tau - 1}^{\rightarrow}$  of  $X_{\tau - 1}$  using Eq. 2, where each row is a linear
Figure 4: The architecture of LaMemo with look-ahead attention and memory interpolation that refresh the memory dynamically with both the current contexts and the long-term history.
combination of the weighted token representations of the previous tokens. In Sec. 3.1, we describe the look-ahead attention which enables  $X_{\tau - 1}$  to attend to the contexts on their right and computes  $C_{\tau - 1}^{\leftarrow}$  using Eq. 9. Here, we formulate the memory interpolation as the interpolation between the old representations  $C_{\tau - 1}^{\rightarrow}$  and the new ones  $C_{\tau - 1}^{\leftarrow}$  with a coefficient vector  $\alpha_{\tau - 1} \in \mathbb{R}^{M}$  controlling the memorization of the past activations:
$$
\begin{array}{l} C _ {\tau - 1} ^ {\leftrightarrow} = \operatorname {M e m - I n t e r p} \left(C _ {\tau - 1} ^ {\rightarrow}, C _ {\tau - 1} ^ {\leftarrow}, \alpha_ {\tau - 1}\right) \\ = \boldsymbol {\alpha} _ {\tau - 1} \operatorname {s g} \left(\boldsymbol {C} _ {\tau - 1} ^ {\rightarrow}\right) + (1 - \boldsymbol {\alpha} _ {\tau - 1}) \boldsymbol {C} _ {\tau - 1} ^ {\leftarrow}. \tag {10} \\ \end{array}
$$
The resulted  $C_{\tau -1}^{\leftrightarrow}$  which attend to contexts from both directions, are further fed to the non-linear transformation defined in Eq. 3 to update representations in higher layers.
For  $\alpha_{\tau -1}$ , we define it to be the sum of the normalized attention weights on the previous tokens when calculating  $C_{\tau -1}^{\rightarrow}$  (Eq. 2):
$$
\boldsymbol {\alpha} _ {\tau - 1} = \frac {\operatorname {s g} \left(\boldsymbol {s} _ {\tau - 1} ^ {\rightarrow}\right)}{\operatorname {s g} \left(\boldsymbol {s} _ {\tau - 1} ^ {\rightarrow}\right) + \boldsymbol {s} _ {\tau - 1} ^ {\leftarrow} + \varepsilon}, \tag {11}
$$
where  $s_{\tau - 1}^{\rightarrow}$  is the sum of the unnormalized attention score of  $C_{\tau - 1}^{\rightarrow}$ , which is the denominator of the softmax in Eq. 2. Similarly,  $s_{\tau - 1}^{\leftarrow}$  is the denominator of the softmax in Eq. 9.  $\varepsilon$  is a small value to prevent zero division error in practice. Then Eq. 10 can be derived into a form that resembles the bi-directional attention with the queries attending to positions on both sides (Appendix A). Figure 4 shows the architecture of LaMemo.
4Note that the query vectors for the past and the future are under different contextualization in higher layers of the model.
Note that the difference between the hidden state reuse in the recurrence memory and our memory interpolation is that they simply reuse the static representations to extend the contexts for attention while we update the memory representations by aggregating weighted attention sum of the history without the need to recompute them.
3.3 Disentangled Relative Positional Encodings
As the look-ahead attention allows the memory to attend to future tokens on its right, we need a relative positional encoding scheme that can generalize to this setting. We start by considering the relative positional encoding in Transformer-XL, as described by Eq. 6. When the  $i$ -th query vector attending to a position  $j = i + \Delta > i$ , we have  $\pmb{R}_{i-j} = \pmb{R}_{-\Delta}$ . As defined by Vaswani et al. (2017),  $\pmb{R}_{\Delta} \in \mathbb{R}^{D}$  is composed of sine and cosine functions with different frequencies. Since the sine function is odd,  $\sin(-\omega \Delta) = -\sin(\omega \Delta)$ , we have  $\pmb{R}_{-\Delta} \neq \pmb{R}_{\Delta}$  so that it can represent attention in different directions ( $\pm$  sign of  $\Delta$ ) with the same relative distance (absolute value of  $\Delta$ ).
However, this approach solely relies on the fixed sinusoid encodings to represent the relative distance and the attention direction. We argue that disentangling them is more effective in capturing these two types of temporal biases and also mitigates the numerical unstability issue. Specifically, we propose to learn two direction-aware global position biases to parameterize the sign and query  $\mathbf{R}$  with the absolute value of the relative distance:
$$
\begin{array}{l} \boldsymbol {A} _ {i, j} ^ {d i s} = \boldsymbol {X} _ {i} ^ {\top} \boldsymbol {W} _ {q} ^ {\top} \boldsymbol {W} _ {k} ^ {E} \boldsymbol {X} _ {j} + \boldsymbol {X} _ {i} ^ {\top} \boldsymbol {W} _ {q} ^ {\top} \boldsymbol {W} _ {k} ^ {R} \boldsymbol {R} _ {| i - j |} \\ + \boldsymbol {u} ^ {\top} \boldsymbol {W} _ {k} ^ {E} \boldsymbol {X} _ {j} + \boldsymbol {v} _ {i - j} ^ {\top} \boldsymbol {W} _ {k} ^ {R} \boldsymbol {R} _ {| i - j |}, \tag {12} \\ \end{array}
$$
where  $\pmb{v}_{i-j} = \pmb{v}_+$  if  $i \geq j$  else  $\pmb{v}_-$ . The global positional bias now explicitly separates the contributions of  $\operatorname{sgn}(i - j)$  and  $|i - j|$ , which can better generalize to long distance in both forward and backward directions.
To illustrate the numerical unstability caused by adapting Eq. 6 to  $j > i$ , we derive the variance of the dot product  $\boldsymbol{x}^T\boldsymbol{R}_{i-j}$  where  $\boldsymbol{x}$  is a random vector. We show that the variance undergoes an oscillation and cannot be properly bounded everywhere when  $i$  shifts from  $i \geq j$  to  $i < j$ . Detailed analysis are presented in Appendix B.
We evaluate LaMemo on both word-level and character-level language modeling tasks and compare with existing Transformer baselines augmented with different types of memory.
For word-level language modeling task, we consider Wikitext-103 (Merit et al., 2017), which is the most widely used word-level language modeling benchmark. It contains 103 million tokens for training from 28 thousand wikipedia articles, with an average length of 3.6 thousand tokens per article and a vocabulary size around 260K. We report perplexity (ppl) on the dev and test set.
We also evaluate on two character-level language modeling benchmarks enwik8 and text8 (Mahoney, 2011). Both datasets contain 100 million Wikipedia characters. While enwik8 is unprocessed, text8 is preprocessed by case lowering and filtering to include only 26 letters from a to z and space. On both datasets, we report bit per character (bpc) on the dev and test set.
To directly compare with different types of memory, we consider Transformer-XL and its variations with the same model architecture but different memory mechanism.
Transformer+RPE is the vanilla Transformer (Vaswani et al., 2017) that uses relative positional encodings from Dai et al. (2019) but does not extend the context with additional memory.
Transformer-XL (Dai et al., 2019) is a Transformer model equipped with relative positional encodings and a recurrence memory comprised of hidden states computed in previous time steps to extend the context length of the attention.
Compressive Transformer (Rae et al., 2020) extends Transformer-XL with an external compressive memory that stores compressed hidden states at the temporal level using convolutional networks.
$\infty$ -former (Martins et al., 2021) uses continuous space attention to attend over the external memory which consists of continuous signals. They also updated the external memory with recent hidden states to enable unbounded memory capacity.
Model	#Params	Mem size	Ext mem size	#FLOPS	dev ppl	test ppl
Transformer+RPE	151M	0	0	148M	28.11	29.14
Transformer-XL (Dai et al., 2019)	151M	150	0	157M	23.42	24.56
Compressive Transformer (Rae et al., 2020)	161M	150	150	169M	-	24.41
∞-former (Martins et al., 2021)	160M	150	150	235M	-	24.22
LaMemo	151M	150	0	191M	22.98	23.77