Efficient Transformers with Dynamic Token Pooling
Year: 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.

heading: 1 Introduction
text: The Transformer architecture (Vaswani et al., 2017) lies at the heart of cutting-edge generative models, 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 complexity 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 selfattention can be sparsified (Child et al., 2019; Roy
et al., 2021; Ren et al., 2021) or linearly approximated (Beltagy et al., 2020). Hourglass Transformers (Nawrot et al., 2022) provide an alternative solution, 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 original length in order to generate sequences in an auto-regressive fashion (Ronneberger et al., 2015).
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 (Kudo, 2018); 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 (Maddison et al., 2017; Jang et al., 2017); 4) use natural data
6403
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 (Wenzek et al., 2020), and wiki40b (Guo et al., 2020), 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.

heading: 2 Background

heading: 2.1 Language Modelling with Transformers
text: Let x = (x1, . . . , xl) denote the input sequence. A language model assigns a probability value to any possible sequence of tokens from a vocabulary V . The parameters of a model θ are optimised to maximise the aggregate probability of all x ∈ V∗ in the training set D:
argmax θ
∑
x∈D
l∑
t=1
log p(xt | x<t,θ), (1)
where t indexes time steps. In our experiments, θ corresponds to the parameters of an autoregressive Transformer model (Vaswani et al., 2017).
A key advantage of Transformers is their ability to scale, which ultimately reaps the largest benefits according to (Sutton, 2019)’s ‘bitter lesson’ and reveals surprising emergent capabilities of language models (Kaplan et al., 2020; Wei et al., 2022). Nevertheless, the algorithmic complexity of selfattention, O(l2) where l is the length of the sequence, creates a bottleneck. To alleviate this cost, previous work (Clark et al., 2022; Tay et al., 2022; Nawrot et al., 2022) proposed to reduce the sequence length after the initial layers by pooling together groups of tokens. A single shortening by a factor k reduces the complexity to O( l2
k2 ). This
allows for increasing either the model efficiency or its depth within the same compute budget.

heading: 2.2 Hourglass Transformer
text: Naïve length reduction through pooling would reduce the length of output, however language models operate with the same input and output resolutions. For this reason, (Nawrot et al., 2022) introduced the Hourglass Transformer composed of three blocks of Transformer layers, which downsample, process, and upsample the tokens back to the original granularity. The first block encodes each input token xt into ht. Afterwards, groups of adjacent tokens of fixed length k are mean-pooled
to form ⌈ lk⌉ representations s:
sm = 1
k
mk∑
i=mk−k+1 hi (2)
Next, each pooled representation sm is processed by the middle block of Transformer layers, which operates with complexity O( l2
k2 ), yielding s′m. This
sequence is up-sampled to its original resolution by duplication: ut = s′⌈ t−k+1
k ⌉, and added to the
hidden representations h from before shortening through a skip connection, and passed to the third block.
Note that we subtract k− 1 from the index. This is because pooling and up-sampling in an autoregressive model pose a risk of data leakage from the future to the past. In fact, up-sampled representations might encompass future tokens if no measures are taken to prevent this. As a remedy, Hourglass Transformer shifts the up-sampled sequence to the right by k−1 positions, and pads it with a learnable null-group representation u0 at the beginning. This is sufficient to satisfy the autoregressive property in the fixed pooling scenario.1
Hourglass Transformer was shown to improve time and space complexity in a number of language and image modelling tasks, for a given parameter count. However, this came at the expense of degrading the perplexity of the language model, especially with shortening factors k > 2. We conjecture that this undesirable side effect is due to two main reasons. Firstly, the distribution of lengths of natural units of meaning such as morphemes and phrases in natural languages is uneven: for instance, word length is correlated with its frequency (Zipf, 1949; Bentz and Ferrer-i Cancho, 2016). Secondly, information content tends to be distributed uniformly across units of meaning (Meister et al., 2021).
As a consequence, fixed pooling creates segments with incongruous boundaries and unequal information content. For instance, in speech, this results in giving silence and voice the same importance. Instead, an ideal model should allocate compute conditionally on the information content of a given token. This would also ultimately lead to interpreting language hierarchically based on the same abstractions that humans adopt for language processing. Hence, we present a method to enable variable-length pooling and up-sampling in autoregressive language models.
1We refer to (Nawrot et al., 2022) for more details.

