LAIT : Efficient Multi-Segment Encoding in Transformers with Layer-Adjustable Interaction (2023)

Abstract: Transformer encoders contextualize token representations by attending to all other tokens at each layer, leading to quadratic increase in compute effort with the input length. In practice, however, the input text of many NLP tasks can be seen as a sequence of related segments (e.g., the sequence of sentences within a passage, or the hypothesis and premise in NLI). While attending across these segments is highly beneficial for many tasks, we hypothesize that this interaction can be delayed until later encoding stages. To this end, we introduce Layer-Adjustable Interactions in Transformers (LAIT). Within LAIT, segmented inputs are first encoded independently, and then jointly. This partial two-tower architecture bridges the gap between a Dual Encoder’s ability to pre-compute representations for segments and a fully self-attentive Transformer’s capacity to model cross-segment attention. The LAIT framework effectively leverages existing pretrained Transformers and converts them into the hybrid of the two aforementioned architectures, allowing for easy and intuitive control over the performance-efficiency tradeoff. Experimenting on a wide range of NLP tasks, we find LAIT able to reduce 30-50% of the attention FLOPs on many tasks, while preserving high accuracy; in some practical settings, LAIT could reduce actual latency by orders of magnitude.

## introduction although the meaning of a sentence may depend on the context in which it appears, sentences still have meaning per se. however, in tasks involving reasoning across multiple sentences or text segments -like natural language inference (nli), fact verification, question answering (qa), semantic similarity (sts), etc. -the common setting is to concatenate and jointly process all tokenized segments as input to a neural model, most often some form of bidirectional transformer-based architecture @xcite . in this setting, the self-attention blocks of the transformer layers contextualize the per-token representations against all other input tokens, including those of different input segments. the potential for independent sentence-level semantics is largely ignored. while this practice has shown to achieve high accuracy, it is computationally expensive due to the quadratic increase in cost with the input length. and in practical settings, such as large-scale citation retrieval @xcite or documentlevel nli @xcite , where a given segment may occur multiple times, the full cartesian product of the sets of text segments must be processed, e.g., @xcite processes all sentence pairs from two wikipedia articles around one subject but in two different languages to identify potential discrepancies. this leads to yet another quadratic increase in cost. our goal is to reduce both of these computational burdens, rendering transformer architectures more efficient for large-scale multi-segment reasoning. in this paper, we present lait (/leit/), a late interaction transformer model with easy to implement layer-adjustable interactions. lait includes encoder layers that process each segment locally and independent of the other segments, followed by traditional transformer layers, in a simple but effective way. unlike the late interaction components of other models, such as colbert @xcite , which are specifically geared toward measuring a similarity score between two text segments, lait generally supports any sequence-to-sequence task and any number of input segments. lait enables several desirable properties for an efficient encoder: it (1) is easy to train on top of existing pretrained language models; (2) readily supports any seq-2-seq task, and any segmentation of the input; (3) improves the encoding efficiency by skipping a large number of attention computations; (4) disentangles independent segment representations from joint processing to allow caching of intermediate segment representations for repeated computations; and (5) provides an easy-to-tune hyperparameter for controlling the efficiency-performance tradeoff. ## layer-adjustable interactions we argue that both fsa and de transformer models can be seen as special cases of a general architecture with adjustable layer depths for both segment-independence and segment-interaction, which we will call a "layer-adjustable interaction transformer" (lait). for a transformer with l layers and an input with n segments, lait is a set of n independent stacks of p layers each, followed by l -p fully self-attentive encoder layers. any function can be used after the encoder. thus a typical fully self-attentive encoder-decoder transformer is a lait where @xmath0, and a shared-parameter dual encoder is a lait where @xmath1. in the fully self-attentive transformer, each token in each segment is interacting with each token in each other segment throughout the entire depth of the encoder; in a dual encoder, each segment is treated independently throughout the encoder. lait without caching the lait framework allows us to make the core questions of this work precise: (1) to what extent are interactions across multiple input text segments necessary? and (2) if they are not always necessary, how can we take advantage of this fact to perform multi-segment modeling efficiently at scale? specifically, given an input x with m tokens that is split into n segments s i . . . s n of possibly different lengths, the lait encoder is defined as: lait(s 1 , s 2 , . . . , s n ) = enc l-p ([enc p (s 1 ); enc p (s 2 ); . . . ; enc p (s n )]), where [x; y] denotes concatenating vectors x and y, and enc k (•) denotes a transformer encoder with k layers. the rule for splitting the input into segments r(x 1 , . . . , x m ) → s 1 , . . . , s n is predefined for each task, based either on prior knowledge of the input structure, or on a simple segmentation function. for example, in nli we can simply use the hypothesis and premise as two segments. in passage-level qa, we can use the question as one segment and the passage as another. however, splitting the passage into multiple shorter segments could help further reduce compute. for instance, we can split the passage by sentences to k segments, leading to a total of k + 1 segments. for @xmath2, l], lait interpolates between an n -encoder model and a fully-self attentive transformer. because interaction between segments is delayed, representations computed at layer p of the model can be stored or cached for later reuse as they are independently generated. figure 2 demonstrates the basic lait architecture, as well as possibilities for partial caching (for instance, multiple unique questions about the same passage), or full caching (for instance, nli-based cross-document reasoning @xcite ). similar to general text-to-text models, the outputs of the lait encoder, consisting of m contextualized representations for m tokens, are passed to the transformer-decoder for generating the output sequence. similarly, the decoder may be replaced with a classification head, or any other module. ## experimental setting below, we describe our evaluation setting, tasks, used metrics, and baselines. ## results to study the performance-efficiency tradeoff, we consider multiple configurations of lait to fully interpolate between a dual encoder and a fully self- as can be seen in tables 2 and 3 , which compare best validation-set performance across models, lait either matches, nearly-matches, or outperforms the t5-base baseline for every task. this holds even in configurations where cross-segment interaction is delayed to the last few layers of the encoder. as long as there are a few cross-segment interactions later in the model, performance remains relatively stable even as the architecture becomes increasingly efficient; crucially, lait can delay cross-segment interaction by 8-10 layers without a notable decrease in performance. we specifically focus on the most efficient lait models that: (1) achieve within 99% of lait-0 performance, which we call lait-99%; (2) achieve within 95% of lait-0 perfomrance, called lait-95%; and (3) achieve within the 95% confidence interval within lait-0 performance, called lait . to select these models with higher validity, we perform a synthetic dev/test split of the validation sets and report the held-out validation performance of the lait models with the highest performance on the synthetic dev set, reported in appendix b. these results also suggest differences in the proportion of cross-segment processing necessary for different tasks. sentence and word representation tasks (i.e., answer equivalence, stsb, and wic) have much better lait models than reasoningintensive tasks, such as mnli, boolq, and vita-minc. we note that fever appears to be easier for lait than other "reasoning" tasks, which we explore further in section 5.3. we also note that some degree of cross-segment processing is necessary for all tasks, evidenced by the steep drop in performance as p approaches 12 (see figure 4 ). ## conclusion we present layer-adjustable interactions in transformers (lait) to allow simple-but-effective efficiency gains over a wide range of nlp tasks. the lait framework leverages existing pretrained transformers such as t5, and converts them during finetuning into a hybrid model that combines parallel independent encoding of multiple segments, followed by fully-attentive layers to allow crosssegment reasoning. we evaluate lait on a large set of 10 wellknown datasets, involving different examined capabilities, number of segments, input lengths, output spaces, and difficulty levels. we find lait to consistently provide significant reduction in encoder attention complexity while preserving high accuracy. furthermore, we show that the parallel independent segment encoding of lait enables additional inference-time compute savings by caching representations of repeated segments in large scale real-world settings. lait demonstrates that transformers can achieve high performance even without crosssegment interaction at every layer; essentially, that sentences can be just as effectively encoded if first processed separately, and then processed jointly.