SlowGuess/ACL_2020_OCR huggingface dataset:
Thanh-Tung Nguyen\*, Xuan-Phi Nguyen\*, Shafiq Joty\*, Xiaoli Li†  
\*Nanyang Technological University  
\*Salesforce Research Asia  
\*Institute for Infocomm Research, A-STAR  
Singapore  
{ng0155ng@e.;nguyenxu002@e.;srjoty@}ntu.edu.sg  
xlli@i2r.a-star.edu.sg
We propose Differentiable Window, a new neural module and general purpose component for dynamic window selection. While universally applicable, we demonstrate a compelling use case of utilizing Differentiable Window to improve standard attention modules by enabling more focused attentions over the input regions. We propose two variants of Differentiable Window, and integrate them within the Transformer architecture in two novel ways. We evaluate our proposed approach on a myriad of NLP tasks, including machine translation, sentiment analysis, subject-verb agreement and language modeling. Our experimental results demonstrate consistent and sizable improvements across all tasks.
Computing relative importance across a series of inputs can be regarded as one of the important advances in modern deep learning research. This paradigm, commonly known as attention (Bahdanau et al., 2015), has demonstrated immense success across a wide spectrum of applications. To this end, learning to compute contextual representations (Vaswani et al., 2017), to point to the relevant part in the input (Vinyals et al., 2015), or to select windows or spans (Wang and Jiang, 2017) from sequences forms the crux of many modern deep neural architectures.
Despite aggressive advances in developing neural modules for computing relative relevance (Luong et al., 2015; Chiu and Raffel, 2018), there has been no general purpose solution for learning differentiable attention windows. While span selection-based pointer network models typically predict a start boundary and an end boundary (Wang and Jiang, 2017; Seo et al., 2017), these soft predictions generally reside at the last layer of the net
work and are softly optimized. To the best of our knowledge, there exists no general purpose component for learning differentiable windows within networks.
Although the practical advantages of learning differentiable windows are plenty, this paper focuses on improving attentions with differentiable windows. The key idea is to enable more focused attention, leveraging dynamic window selection for limiting (and guiding) the search space for the standard attention modules to work within. This can also be interpreted as performing a form of dynamic local attention.
We make several key technical contributions. First, we formulate the dynamic window selection problem as a problem of learning a discrete mask (i.e., binary values representing the window). By learning and composing left and right boundaries, we show that we are able to parameterize the (discrete) masking method. We then propose soft adaptations of the above mentioned, namely trainable soft masking and segment-based soft masking, which are differentiable approximations that can not only be easily optimized in an end-to-end fashion, but also inherit the desirable properties of discrete masking.
While these modules are task and model ag- nostic, we imbue the state-of-the-art Transformer (Vaswani et al., 2017) model with our differentiable window-based attention. To this end, we propose two further variants, i.e., multiplicative window attention and additive window attention for improving the Transformer model. Within the context of sequence transduction and self-attention based encoding, learning dynamic attention windows are beneficial because they can potentially eliminate noisy aggregation and alignment from large input sequences. On the other hand, it is good to note that hard attention (Xu et al., 2015b), which replaces the weight average of soft attention with a stochas
6589
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 6589-6599 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
tic sampling model, tries to achieve similar ends, albeit restricted to token-level selection. Hence, our proposed differentiable windows are more flexible and expressive compared to hard attentions.
We evaluate our Transformer model with differentiable window-based attention on a potpourri of NLP tasks, namely machine translation, sentiment analysis, language modeling, and subject-verb agreement. Extensive experimental results on these tasks demonstrate the effectiveness of our proposed method. Notably, on the English-German and English-French WMT'14 translation tasks, our method accomplishes improvements of 0.63 and 0.85 BLEU, respectively. On the Stanford Sentiment Treebank and IMDB sentiment analysis tasks, our approach achieves  $2.4\%$  and  $3.37\%$  improvements in accuracy, respectively. We further report improvements of  $0.92\%$  in accuracy and 2.13 points in perplexity on the subject-verb agreement and language modeling tasks, respectively. We make our code publicly available at https://ntunlpsg.github.io/project/dynamic-attention/.
The attention mechanism enables dynamic selection of relevant contextual representations with respect to a query representation. It has become a key module in most deep learning models for language and image processing tasks, especially in encoder-decoder models (Bahdanau et al., 2015; Luong et al., 2015; Xu et al., 2015a).
2.1 Transformer and Global Attention
The Transformer network (Vaswani et al., 2017) models the encoding and decoding processes using stacked self-attention and cross-attention (encoder-decoder attentions). Each attention layer uses a scaled multiplicative formulation defined as:
$$
\operatorname {s c o r e} (\boldsymbol {Q}, \boldsymbol {K}) = \frac {\left(\boldsymbol {Q} \boldsymbol {W} ^ {Q}\right) \left(\boldsymbol {K} \boldsymbol {W} ^ {K}\right) ^ {T}}{\sqrt {d}} \tag {1}
$$
$$
\operatorname {a t t} (Q, K, V) = \mathcal {S} (\operatorname {s c o r e} (Q, K)) (V W ^ {V}) \quad (2)
$$
where  $S(\mathbf{A})$  denotes the softmax operation over each row of matrix  $\mathbf{A}$ ,  $Q \in \mathbb{R}^{n_q \times d}$  is the matrix containing the  $n_q$  query vectors, and  $\mathbf{K}, \mathbf{V} \in \mathbb{R}^{n \times d}$  are the matrices containing the  $n$  key and value vectors respectively, with  $d$  being the number of vector dimensions;  $\mathbf{W}^Q, \mathbf{W}^K, \mathbf{W}^V \in \mathbb{R}^{d \times d}$
are the associated weights to perform linear transformations.
To encode a source sequence, the encoder applies self-attention, where  $Q$ ,  $K$  and  $V$  contain the same vectors coming from the output of the previous layer. In the decoder, each layer first applies masked self-attention over previous-layer states. The resulting vectors are then used as queries to compute cross-attentions over the encoder states. For cross-attention,  $Q$  comprises the decoder self-attention states while  $K$  and  $V$  contain the encoder states. The attention mechanism adopted in the Transformer is considered global since the attention context spans the entire sequence.
2.2 Windows in Attentions
In theory, given enough training data, global attention should be able to model dependencies between the query and the key vectors well. However, in practice we have access to only a limited amount of training data. Several recent studies suggest that incorporating more focused attention over important local regions in the input sequence as an explicit inductive bias could be more beneficial.
In particular, Shaw et al. (2018) show that adding relative positional biases to the attention scores (Eq. 1) increases BLEU scores in machine translation. Specifically, for each query  $\pmb{q}_i \in \pmb{Q}$  at position  $i$  and key  $k_j \in K$  at position  $j$ , a trainable vector  $\pmb{a}_{i,j} = \pmb{w}_{\text{max}(-\tau, \min(j - i, \tau))}$  is added to the key vector before the query-key dot product is performed. The window size  $\tau$  is chosen via tuning. Sperber et al. (2018) also consider local information by restricting self-attention to neighboring representations to improve long-sequence acoustic modeling. Although shown to be effective, their methods only apply to self-attention and not to cross-attention where the query vectors come from a different sequence.
That said, Luong et al. (2015) are the first to propose a Gaussian-based local attention for crossattention. At each decoding step  $t$ , their model approximates the source-side pivot position  $p_t$  as a function of the decoding state and the source sequence length. Then, local attention is achieved by multiplying the attention score with a confidence term derived from a  $\mathcal{N}(p_t, \sigma^2)$  distribution. The aligned pivot  $p_t$  and the variance  $\sigma^2$  (a hyperparameter) respectively represent the center and the size of the local window.
${}^{1}$  Initially,  $Q,K$  ,and  $V$  contain the token embeddings.
Meanwhile, Yang et al. (2018) improve the method of Luong et al. (2015) by assigning a soft window weight (a Gaussian bias) to obtain a flexible window span. Despite effective, the aligned pivot position in the source is determined only by the decoder state, while the encoder states are disregarded - these should arguably give more relevant information regarding the attention spans over the source sequence. Besides, the confidence for local attention span may not strictly follow a normal distribution, but rather vary dynamically depending on the relationship between the query and the key. Furthermore, the approach of Luong et al. (2015) is only applicable to cross-attention while the one of Yang et al. (2018) works better only for encoder self-attention as shown in their experiments.
Our proposed differentiable window approach to local attention addresses the above limitations of previous methods. Specifically, our methods are dynamic and applicable to encoder and decoder self-attention as well as cross-attention, without any functional constraints. They incorporate encoder states into the local window derivation. They are also invariant to sequence length, which removes the dependence on global features from the local context extraction process.
3 Dynamic Differentiable Window
Our proposed attention method works in two steps: (i) derive the attention span for each query vector to attend over, and (ii) compute the respective attention vector using the span. In this section, we present our approaches to step (i) by proposing trainable soft masking and segment-based soft masking. In the next section, we present our methods to compute the attention vectors. To give the necessary background to understand what can be expected from our method, we first present the discrete masking case.
3.1 Discrete Window Masking
In this context, we seek to dynamically derive a boolean mask vector for each query that will indicate the window in the key-sequence over which the query should attend. In other words, attentions are only activated on the consecutive positions where the mask vector element is 1, and the positions with 0 are canceled out. Let the query vector and the key-sequence be  $\pmb{q} \in \mathbb{R}^d$  and  $\pmb{K} = (k_1, k_2, \dots, k_n)$ , respectively. Formally, we define the local attention mask vector  $m_q \in \{0, 1\}^n$  for the query  $q$  as
Figure 1: Example of  $\phi, f$ , and  $\pmb{g}$  vectors and how the mask vector  $\pmb{m}_q$  can be derived for  $l_q = 3$  and  $r_q = 8$ .
follows.
$$
\boldsymbol {m} _ {q} ^ {i} = \left\{ \begin{array}{l l} 1, & \text {i f} l _ {q} \leq i \leq r _ {q} \\ 0, & \text {o t h e r w i s e} \end{array} \right. \tag {3}
$$
where  $l_{q}$  and  $r_q$  denote the left and right positional indices that form a discrete window  $[l_q,r_q]$  over which the query attends. As such, in the standard global attention,  $l_{q} = 1$  and  $r_q = n$  for all the query vectors, and in decoder self-attention,  $l_{q} = 1$  and  $r_q = t$  for the query vector at decoding step  $t$ . To facilitate the construction of  $\pmb{m}_{q}$ , we first define vectors  $\phi_k, f_k, g_k$  and matrix  $L_{n}$  with entries as:
$$
\phi_ {k} ^ {i} = \left\{ \begin{array}{l l} 1, & \text {i f} i = k \\ 0, & \text {o t h e r w i s e} \end{array} ; \quad \boldsymbol {f} _ {k} ^ {i} = \left\{ \begin{array}{l l} 1, & \text {i f} i \geq k \\ 0, & \text {o t h e r w i s e} \end{array} \right. \right.
$$
$$
\boldsymbol {g} _ {k} ^ {i} = \left\{ \begin{array}{l l} 1, & \text {i f} i \leq k \\ 0, & \text {o t h e r w i s e} \end{array} ; \quad \boldsymbol {L} _ {n} ^ {i, j} = \left\{ \begin{array}{l l} 1, & \text {i f} i \leq j \\ 0, & \text {o t h e r w i s e} \end{array} \right. \right. \tag {4}
$$
where  $\phi_k \in \{0,1\}^n$  denotes the one-hot representation for a boundary position  $k$  (from the left or right of a sequence), and  $\pmb{f}_k, \pmb{g}_k \in \{0,1\}^n$  are the 'rightward' mask vector and 'leftward' mask vector, respectively;  $L_n \in \{0,1\}^{n \times n}$  denotes a unit-value (1) upper-triangular matrix with  $i$  and  $j$  being the row and column indices respectively. Figure 1 visualizes how these entities appear. Specifically,  $\pmb{f}_k$  has entry values of 1's for position  $k$  and its right positions, while  $\pmb{g}_k$  has entry values of 1's for position  $k$  and its left positions. As such,  $\pmb{f}_k$  and  $\pmb{g}_k$  can be derived from  $\phi_k$  and  $L_n$  as follows.
$$
\boldsymbol {f} _ {k} = \phi_ {k} ^ {T} \boldsymbol {L} _ {n}; \quad \boldsymbol {g} _ {k} = \phi_ {k} ^ {T} \boldsymbol {L} _ {n} ^ {T} \tag {5}
$$
Note that  $\pmb{f}_k$  can be interpreted as the cumulative sum across  $\phi_k$ , while  $\pmb{g}_k$  as the inverse cumulative sum across  $\phi_k$ .
Given the above definitions, the mask vector  $\pmb{m}_q$  for a query  $\pmb{q}$  to attend over the window  $[l_q, r_q]$  in
6591
the key sequence such that  $1 \leq l_{q} \leq r_{q} \leq n$  can be achieved by:
$$
\boldsymbol {m} _ {q} = \boldsymbol {f} _ {l _ {q}} \odot \boldsymbol {g} _ {r _ {q}} = \left(\phi_ {l _ {q}} ^ {T} \boldsymbol {L} _ {n}\right) \odot \left(\phi_ {r _ {q}} ^ {T} \boldsymbol {L} _ {n} ^ {T}\right) \quad (6)
$$
where  $\odot$  denotes element-wise multiplication. As shown in Figure 1,  $m_q$  represents the intersection between  $f_{l_q}$  and  $g_{r_q}$ , and forms a masking span for the attention.
3.2 Trainable Soft Masking
The above masking method is non-differentiable as  $\phi$  is discrete, which makes it unsuitable in an end-to-end neural architecture. In our trainable soft masking method, we approximate the discrete one-hot vector  $\phi$  with a pointing mechanism (Vinyals et al., 2015). Specifically, given the query  $q$  and the key-sequence  $K$  as before, we define confidence vectors  $\hat{\phi}_{l_q},\hat{\phi}_{r_q}\in \mathbb{R}^n$  as follows.
$$
\hat {\phi} _ {l _ {q}} = \mathcal {S} \left(\frac {\boldsymbol {q} ^ {T} \boldsymbol {W} _ {L} ^ {Q} (\boldsymbol {K} \boldsymbol {W} _ {L} ^ {K}) ^ {T}}{\sqrt {d}}\right) \tag {7}
$$
$$
\hat {\phi} _ {r _ {q}} = \mathcal {S} \left(\frac {\boldsymbol {q} ^ {T} \boldsymbol {W} _ {R} ^ {Q} (\boldsymbol {K} \boldsymbol {W} _ {R} ^ {K}) ^ {T}}{\sqrt {d}}\right) \tag {8}
$$
where  $\mathcal{S}$  is the softmax function as defined before, and  $W_{L}^{Q}, W_{L}^{K}, W_{R}^{Q}, W_{R}^{K} \in \mathbb{R}^{d \times d}$  are trainable parameters. Eq. 7-8 approximate the left and right boundary positions of the mask vector for the query  $q$ . However, contrary to the discrete case, they do not enforce absolute cancellation or activation of attention weights on any position in the key-sequence. Instead, they assign a confidence score to each position. This allows the model to gradually correct itself from invalid assignments. Moreover, the softmax operations enable differentiability while maintaining the gradient flow in an end-to-end neural architecture.
Note however that the left and right boundary concepts have now become ambiguous since the positions  $l_{q} = \arg \max (\hat{\phi}_{l_{q}})$  and  $r_q = \arg \max (\hat{\phi}_{r_q})$  are not guaranteed to conform to the constraint  $l_{q}\leq r_{q}$ . To understand its implication, lets first consider the discrete case in Eq. 6; the elementwise multiplication between  $\pmb{f}_{l_q}$  and  $\pmb{g}_{r_q}$  results in a zero vector for  $m_{q}$  if  $l_{q} > r_{q}$ , canceling out the attention scores entirely. Although not absolute zeros,
in the continuous case,  $\pmb{m}_q$  would potentially contain significantly small values, which renders the attention implausible. To address this, we compute the soft mask vector  $\hat{\pmb{m}}_q$  as follows.
$$
\hat {\boldsymbol {m}} _ {q} = \left(\hat {\phi} _ {l _ {q}} ^ {T} \boldsymbol {L} _ {n}\right) \odot \left(\hat {\phi} _ {r _ {q}} ^ {T} \boldsymbol {L} _ {n} ^ {T}\right) + \left(\hat {\phi} _ {r _ {q}} ^ {T} \boldsymbol {L} _ {n}\right) \odot \left(\hat {\phi} _ {l _ {q}} ^ {T} \boldsymbol {L} _ {n} ^ {T}\right) \tag {9}
$$
This formulation has two additive terms; the former constructs the mask vector when  $l_{q} \leq r_{q}$ , whereas the latter is activated when  $l_{q} > r_{q}$ . This ensures a non-zero result regardless of  $l_{q}$  and  $r_{q}$  values. It can be shown that the values in  $\hat{m}_q$  represent the expected value of the discrete flags in  $m_q$ , i.e.,  $\hat{m}_q = \mathbb{E}(m_q)$ ; see Appendix for a proof.
We concatenate the mask vectors horizontally for all the query vectors in  $Q \in \mathbb{R}^{m \times d}$  to get the mask matrix  $M \in \mathbb{R}^{m \times n}$ . Since the pointing mechanism is invariant to sequence length, the computation of the mask vectors enjoys the same advantages, enabling our models to efficiently perform attentions on any arbitrarily long sequences. In addition, the method is applicable to all attention scenarios – from decoder to encoder cross-attention, encoder self-attention, and decoder self-attention.
3.3 Segment-Based Soft Masking
The soft masking introduced above modulates the attention weight on each token separately which may result in unsmooth attention weights on neighbouring tokens. However, words in a sentence are related and they often appear in chunks or phrases, contributing to a shared meaning. Thus, it may be beneficial to assign identical mask values to the tokens within a segment so that they are equally treated in the window selection method. In this section, we propose a novel extension to our soft masking method that enables the mask vector to share the same masking values for the tokens within a segment in a key-sequence.
The main idea is to divide the key-sequence  $\mathbf{K} = (\mathbf{k}_1, \mathbf{k}_2, \dots, \mathbf{k}_n)$  into  $\lceil n / b \rceil$  consecutive segments and to assign the same masking value to the tokens in a segment. The segment size  $b$  is considered a hyper-parameter. We compute the segment-based mask vector  $\mathbf{m}_q'$  similarly as in Eq. 9, but with  $\mathbf{L}_n$  replaced by  $\mathbf{J}_n \in \mathbb{R}^{n \times n}$  defined as follows.
$$
\boldsymbol {J} _ {n} ^ {i, j} = \left\{\begin{array}{l l}1,&\text {i f} i \leq b \left\lceil \frac {j}{b} \right\rceil\\0,&\text {o t h e r w i s e}\end{array}\right. \tag {10}
$$
However, unlike the standard pointer network, in our case there is no direct supervision for learning the pointing function. Our network instead learns it from the end prediction task.
6592
Figure 2: Segment-based masking for segment size  $= 2$ . Instead of pointing to the left and right indices of the tokens, the soft segment-based method (approximately) points to the left and right boundaries of the segments, respectively.
$$
\boldsymbol {m} _ {q} ^ {\prime} = \left(\hat {\phi} _ {l _ {q}} ^ {T} \boldsymbol {J} _ {n}\right) \odot \left(\hat {\phi} _ {r _ {q}} ^ {T} \boldsymbol {J} _ {n} ^ {T}\right) + \left(\hat {\phi} _ {r _ {q}} ^ {T} \boldsymbol {J} _ {n}\right) \odot \left(\hat {\phi} _ {l _ {q}} ^ {T} \boldsymbol {J} _ {n} ^ {T}\right) \tag {11}
$$
Eq. 10 - 11 ensure that all the items in a segment share the same masking value, which is the cumulative sum of the confidence scores in  $\hat{\phi}_{l_q}$  and  $\hat{\phi}_{r_q}$ . For instance, suppose  $\hat{\phi}_{l_q} = (a_1, a_2, a_3, \ldots, a_n)$  and segment size  $b = 2$ , then the term  $\hat{\phi}_{l_q}^T J_n$  evaluates to  $(\sum_{i=1}^{2} a_i, \sum_{i=1}^{2} a_i, \sum_{i=1}^{4} a_i, \ldots)$ , and  $\hat{\phi}_{l_q}^T J_n^T$  evaluates to  $(\sum_{i=1}^{n} a_i, \sum_{i=1}^{n} a_i, \sum_{i=3}^{n} a_i, \ldots)$ . Similarly,  $\hat{\phi}_{r_q}^T J_n^T$  and  $\hat{\phi}_{r_q}^T J_n$  will have segment-level effects on the cumulative sums. Figure 2 visualizes the method with an example for  $b = 2$ .
One advantage of this approach is that it allows us to control the masking behavior (by varying  $b$ ) without increasing the number of parameters compared to the token-based masking. We also show its effectiveness in our experiments.
4 Dynamic Window Attention Methods
Having presented our method to compute the mask vector that defines the attention spans, we now present our methods to incorporate the mask vectors into the attention layers.
4.1 Multiplicative Window Attention
In this approach, the attention weights (Eq. 2) are (element-wise) multiplied by the mask matrix  $M$  to confine their attention scope defined by the mask. Formally, the attention scores and outputs are defined as follows.
$$
\begin{array}{l} \operatorname {s c o r e} = \frac {\left(\boldsymbol {Q} \boldsymbol {W} ^ {Q}\right) \left(\boldsymbol {K} \boldsymbol {W} ^ {K}\right) ^ {T}}{\sqrt {d}} (12) \\ \operatorname {a t t} _ {\mathrm {M W}} = \left(S (\text {s c o r e}) \odot M\right) \left(\boldsymbol {V} \boldsymbol {W} ^ {V}\right) (13) \\ \end{array}
$$
In this approach, the standard global attention weights are suppressed and partially overshadowed by the attention window imposed by  $M$ . Thus, it can be interpreted as a local attention method similar to Luong et al. (2015). However, instead of using a static Gaussian bias, we use a dynamic mask to modulate the attention weights.
4.2 Additive Window Attention
Having a local attention window could be beneficial, but it does not rule out the necessity of global attention, which has been shown effective in many applications (Vaswani et al., 2017; Devlin et al., 2019). Thus, we also propose an additive window attention, which implements a combination of global attention and local attention. The attention output in this method is formally defined as
$$
s _ {\mathrm {g l b}} = \left(\boldsymbol {Q} \boldsymbol {W} _ {\mathrm {g l b}} ^ {Q}\right) \left(\boldsymbol {Q} \boldsymbol {W} _ {\mathrm {g l b}} ^ {K}\right) ^ {T} \tag {14}
$$
$$
s _ {\text {l o c}} = \left(\boldsymbol {Q} \boldsymbol {W} _ {\text {l o c}} ^ {Q}\right) \left(\boldsymbol {Q} \boldsymbol {W} _ {\text {l o c}} ^ {K}\right) ^ {T} \odot \boldsymbol {M} \tag {15}
$$
$$
\operatorname {s c o r e} _ {\mathrm {A W}} = \frac {s _ {g l b} + s _ {l o c}}{\sqrt {d}} \tag {16}
$$
$$
\operatorname {a t t} _ {\mathrm {A W}} = \mathcal {S} \left(\operatorname {s c o r e} _ {\mathrm {A W}}\right) \left(\boldsymbol {V} \boldsymbol {W} ^ {V}\right) \tag {17}
$$
where  $W_{\mathrm{glb}}^{Q}, W_{\mathrm{glb}}^{K}, W_{\mathrm{loc}}^{Q}$ , and  $W_{\mathrm{loc}}^{K} \in \mathbb{R}^{d \times d}$  are the weight matrices for global and local attentions.
Compared to the multiplicative window attention where the mask re-evaluates the global attention weights, additive window attention applies the mask vector to the local attention scores  $(s_{\mathrm{loc}})$ , which is then added to the global attention scores  $(s_{\mathrm{glb}})$  before passing it through the softmax function. In this way, the mask-defined local window does not suppress the global context but rather complements it with a local context. Moreover, the resulting attention weights add up to one, which avoids attention weights diminishment that could occur in the multiplicative window attention. Additive merger of global and local window components may also facilitate more stable gradient flows.
4.3 Implementation in Transformer
We now describe how the proposed dynamic window attention methods can be integrated into the Transformer.