heading: 3 Dynamic-Pooling Transformer

heading: 3.1 Boundary Prediction
text: In order to augment the Hourglass architecture with variable-size pooling, we seek to find a sequence of segment boundaries b ∈ {0, 1}l for every input x. Let bt = 1 denote a segment boundary between elements xt and xt+1. The boundary predictor is implemented as a Multi-Layer Perceptron with parameters ϕ. As shown in Figure 1, this module maps each representation ht encoded by the first stack of Transformer layers into a Bernoulli probability distribution:
b̂t = p(bt=1) = sigmoid (MLPϕ (ht)) . (3)
Since segment boundaries are discrete, sampling from this distribution is not differentiable with respect to the model perplexity. Hence, we optimise this latent variable through stochastic reparametrisation (Jang et al., 2017; Maddison et al., 2017) via hard Gumbel-sigmoid (Section 3.1.1), jointly learning the language model and boundary predictor. We favour this solution over a scorefunction estimator of the gradient, as it suffers from high variance and computation costs due to sampling (Schulman et al., 2015).
As an alternative, we explore training the boundary predictor module with a binary cross-entropy loss with respect to two different sources of supervision: a Unigram tokenizer (Section 3.1.2) and spikes in conditional entropy (Section 3.1.3). Finally, we consider resorting to linguistically inspired boundaries (Section 3.1.4). During training and evaluation, we perform maximum likelihood inference for these variables. In other words, each b̂t from Equation (3) is rounded to the closest binary scalar such that bt = ⌊b̂t⌉.

heading: 3.1.1 Segmenting with Gumbel-Sigmoid
text: In order to learn the input segmentation end-toend based on the model perplexity, we can reparameterise the Bernoulli distribution of Equation (3) by injecting stochasticity in this form:
b̂t = sigmoid  log b̂t u
(1− b̂t) (1− u)
1/τ  
u ∼ Uniform(0, 1). (4) where τ is the temperature, a hyper-parameter. This estimator, however, is biased and might lead to sub-optimal results. As a consequence, we also propose methods based on supervised learning of the boundary predictor in the following sections.

heading: 3.1.2 Segmenting with Subword Tokenizers
text: Widespread algorithms for extracting variablelength boundaries for text are subword tokenizers, including Unigram (Kudo, 2018), Byte Pair Encoding (BPE; Sennrich et al., 2016), and WordPiece (Schuster and Nakajima, 2012). However, these create subwords greedily, and might change the segmentation of a given sequence prefix after more tokens are observed. For instance, consider the phrase ‘civil aviation’. A Unigram model might segment its prefix ‘civil aviatio’ differently before and after observing the character ‘n’:
_civil _a vi ati o _civil _a vi ation
During training an entire sentence is tokenized, but during inference a prefix is extended one character at a time and re-tokenized, possibly changing the boundaries like in the example above. Hence, deploying off-the-shelf tokenizers naïvely during inference does not recover the oracle segments and creates a mismatch between training and evaluation boundaries.
As a remedy, we provide the training tokenization as supervision to our autoregressive boundary predictor instead. More specifically, we employ a Unigram tokenizer (Kudo, 2018), as it aligns with morphological units better than other algorithms (Bostrom and Durrett, 2020). To prevent subword units from crossing word boundaries, we split the text on whitespace characters beforehand. Vocabulary size is a tunable hyper-parameter which
offers different efficiency–performance trade-offs.

