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

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 10251–10269 July 9-14, 2023 ©2023 Association for Computational Linguistics LAIT: Efficient Multi-Segment Encoding in Transformers with Layer-Adjustable Interaction Jeremiah Milbauer1,2∗, Annie Louis1, Mohammad Javad Hosseini1, Alex Fabrikant1, Donald Metzler1, Tal Schuster1 1Google Research 2Carnegie Mellon University Abstract Transformer encoders contextualize token rep- resentations 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 re- lated segments (e.g., the sequence of sen- tences 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 in- dependently, and then jointly. This par- tial 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 pre- trained Transformers and converts them into the hybrid of the two aforementioned architec- tures, allowing for easy and intuitive control over the performance-efficiency tradeoff. Ex- perimenting on a wide range of NLP tasks, we find LAIT able to reduce 30-50% of the at- tention FLOPs on many tasks, while preserv- ing high accuracy; in some practical settings, LAIT could reduce actual latency by orders of magnitude. 1 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 seg- ments — like natural language inference (NLI), fact verification, question answering (QA), seman- tic similarity (STS), etc. — the common setting is to concatenate and jointly process all tokenized ∗ Work done as an intern at Google Research. EVIDENCE: L EVIDENCE: all geeks are mortal all geeks are humans all humans are mortal CLAIM: Fully parallel encoding: L all geeks are mortal all geeks are humans all humans are mortal CLAIM: EVIDENCE: EVIDENCE: Fully self-attentive encoding: CLAIM: P L P - all geeks are mortal all geeks are humans all humans are mortal EVIDENCE: EVIDENCE: Layer-adjustable interaction: Figure 1: A comparison of three approaches to multi- segment modeling for an arbitrary claim verification task. a) Fully-self attentive architecture, with each to- ken attending to each other token over L layers. b) Gen- eralized dual encoder, with each segment encoded sep- arately by an L-layer Transformer and representations concatenated. c) Layer-adjustable interactions (ours), with N layers of independent segment encoding and L − P layers of fully self-attentive segment encoding. segments as input to a neural model, most often some form of bidirectional Transformer-based ar- chitecture (Vaswani et al., 2017). In this setting, the self-attention blocks of the Transformer layers contextualize the per-token representations …

Content truncated to 500 words