Efficient Transformers with Dynamic Token Pooling (2023)

Abstract: Transformers achieve unrivalled performance in modelling language, but remain inefficient in terms of memory and time complexity. A possible 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 boundaries in an autoregressive fashion. We compare several methods to infer boundaries, including end-to-end learning through stochastic re-parameterisation, supervised learning (based on segmentations from subword tokenizers or spikes in conditional entropy), as well as linguistically motivated boundaries. We perform character-level evaluation on texts from multiple datasets and morphologically diverse languages. 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.

## introduction the transformer architecture @xcite lies at the heart of cutting-edge generative models, such as gpt-3 @xcite ) for text and dall•e 2 @xcite for images. its success can be largely attributed to the ability to leverage a considerable amount of data, which yields performance gains @xcite and emergent abilities @xcite in accordance with well-established scaling laws. nonetheless, the time and memory efficiency of transformers remains constrained by their algorithmic complexity of o(l 2 n), 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 selfattention can be sparsified @xcite @xcite or linearly approximated @xcite . hourglass transformers @xcite provide an alternative solution, where the sequence length is reduced in the intermediate layers by merging fixed-size groups of tokens, similar to @xcite . these pooled representations are up-sampled back to the original length in order to generate sequences in an auto-regressive fashion @xcite . nevertheless, pooling groups of fixed size is suboptimal 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 sequence may carry different degrees of information (for instance, silence and voice in speech). ideally, the model should perform hierarchical computation, relying on the same abstractions as human processing of language, and conditional, by allocating resources to sub-sequences in proportion to the model uncertainty. in this work, we demonstrate that dynamic pooling results not only in higher shortening rates of input sequences, and thus increased efficiency, but also superior performance in next token prediction due to adopting the correct inductive bias in grouping tokens. to this end, we propose a new transformer variant that jointly learns token sequences and dynamically pools them into latent groupings of variable size (figure 1 ). crucially, the segmentation must preserve the auto-regressive property, and typical subword tokenizers cannot be applied to incomplete sequences during generation. rather, we learn a neural boundary predictor during training: 1) supervised by tokenizers such as unigram @xcite ; 2) supervised by spikes in the conditional entropy of the predictive distribution, which ensure that the computation is adaptive to the level of uncertainty of the sequence model; 3) end-to-end through stochastic re-parameterisation @xcite ; 4) use natural data boundaries such as whitespaces, which separate words in many scripts, without a predictor. to validate our model, we experiment with character-level language modelling of text in several english benchmarks, including text8 (mahoney, 2006), cc-100 @xcite , and wiki40b @xcite , as well as in a series of languages representing different morphological types: finnish, hebrew, and vietnamese. we find that dynamic pooling not only achieves lower time and memory complexity, but even surpasses the performance of vanilla transformers and fixed-size pooling transformers in most benchmarks by statistically significant margins. overall, our results indicate a promising direction to further accelerate training and therefore facilitate scaling. a faq section about our methods, findings, and the experimental setup is available in appendix a. we release the code at https:// github.com/piotrnawrot/dynamic-pooling . ## results the results for the experiments on character-level language modelling are shown in figure 5 shows how higher sf translates to lower training time and memory consumption in practice, as measured on a common gpu with an optimised model implementation. ## conclusions we proposed a new family of language models that pool variable-size segments of tokens in the intermediate layers in order to enhance the efficiency and performance of the transformer architecture. in particular, we learn a boundary predictor either end-to-end through stochastic re-parameterisation, through supervision (obtained from subword tokenization or spikes in the conditional entropy), or based on linguistic boundaries such as words. we evaluate this model extensively on multiple language modelling benchmarks in english and in other typologically diverse languages: finnish, hebrew, and vietnamese. compared to vanilla transformers and fixed pooling, we observe a significant decrease in model perplexity as well as time and space complexity. this opens up the perspective to develop transformer models capable of computing language both hierarchically, with the same abstractions humans perform at different levels of linguistic structure, and conditionally on the information content of each segment. in the future, our dynamic-pooling transformer can be combined with methods relying on external memory @xcite , encoders operating at a fine resolution @xcite , and more generally any task with long-context inputs @xcite . this may further facili-tate the scalability of current language modelling architectures.