heading: 3.1.3 Segmenting with Entropy Spikes
text: As an alternative to providing supervision through Unigram, we also propose a new segmentation method based on spikes of conditional entropy, which is agnostic about the presence of natural boundaries (such as whitespaces) or the availability of tokenizers. These properties make it suitable for other modalities in addition to text, such as speech and vision. Moreover, this enables top-down supervision and end-to-end training without external dependencies.
Intuitively, in natural language the information content tends to be spread evenly throughout a sentence, to facilitate communication. The conditional entropy is the expectation of such information content over the tokens in the vocabulary: H(xt | x<t) = ∑
x∈V p(xt | x<t) (− log p(xt | x<t))︸ ︷︷ ︸
information content (5)
Therefore, peaks in this conditional entropy provide indications of surprisal, and can serve as natural boundaries between segments. More formally, let Ht be the conditional entropy at time t. We select local spikes by comparing their value within a (left) window of size k. We place boundaries according to the following conditions:
bt = { 1 if Ht > Hi ∀i ∈ {t− k, . . . , t− 1} 0 otherwise. (6) Empirically, entropy spikes in language models overlap with word boundaries to a significant degree (Hutchens and Alder, 1998). However, they are also more flexible as they enable conditional computation based on the model’s confidence about its next token prediction. As an example of segmentation based on entropy spikes, consider Figure 2.

heading: 3.1.4 Linguistically Inspired Segments
text: Finally, perhaps the most straightforward source of segmentation is word boundaries. In fact, in many scripts, these are marked by whitespace characters.2 The simplicity of this method of segmentation comes with the obvious drawback of not providing control over the rate of shortening, while we found that the optimal rate varies with the language. Hence its efficiency–performance trade-off is not tunable.
2Several scripts such as Chinese characters, however, do not adopt this convention.
Segment boundaries are placed in between two symbols. In our experiments, we put a boundary after a whitespace character. Thus, we do not need to train a boundary predictor, since predicting a whitespace character is a signal to close the group in the next iteration of auto-regressive generation. This would not be possible, had we chosen to put a boundary before a whitespace character.

heading: 3.2 Pooling and Up-sampling
text: In the pooling step (Figure 1) a generated sequence of boundaries b is used to pool the tokens belonging to the same segment by averaging. Thus, we form∑l
t=1 bt+1 shortened representations s, which are then passed to the middle block of Transformer layers. Note that for Gumbel-sigmoid, to keep pooling differentiable, we algebraically manipulate b ∈ Rl into B ∈ Rl×1+ ∑ t bt , i.e. a binary matrix that maps from the original length to the shortened length, following (Bhati et al., 2021). The cell Bij is 1 if token i is merged into the j-th group, and 0 otherwise. Thus, s = hB/ ∑ iBi⋆, where the denominator unit-normalises the matrix columns. To obtain the up-sampled representationut while preserving the autoregressive property, we calculate the largest index m so that the output of the middle block s′m does include future information: ut = s ′ m, where m = ∑t i=1 bi. As a consequence, a segment representation s′m can only be added to the last token pooled into group m. For all the other non-final tokens, we take the representation of a previous segment s′m−1. Similar to Hourglass, the representation for the first (null) group s0 is a learnable vector. Afterwards, ut is added to the highway layer representation ht.

heading: 3.3 Auxiliary Objectives
text: In addition to minimising the language modelling loss with respect to the parameters θ as shown in Equation (1), we use auxiliary objectives to train the boundary predictor parameters ϕ. For supervised learning with subword tokenizers and entropy spikes, we minimise the cross-entropy between predicted boundaries b and gold ones. For end-to-end learning with Gumbel softmax, we introduce a regularizer based on a Binomial prior. Let k = ∑ t bt:
Binomial(α; l, k) =
( l
k
) αk(1− α)l−k (7)
where α ∈ [0, 1] is a hyper-parameter. This regularizer prevents the model from collapsing into
trivially predicting each position as a boundary.

heading: 4 Experimental Setup

heading: 4.1 Datasets
text: In addition to English, we evaluate our model on data in three languages, which represent different morphological types: Finnish for agglutinative, Hebrew for introflexive, and Vietnamese for isolating. Thus, we ensure that dynamic pooling is robust to different word length distributions. For English, we use text8 (CC-BY-SA) (Mahoney, 2006), CC-100 (MIT) (Conneau et al., 2020) and wiki40b (CCBY-SA) (Guo et al., 2020) as they are established benchmarks for character-level language models. For the rest of the languages, we use the corresponding subsets of wiki40b. To make results comparable across languages and prevent data imbalance, we limit the size of CC-100 and wiki40b to the first 400M tokens of the training set and the first 2M tokens of the validation set. We retain the original splits for each dataset.
For all datasets and languages, we follow the same pre-processing steps of (Mahoney, 2006) for creating text8. Specifically, for each language we keep only the characters from its script, as well as whitespace and an end-of-line. The text is lowercased, and the digits are spelt out in the target language. For wiki40b, we also remove special structural markers and normalise homoglyphs. Finally, for Hebrew we also remove diacritics as they are not required to understand the text. This way, we filter out excerpts in different languages, which are known to contaminate noisy multilingual texts (Kreutzer et al., 2022). The pre-processing scripts can be found as part of our code.

heading: 4.2 Models
text: All of our experiments, except for the scaling ablation, use 12-layer Hourglass Transformers with 2 layers in the first block, 8 layers in the second block which operates on shortened sequences, and 2 layers in the final block, following (Nawrot et al., 2022). For every Transformer layer, the hidden dimension is 512, the intermediate feed-forward dimension is 2048. Self-attention is split into 8 heads. We use a post-norm architecture, GELU activation function (Hendrycks and Gimpel, 2016) in feed-forward layers and the relative attention parametrisation from Transformer XL (Dai et al., 2019). In total, the model has ~41M parameters.
The boundary predictor is a 2-layer MLP that
takes a hidden state as input and outputs a scalar at every time step. For models with dynamic pooling, this module adds around 1M additional parameters. We use the SentencePiece (Kudo and Richardson, 2018) library to train Unigram segmentation for every dataset separately. We detect spikes in conditional entropy according to a window of size k = 2, which we select from range k=1 . . . 4 for optimal BPC on text8. For Gumbel Sigmoid, we set the prior probability of a boundary α to 0.2 for English, Vietnamese and Hebrew, and 0.37 for Finnish. The Gumbel temperature parameter was set to 0.5 in all experiments. For Unigram vocabulary size, we set |V| = 10000 for English and Vietnamese and |V| = 200 for Finnish and Hebrew. We list training hyper-parameters in Appendix B.

heading: 5 Results
text: The results for the experiments on character-level language modelling are shown in Table 1. In addition to the four proposed segmentation methods, we include a vanilla Transformer and fixed-size pooling Transformers with multiple shortening factors as baselines. Every model is evaluated with respect to two metrics: bits per character (BPC; ↓) and shortening factor (SF; ↑). The former measures the negative log-probability of the language model predictions, and thus its quality; the latter measures the average reduction of the sequence length in intermediate layers, and thus the model efficiency. 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.
Segmentation Methods In all the English evaluation benchmarks (text8, wiki40b, and CC-100), both whitespace-based and Unigram-based segmentations achieve the lowest BPC, outperforming both vanilla and fixed-pooling Transformers by statistically significant margins.3 Moreover, the same two methods achieve the highest degrees of shortening. Note that for equivalent SFs, fixed-size pooling becomes detrimental to performance. The approaches based on entropy spikes and GumbelSigmoid are generally inferior to the alternatives for dynamic pooling. However, for comparable shortening factors, they always outperform vanilla and fixed-pooling Hourglass models. Moreover, they make the fewest assumptions about the data and the
3We indicate with a ⋆ wherever this is the case according to a Paired Student’s t-test with p < 0.05.
availability of external supervision, so they might be appropriate for other domains (such as speech and vision) in future work. In general, providing a Transformer with the correct inductive bias for pooling variable-size segments not only facilitates scaling but also enhances prediction quality.
Notably, the gains resulting from whitespace segmentation are not identical in all languages, due to their inherent differences in morphological types and average word length. Shortening Factors for this method range from 3.8× in introflexive Hebrew, to 7.9× in agglutinative Finnish, whereas isolating Vietnamese and mildly fusional English lie in between with 4.4× and 5.7×, respectively. The larger SFs of dynamic pooling methods translate into higher training speed, from 1.7× for Unigram in Hebrew to over 2.5× for whitespaces in English, while simultaneously lowering BPC. Overall, the gains from dynamic pooling are robust cross-lingually, but the optimal segmentation method may vary.
Efficiency–Performance Pareto Front While both low BPC and high SF are desirable, there exists a trade-off between them which is specific to each boundary prediction method. Hence, the ideal model should strike the right balance to improve in both respects simultaneously. Intuitively, vocab-
ulary size in Unigram and the prior α in GumbelSigmoid provide easily controllable knobs to study this interaction: as they change, so does the shortening factor. In Figure 3, we plot BPC and SF for six vocabulary sizes (200, 500, 1k, 3k, 5k, 10k) and five α values (0.20, 0.25, 0.30, 0.37, 0.45) and compare them with fixed-size pooling in Hourglass Transformers. Manifestly, dynamic pooling enhances the Pareto front by finding more optimal trade-offs between efficiency and performance. Moreover, while fixed pooling follows a similar trend cross-lingually, dynamic pooling behaves more idiosyncratically: e.g. BPC in Vietnamese and English surprisingly improves with higher SFs. During our study of the Efficiency–Performance Pareto Front, we noticed that the Gumbel-Sigmoid pooling approach exhibits greater instability compared to the Unigram-based pooling method. This can be observed through artifacts such as the spikes in BPC for Hebrew, depicted in Figure 3.
Time and Space Complexity To capture the concrete gains in efficiency of models with higher SFs, we have measured the memory consumption and training time of our PyTorch implementation of text8 models on a typical GPU (NVIDIA GV100 32GB). The results in Figure 5 apply to dynamicpooling (Gumbel, Whitespace, Unigram, and Entropy), fixed-pooling, and vanilla Transformers (only for SF=1). Note that these results are identical for both fixed-pooling and dynamic-pooling Hourglass for the same SF as the cost of the boundary predictor is negligible. With a shortening factor SF = 2, the model reduces both memory consumption and training time by over 40%, compared to a vanilla Transformer. At SF = 4, where
dynamic-pooling Hourglass still achieves superior BPC scores, resource consumption is reduced between 50% and 60% and training is 2.5× faster. This allows models to increase in size with the same compute budget (which depends on the hardware), while simultaneously benefiting their performance.
Scaling the Model We investigate if dynamicpooling Transformers scale well in terms of model size, by adding more layers in the middle block (Figure 4). We focus on this block as it increases the model depth (and hence its capacity) while retaining a higher efficiency due to operating on shortened sequences. We find that the gains from dynamic pooling are consistent across all numbers of layers. Extrapolating from the trends, dynamic pooling holds promise to continue providing benefits even in extremely large language models.
Average-pooling vs Sub-sampling As an ablation, we also compare two different methods to represent groups of tokens when shortening the input sequence length: average pooling, used in our experiments, and sub-sampling, i.e. selecting only the last token as a representative for each group. As it emerges from Table 2, average pooling yields superior performance in all models, including both fixed and dynamic pooling Transformers.
Other Efficient Transformer Models Finally, we remark that our method differs from most efficient Transformer algorithms, which reduce the quadratic complexity of attention (Child et al., 2019; Lee-Thorp et al., 2022; Choromanski et al., 2021; Wang et al., 2020), as it focuses on length reduction. While previous efficient variants tend to trade quality for efficiency, we have shown that the
dynamic-pooling mechanism improves both simultaneously in our experiments. Moreover, Nawrot et al. (2022) has shown that combining both strategies yields further gains.

heading: 6 Related Work
text: Dynamic RNNs Our approach is inspired by variants of RNNs that process sequences at varying time scales by introducing a hierarchy of hidden units. For instance, RNNs that mimic speedreading by introducing hidden units that can skip over some input elements (Campos et al., 2018; Seo et al., 2018). Similarly, (Chung et al., 2017) discovers the latent hierarchy of an input sequence using a stack of LSTMs. Each layer is equipped with a binary gate responsible for hard boundary detection, where lower-level boundaries determine state updates made by higher-level layers. Whenever the detector ends a segment, its representation is fed to the upper layer.
Early slow- and fast-changing units were already
described by (Hihi and Bengio, 1995). Similarly, Clockwork RNN (Koutnik et al., 2014) introduces a hierarchy of hidden state units that make transitions at a set of different, fixed frequencies. Adaptive Computation Time networks perform a different amount of computation on each sequence item (Graves, 2016). Both ideas were combined in FastSlow RNNs (Mujika et al., 2017) which can choose a heavy or light transition between timesteps.
Pooling Transformer models While pooling blocks in Transformers are related to slowly varying units in RNNs, their operation is different. RNNs suffer from unreliable transport of information across long time spans. Units that act like skipconnections over time can help them to carry information (Krueger et al., 2017). In a Transformer network, a unit at time t can directly communicate with any other unit, including previous ones, and we find it important to confirm the benefits of dynamic pooling in Transformer models.
Perhaps the most similar approach to ours is Funnel Transformer (Dai et al., 2020) which uses a similar, hourglass-shaped Transformer architecture. After passing through the first block, the data is pooled at a fixed rate, processed by the deep middle Transformer block, and up-sampled for the last block. Canine (Clark et al., 2022) has a similar three-part architecture, and processes Unicode inputs, which are downsampled with Transformer and convolution layers. (Tay et al., 2022) implements gradient-based subword tokenization within a Transformer model, which learns dynamic groupings of tokens into fixed-size groups. In (Bai et al., 2021), sentence and paragraph boundaries were used as additional conditioning for the model.
Boundary Detection We investigate boundaries provided by an external model, derived directly from the data, or top-down from the model’s entropy. (Kreuk et al., 2020) shows a bottom-up approach to phoneme segmentation task combining contrastive learning (van den Oord et al., 2019) with a method for boundary detection based on dissimilarity between subsequent frames. It was later extended by (Bhati et al., 2021) to segment the sequence of speech frames dynamically. Recently, (Cuervo et al., 2022) introduced a hierarchical sequence processing model in which units in the upper layer operate on a dynamically shortened sequence, with the shortening guided by a boundary prediction model.
(Rocki et al., 2016) control the activity of LSTM gates with the model’s output cross-entropy. (Alpay et al., 2019) used a similar mechanism based on information content to guide the copying of individual activations in an LSTM network. Similarly, we employ the entropy of model predictions to choose where to insert boundaries.

heading: 7 Conclusions
text: 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 (Wu et al., 2022), encoders operating at a fine resolution (Xue et al., 2022; Tay et al., 2022), and more generally any task with long-context inputs (Shaham et al., 2022). This may further facili-
tate the scalability of current language modelling architectures.

heading: 8 Limitations
text: Linguistic variation Our results are highly dependent on the target language and its morphology. For example, word boundaries might seem like an obvious choice for dynamic segmentation, and in fact they achieve the best performance in English and Vietnamese. However, for some languages like agglutinative Finnish, whitespaces are less frequent, which is detrimental to model performance. Explicit word boundaries are not available for all scripts. For example, in Chinese characters, or in modalities other than text like speech or vision, there is no obvious equivalent to whitespaces. However, segmentation based on stochastic re-parameterisation, subword tokenizers and spikes in conditional entropy overcomes these limitations.
Contiguous segments In its current formulation, dynamic pooling only allows for merging contiguous segments of tokens in a sequence. However, this is not ideal for morphology types like Hebrew where morphemes are discontinuous: vowels are interspersed between consonant roots for inflection. Moreover, future works should consider higher levels of linguistic structure than words, such as dependency trees, for pooling. In this case, discontinuous segments may be necessary to handle non-projective syntactic dependencies.
Independent boundary decisions The decision to emit a boundary at time step t depends on previous boundaries only indirectly through the hidden representation of the first Transformer block, as this preserves the efficiency of the boundary predictor. Instead, a recurrent model could be explicitly conditioned on previous boundary decisions, which however would negatively affect the time complexity of the language model.
Work contribution of authors
The idea of training the models with pooling of variable-length segments was discussed among the authors while Jan Chorowski was at the University of Wrocław. Experiments were performed by Piotr Nawrot while he was employed in a research grant at the University of Wrocław, under the supervision of Adrian Łańcucki and Edoardo M. Ponti. The manuscript was written by Piotr Nawrot, Adrian Łańcucki and Edoardo M. Ponti.

heading: Acknowledgements
text: This work was supported in part by the UKRI Centre for Doctoral Training in Natural Language Processing, funded by the UKRI (grant EP/S022481/1) and the University of Edinburgh, School of Informatics and School of Philosophy, Psychology & Language Sciences; and the Polish National Science Center under the OPUS-18 2019/35/B/ST6/04379 grant.

heading: A Frequently Asked Questions
text: A.1 Pros and Cons of shortening methods
A.2 What is the ultimate segmentation method? While Whitespace offers the best performance in many cases, this is not always true even in the linguistic domain. In agglutinative languages (e.g., Finnish), words are longer than in English, which has a detrimental effect on the Whitespace method. For such languages, other dynamic methods that allow for controlling the shortening factor (SF), such as Unigram, are better suited. Moreover, languages with non-Latin scripts (like Chinese) may lack explicit whitespaces. For modalities different from text, such as speech and vision, Gumbel and Entropy are to be favoured as they do not assume the discreteness of the input sequence.
A.3 Why evaluating on language modelling rather than downstream tasks? Since we present a proof of concept for dynamic-pooling Transformers, we limit the experiments to language modelling because: 1) it is a foundational NLP task; 2) previous efficient Transformer variants were evaluated on similar benchmarks. Crucially, there is a strong correlation between performance in language modelling and downstream tasks.
A.4 How do you ensure that the results are reliable? Our code is based on the optimised, open-source implementation of Transformer-XL from NVIDIA (Apache 2.0 License), which reproduces the scores reported by (Dai et al., 2019). Our implementation of the fixed-pooling Hourglass Transformer model similarly reproduces the results from (Nawrot et al., 2022). We make our code publicly available, under the Apache 2.0 License, inheriting from the original source, to ensure the reproducibility of our results. Moreover, memory utilisation was measured by controlling resource allocation on GPUs (Figure 5) rather than through a naive nvidia-smi readout, as this would overestimate the reserved buffers.

