Efficient Long Sequence Encoding via Synchronization

Transformer-based encoders Vaswani et al. (2017) have been widely used in natural language processing with successes. The pre-trained language models based on Transformer, such as BERT Devlin et al. (2019), GPT-3 Brown et al. (2020), T5 Raffel et al. (2019) and BART Lewis et al. (2019), further make it a dominating architecture in NLP.

Despite its successes, the Transformer models suffer from a major challenge in encoding long sequences. It is due to the fact that the self-attention mechanism used in each Transformer layer requires to compute attention for each pair of input words. Such computations lead to $O(l^{2})$ complexity in time and space in each Transformer layer, where $l$ is the sequence length. This limits Transformer’s roles in the increasingly important long sequence encoding for two common scenarios: (1) encoding of a single long document with lengths exceeding the input limitation, and (2) joint-encoding of multiple related documents for tasks that require synthesizing scattered pieces of evidence, e.g., multi-hop reasoning and multi-document summarization. Figure 1 gives an example of the necessary information exchange in multi-hop QA. Each paragraph provides a partial clue to solve the task (shown as the connected entities). Intuitively, an effective global encoding should allow the entities (e.g., Dutch-Belgian) appearing in multiple paragraphs to share information across all their occurrences. In this way, the embedding of Dutch-Belgian in the second paragraph can be aware of partial evidence from the first one and resolve the required information of House of Anubis.

To overcome this difficulty, many techniques have been proposed. The existing solutions can be categorized into two classes. The first class is hierarchical encoding. The idea is to either explicitly split the input into multiple short segments for fast encoding of each segment and then exchange information on top of their embeddings following sample-agnostic strategies Ainslie et al. (2020); Wang et al. (2020); or implicitly constrain the information exchange among tokens with a sparse attention map Beltagy et al. (2020); Zaheer et al. (2020). The essential of these methods is to find efficient ways to pass information among segments and compensate for the loss of important global context across segments. One solution is introducing a pseudo token for each segment and encouraging the pseudo tokens to attend one another during encoding Ainslie et al. (2020) for inter-segment interactions. The second class is post-hoc aggregation in the generative framework, such as Fusion-in-Decoder (FiD, Izacard and Grave 2020) with BART model. The input is split into segments that are encoded independently by the encoder. Then the decoder casts global attention over all the segments and generates the prediction. This approach allows for only shallow information exchange because the encoding is purely localized; yet is proved empirically very powerful in many NLP tasks.

We propose an orthogonal direction towards an efficient encoding of long sequences. Our method starts with local segments in the post-hoc aggregation approaches and relies on our proposed synchronization mechanisms to swap useful information from the other relevant segments during encoding, so as to maintain global information. Formally, our approach first identifies a set of anchors in the segments and puts them into different groups based on the similarity in their semantic units or the roles they play in the original input sequence. The identified anchors and groups connect different segments logically and naturally. Our synchronization is applied only to the encoding stage where inside each Transformer layer of the encoder, we perform an additional embedding update for each anchor using other anchor embeddings in the same group after a normal local encoding. The local encoding and anchor synchronization happen iteratively so that the global information is propagated deeply among segments with anchors as bridges.

Compared to previous hierarchical encoding approaches with fixed communication designs, our approach is more powerful and flexible. First, our approach provides a finer-grained information exchange mechanism. Second, our approach is a general framework that reduces the problem of global encoding to synchronization schema design. For any new applications or tasks, it is easy to infuse human prior knowledge to the model by identifying task-specific anchors and anchor grouping.

We evaluate our approach on two different tasks that require encoding of long sequences: NarrativeQA Kočiskỳ et al. (2018), where each input is a single long story summary; and a wild multi-hop QA task adapted from HotpotQA Yang et al. (2018), where the evidence annotation is assumed unavailable and the input documents are treated more independently from each other. The two settings correspond to the representative examples of the aforementioned long sequence encoding scenarios (1) & (2). Results show that our approach significantly improves the performance while remaining efficient. Moreover, building on top of FiD, the state-of-the-art hierarchical method ETC Ainslie et al. (2020) does not bring further improvement as we observe but our approach improves consistently.

In this section, we propose our TranSync   framework which extends Transformer layer with an embedding synchronization module attached to the end. Given a long context sequence $C$, we divide it into segments, i.e. $C=[s_{1};s_{2};...;s_{n}]$ where $s_{i}$ is the $i$-th segment of $C$ and $n$ is the number of segments. A segment can represent a natural sentence or a sequence in a certain length. Together with the question $q$, we re-organize the input to the Transformer and form a set of question-prefixed segments $\{s_{i}^{q}\}_{i=1}^{n}$, s.t.

$\displaystyle s_{i}^{q}=[q;\text{<SEP>};s_{i}]$

(1)

where <SEP> is a special token. An embedding layer converts the text segments $\{s_{i}^{q}\}_{i=1}^{n}$ into their corresponding question-aware embeddings $\{\mathbf{e}_{i}^{q}\}_{i=1}^{n}$, s.t.

$\displaystyle\mathbf{e}_{i}^{q}=[\mathbf{t}_{i}^{1};\mathbf{t}_{i}^{2};...;\mathbf{t}_{i}^{l_{i}}]\in\mathbb{R}^{l_{i}\times d}$

(2)

where $l_{i}$ is the length of $s_{i}^{q}$’s token sequence, $d$ is the dimension of the feature vector and $\mathbf{t}_{i}^{j}\in\mathbb{R}^{d}$ is the embedding for the $j$-th token in the $i$-th segment.

For genericity, our synchronization is performed between the target anchor and the incoming anchors, following the idea of message passing. The values of the target anchor embedding $\mathbf{a}_{t}$ are updated with the weighted sum of the incoming anchor embeddings and itself, i.e.,

$\displaystyle\mathbf{a}_{t}^{\prime}=\sum_{k}\alpha_{k}\mathbf{a}^{k},\quad s.t.\sum_{k}\alpha_{k}=1$

(3)

where $\{\mathbf{a}^{k}\}$ are the embedding spans¹¹1Some words may correspond to multiple tokens due to the byte pair encoding (BPE) algorithm. of the same length within the same anchor group; $\alpha_{k}$ is the normalized weight. In this work, for each anchor group, we form a new sequence from the selected anchor embeddings, i.e., $[\mathbf{a}^{1};\mathbf{a}^{2};...;\mathbf{a}^{k}]$, and use a self-attention module to compute the weights and update the embedding values.

Our TranSync  framework is embedded in Transformer layer. The synchronization is performed between the local self-attention and normalization steps to achieve deep information exchanging. At the end of the last Transformer layer, the synchronized segment embeddings are fused into one by concatenating one another as follows:

$\displaystyle[\mathbf{e}_{1}^{q};\mathbf{e}_{2}^{q};...;\mathbf{e}_{n}^{q}]\in\mathbb{R}^{\sum_{1}^{n}l_{i}\times d}$

(4)

The flexibility of our TranSync  framework is granted by the manifold strategies of identifying anchors in the segments and the heterogeneous message passing directions. The schema will be detailed in Section 3.

In this section, we introduce two experiments performed to verify the feasibility and flexibility of our TranSync  framework.

In this work, we propose TranSync  framework with flexible synchronization mechanisms for encoding long sequences. We demonstrate the feasibility of our method in reasoning tasks with long context, and also show its high adaptability to different scenarios. We consider our work to be valuable as an easy solution to address the long context issue in QA, and to be potentially applicable to other long sequence modeling tasks.