Efficient Transformers with Dynamic Token Pooling

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 6403–6417 July 9-14, 2023 ©2023 Association for Computational Linguistics Efficient Transformers with Dynamic Token Pooling Piotr Nawrot† Jan Chorowski‡ Adrian Ła´ncucki⋄♣ Edoardo M. Ponti† †University of Edinburgh ‡Pathway ⋄NVIDIA ♣University of Wrocław piotr.nawrot@ed.ac.uk Abstract Transformers achieve unrivalled performance in modelling language, but remain inefficient in terms of memory and time complexity. A possi- ble remedy is to reduce the sequence length in the intermediate layers by pooling fixed-length segments of tokens. Nevertheless, natural units of meaning, such as words or phrases, display varying sizes. To address this mismatch, we equip language models with a dynamic-pooling mechanism, which predicts segment bound- aries in an autoregressive fashion. We com- pare several methods to infer boundaries, in- cluding end-to-end learning through stochastic re-parameterisation, supervised learning (based on segmentations from subword tokenizers or spikes in conditional entropy), as well as lin- guistically motivated boundaries. We perform character-level evaluation on texts from multi- ple datasets and morphologically diverse lan- guages. The results demonstrate that dynamic pooling, which jointly segments and models language, is both faster and more accurate than vanilla Transformers and fixed-length pooling within the same computational budget. 1 Introduction The Transformer architecture (Vaswani et al., 2017) lies at the heart of cutting-edge generative mod- els, such as GPT-3 (Brown et al., 2020) for text and DALL·E 2 (Ramesh et al., 2022) for images. Its success can be largely attributed to the ability to leverage a considerable amount of data, which yields performance gains (Kaplan et al., 2020) and emergent abilities (Wei et al., 2022) in accordance with well-established scaling laws. Nonetheless, the time and memory efficiency of Transformers remains constrained by their algorithmic complex- ity of O(l2n), where l stands for sequence length and n for the number of layers. To remedy this shortcoming without renouncing the expressivity of a deep model, the quadratic self- attention can be sparsified (Child et al., 2019; Roy et al., 2021; Ren et al., 2021) or linearly approxi- mated (Beltagy et al., 2020). Hourglass Transform- ers (Nawrot et al., 2022) provide an alternative so- lution, where the sequence length is reduced in the intermediate layers by merging fixed-size groups of tokens, similar to (Dai et al., 2020). These pooled representations are up-sampled back to the orig- inal length in order to generate sequences in an auto-regressive fashion (Ronneberger et al., 2015). Nevertheless, pooling groups of fixed size is sub- optimal in several respects. First, these groups are misaligned with linguistic primitives: units of meaning such as morphemes, words, phrases, and clauses vary in size. Second, the elements of a se- quence may carry different degrees of information (for instance, silence and voice in speech). Ideally, the model should perform hierarchical computa- tion, relying on the same abstractions as human processing of language, and conditional, by allocat- ing resources to sub-sequences in proportion to the model uncertainty. In this work, we demonstrate that dynamic pooling results not only in higher shortening …

Content truncated to 500 words