heading: B Hyper-parameters
text: Following (Dai et al., 2019), we train for 2 · 105 steps with a batch size of 8 and a learning rate 2.5 · 10−4 on 2x NVIDIA RTX 3080. Each training run took from approximately 12h to 30h, depending on the configuration. We use a linear warm-up schedule for the first 4k steps, followed by a single-cycle cosine scheduler. We use an Adam optimiser with β1 = 0.9, β2 = 0.999 and ϵ = 1e−8, and clip the gradients at 0.25. We apply a 0.1 dropout rate in the attention matrix and feed-forward layers. Before every epoch, we cyclically shift the text stream, divide it into non-overlapping chunks of 2048, and shuffle. During the evaluation, to provide context to the model, we split the test set into partially overlapping sequences of size l = 2048 with a step size of 512 and calculate the model perplexity only over the last 512 tokens.
ACL 2023 Responsible NLP Checklist

heading: A For every submission:
text: 3 A1. Did you describe the limitations of your work?
8
A2. Did you discuss any potential risks of your work? Not applicable. We propose the generic modification of the Transformer architecture. Potential risks are the same as for any other Transformer based model.
3 A3. Do the abstract and introduction summarize the paper’s main claims? 1
7 A4. Have you used AI writing assistants when working on this paper? Left blank.
B 3 Did you use or create scientific artifacts? 4.1

heading: 3 B1. Did you cite the creators of artifacts you used?
text: 4.1

heading: 3 B2. Did you discuss the license or terms for use and / or distribution of any artifacts?
text: 4.1 A.4
B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? Not applicable. Dataset we used to train the model were released for research purposes.
7 B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? For our experiments we used CC-100 (MIT License) and Wikipedia (CC-BY-SA) datasets (text8 and wiki40b) that are publicly available datasets released for research purposes with a goal to give the community easier access to the information. We assumed that the authors of these datasets took all necessary steps to not allow for any undesirable situations such as leak of private informations.

heading: 3 B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and
text: linguistic phenomena, demographic groups represented, etc.? 4.1
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. 4.1 B
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.
C 3 Did you run computational experiments? 4, 5

heading: 3 C1. Did you report the number of parameters in the models used, the total computational budget
text: (e.g., GPU hours), and computing infrastructure used? 4.2 B

heading: 3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found
text: hyperparameter values? 4.2 B

heading: 3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary
text: statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? Figure 4 Table 1

heading: 3 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did
text: you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? 4.2 B We also report the libraries and settings in our repository.
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? No response.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? No response.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? No response.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? No response.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? No response.

heading: Future Work
text: Moreover, future works should consider higher levels of linguistic structure than words, such as dependency trees, for pooling. In this case, discontinuous segments may be necessary to handle non-projective syntactic dependencies. Independent boundary decisions The decision to emit a boundary at time step t depends on previous boundaries only indirectly through the hidden representation of the first Transformer block, as this preserves the efficiency of the boundary predictor. Instead, a recurrent model could be explicitly conditioned on previous boundary decisions, which however would negatively affect the time complexity of the language model. Work contribution of authors
The idea of training the models with pooling of variable-length segments was discussed among the authors while Jan Chorowski was at the University of Wrocław. Experiments were performed by Piotr Nawrot while he was employed in a research grant at the University of Wrocław, under the supervision of Adrian Łańcucki and Edoardo M. Ponti. The manuscript was written by Piotr Nawrot, Adrian Łańcucki and Edoardo M. Ponti