SlowGuess/NAACL_2022_OCR huggingface dataset:
Piotr Nawrot\*, Szymon Tworkowski\*, Michał Tyrolski\*, Łukasz Kaiser\*, Yuhuai Wu\*, Christian Szegedy\*, Henryk Michalewski\*
<sup>1</sup>University of Warsaw, <sup>2</sup>OpenAI, <sup>3</sup>Google Research
{p.nawrot99, szy.tworkowski, michal.tyrolski, lukaszkaiser}@gmail.com, {yuhuai, szegedy, henrykm}@google.com
Transformer models yield impressive results on many NLP and sequence modeling tasks. Remarkably, Transformers can handle long sequences, which allows them to produce long coherent outputs: entire paragraphs produced by GPT-3 or well-structured images produced by DALL-E. These large language models are impressive but also very inefficient and costly, which limits their applications and accessibility. We postulate that having an explicit hierarchical architecture is the key to Transformers that efficiently handle long sequences. To verify this claim, we first study different ways to downsample and upsample activations in Transformers so as to make them hierarchical. We use the best performing upsampling and down-sampling layers to create Hourglass - a hierarchical Transformer language model. Hourglass improves upon the Transformer baseline given the same amount of computation and can yield the same results as Transformers more efficiently. In particular, Hourglass sets new state-of-the-art for Transformer models on the ImageNet32 generation task and improves language modeling efficiency on the widely studied enwik8 benchmark.
Transformer models (Vaswani et al., 2017) are capable of solving many sequence modeling tasks, including classical NLP tasks (Devlin et al., 2019), summarization (Zhang et al., 2020), language modeling (Radford et al., 2019; Brown et al., 2020), code generation (Chen et al., 2021), or even music generation (Huang et al., 2018; Dhariwal et al., 2020) and image generation (Parmar et al., 2018; Chen et al., 2020; Ramesh et al., 2021). One compelling feature of Transformers is their ability to handle long contexts given as part of the input. This is particularly visible in tasks where the output depends on parts of the context that may not be
close-by in the generated sequence, like in summarization, where the summary may need to refer to information scattered across the context, or in large-scale image generation, where pixels belonging to the same object may be far apart in the generation order. Transformers excel at such tasks thanks to self-attention, and they are used with longer and longer contexts.
Figure 1: Bits-per-character vs. training cost for baseline (orange) and hierarchical Transformers (green). We observe significant perplexity improvements on enwik8 over the vanilla Transformer-XL baseline, see text for details.
The ability of Transformers to handle long contexts comes at a price: each self-attention layer, at least in its original form, has complexity quadratic in the length of the context. When a stack of  $n$  Transformer layers is used, both memory and time complexity is equal to  $O(L^{2}n)$  where  $L$  is a sequence length and  $n$  number of decoder blocks. Due to this limitation, vanilla transformers are infeasible to train on tasks with very long input sequences, for instance, on high-resolution images. This issue has been studied extensively, and a number of techniques were introduced that modify attention mechanism without changing overall transformer architecture (Child et al., 2019; Roy et al., 2020; Ren et al., 2021). These sparse attention mechanisms reduce the complexity of self-attention
*Equal contribution. Order determined by coin toss.
Findings of the Association for Computational Linguistics: NAACL 2022, pages 1559 - 1571 July 10-15, 2022 ©2022 Association for Computational Linguistics
but still force the model to operate on the sequence of the same length as the input.
For generative Transformer models, operating at the original scale of the input sequence is necessary, at least in the early and final layers, as the input must be processed at first and generated at the end (Section 4.3). But forcing the models to operate at this granularity throughout the layer stack has both fundamental and practical shortcomings:
- Fundamentally, we aim for the models to create high-level representations of words, entities, or even whole events – which occur at a very different granularity than single letters that the model receives on input.
- On the practical side, even layers with linear complexity can be slow and memory-intensive when processing very long sequences.
To alleviate these issues, we propose to change the Transformer architecture to first shorten the internal sequence of activations when going deeper in the layer stack and then expand it back before generation. We merge tokens into groups using a shortening operation (Section 2.1) and so reduce the overall sequence length, and then up-sample them again combining with the sequence from earlier layers (Section 2.3). The first part is analogous to the Funnel-Transformer architecture (Dai et al., 2020), and the whole architecture takes inspiration from U-Nets (Ronneberger et al., 2015). In contrast to both these architectures, the model we present is autoregressive, which is harder to ensure in hierarchical models than in vanilla Transformers.
The resulting model - which we call Hourglass - is an autoregressive Transformer language model that operates on shortened sequences. It yields significant performance improvements for different attention types (Fig. 6,7). We tested Hourglass with Transformer-XL (Dai et al., 2019) and Reformer (Kitaev et al., 2020) blocks on enwik8 dataset. In both cases, it is not only better in terms of perplexity, but it is faster and uses less memory during training. We also propose a regularization technique for hierarchical Transformers called shorten factor dropout which improves perplexity upon baselines trained with fixed shorten factor (see Section 4.1). Finally, Hourglass achieves the new state-of-the-art among Transformer models for image generation of ImageNet32 (see Tab. 3).
Standard self-attention mechanism uses full token-level sequence representations. In the Hourglass, we bring efficiency to the model by utilizing shortening, which allows us to use the Transformer layers on inputs with significantly smaller lengths. A high-level overview of our proposed model architecture is shown in figures 2 and 3.
Attention type in the vanilla layers and shortened layers is a configurable parameter. By default we use relative attention defined in Transformer-XL (Dai et al., 2019). Any attention module can be used - we show significant efficiency gains when applying Hourglass also for LSH (Kitaev et al., 2020) attention (see Section 3.2 and Fig. 7).
2.1 Methods of shortening the input sequence
Shortening can be defined as any function  $S$  that accepts a tensor  $x$  of shape  $(l, d)$  and returns a tensor  $x'$  of shape  $(\frac{l}{k}, d)$ , where  $k$  is a hyperparameter called shorten factor.
A simple shortening method is 1D average pooling with stride  $k$  and pool size  $k$ , applied along the sequence dimension  $l$ . Another way of shortening is what we will further call linear pooling ( $l$  and  $d$  denote sequence length and  $d_{model}$ ):
Algorithm 2 LinearPooling
$$
\begin{array}{l} x ^ {\prime} \leftarrow R e s h a p e (x, (\frac {l}{k}, k \cdot d)) \\ x ^ {\prime} \leftarrow \operatorname {L i n e a r P r o j e c t i o n} \left(x ^ {\prime}\right) \\ \end{array}
$$
Shortening can be also performed by attention, as was introduced in (Dai et al., 2020):  $x' = S(x) + \text{Attention}(Q = S(x), K = V = x)$  where  $S$  is shortening function, originally  $S = \text{AvgPool}$ . Directly after this attention operation, a position-wise feed-forward with a residual is performed, so that these two layers form a Transformer block (Vaswani et al., 2017). In this work we also try  $S = \text{LinearPool}$  and find it more effective on image tasks (see Tab. 8).
2.2 Shortening and autoregressive property
Information leaks Shortening interferes with the standard causal masking used in Transformer decoders. Namely, in any shortened representation by a factor of  $k$  each shortened token contributes to predicting up to the next  $k$  tokens in the finest scale, that is if  $e$  is the shortened sequence and  $x$  is the sequence on the finest scale,  $e_0$  is not only used
Figure 2: Hourglass - a high-level architecture overview. The arrows denote residual connections.
to generate  $x_0$ ; in fact, the same embedding is used to generate tokens  $x_0, \ldots, x_{k-1}$ .
Therefore, we need to guarantee that  $e_0$  and any other  $e_i$  cannot access information about tokens they will implicitly predict. To ensure that, we apply another shift right by  $k - 1$  tokens, directly before any shortening by a factor of  $k$  (Fig. 4). The shift is the smallest that does not cause an information leak (see Fig. 5 for an example of a shifting that leads to a leak). We included a more detailed analysis of this fact in the Appendix (Section A.2).
Reduced expressivity Let us consider an Hourglass model with shortening by a factor of  $k$  and no transformer blocks operating on the finest scale (that is, a model without vanilla layers).
In this situation
$$
P (x) \quad = \quad \prod_ {i = 0} ^ {n - 1} P (x _ {i} | e _ {0}, \dots , e _ {\left| \begin{array}{c} i \\ k \end{array} \right|}) \quad =
$$
$$
\prod_ {i = 0} ^ {n - 1} P (x _ {i} | x _ {0}, \dots , x _ {\left\lfloor \frac {i}{k} \right\rfloor \cdot k - 1})
$$
because for predicting  $x_{i}$  we combine the processing done on shortened representations  $e$  with token-independent operations. This means token  $x_{i}$  is generated independently from the tokens  $x\lfloor \frac{i}{k}\rfloor .k,\dots ,x_{i - 1}$ . This situation is detrimental to the model's capabilities, though including at least one vanilla layer solves this issue. In the Appendix we provide a detailed example illustrating this problem (Section A.1).
2.3 Upsampling methods
Upsampling is a crucial part of the Hourglass architecture since we need to convert shortened representations back to the full token-level sequence in order to perform language modeling.
A method proposed in (Dai et al., 2020) is repeating each shortened vector shorten factor times. This method is computationally efficient, but it does not distinguish tokens with respect to position inside the group.
Another method is linear upsampling which works analogously to linear pooling - it projects vectors of shape  $\left(\frac{l}{k},d\right)$  to  $\left(\frac{l}{k},k\cdot d\right)$  and then reshapes to  $l$  vectors, each of dimension  $d$ . This method is fast and allows to project shortened embeddings differently for each position in the group. This happens because the  $(k\cdot d)\times d$  projection matrix can be thought of as  $k$  separate  $d\times d$  matrices, one per each position.
We also investigated a method which we further call attention upsampling. It is similar to attention pooling (Dai et al., 2020) and to the aggregation layer from (Subramanian et al., 2020). It works as follows:  $x = U(x, x') + \text{Attention}(Q = U(x, x'), K = V = x')$  where  $x$  are embeddings from just before the shortening,  $x'$  are final shortened embeddings and  $U$  is an arbitrary upsampling function. After the attention operation there is also a residual with a feed-forward layer.
Linear upsampling learns a fixed pattern that is the same for each shortened token. Attention upsampling has the advantage of being content-based - each token can extract relevant information from the shortened embeddings. We set  $U(x, x') = x + \text{LinearUpsampling}(x')$  which allows to explicitly inject group-level information into the attention queries. We experimentally show that variants of attention upsampling lead to the best results for our model across different datasets (see Tab. 7).
In this section, we present experimental results of Hourglass. We start with a quick analysis of time and memory complexity of the approach (Section 3.1). Then we investigate the efficiency gains of applying Hourglass to Transformers with different attention types (Section 3.2). Finally, we use Hourglass with relative attention parametrization from Transformer-XL (Dai et al., 2019), evaluate it on three language modeling tasks, and compare the results with other models. (Sections 3.3, 3.4)
To show cross-domain generalization of our method, we train our model on one dataset related to Natural Language Processing and two from the Computer Vision field.
To ensure consistency in presenting configurations of our model, we introduce a notation describing hierarchy of our architecture:  $(N_{1}@\mathbf{f}_{1},\dots ,N_{k}@\mathbf{f}_{k})$  where each entry  $(N_{j}@\mathbf{f}_{j})$  means  $N_{j}$  layers shortened by factor  $f_{j}$ .
Our model implementation is open source.
Figure 4: An overview of our shortening approach. Different colors denote token positions. Initially, we shift right by one, which is a standard step in TransformerLM. Then, just before performing shortening, we additionally shift the tokens right by shorten factor -1 to preserve the autoregressive property of the model.
Figure 3: The architecture starts with pre vanilla layers - a stack of Transformer blocks operating on the full token-level sequence. After them we insert shortening layer where  $k$  is the shorten factor parameter (Fig. 4). The sequence is shifted right before shortening to prevent information leak (Fig. 5). Then we recursively insert another Hourglass block operating on  $k$  times smaller scale. On the final level of shortening, we apply shortened layers - Transformer blocks operating on the smallest scale. Upsampling layer brings the resulting activations  $x'$  back to the original resolution. After upsampling and residual, the activations are processed by token-level post vanilla layers.
Figure 5: An example of information leak. If the shift right factor is too small, after upsampling the knowledge from the next tokens leaks to previous ones violating autoregressiveness and making decoding impossible.
3.1 Computational cost analysis
In vanilla Transformers, the number of parameters can indicate the computation required to train the model. This is not true for Hourglass – for instance, it can have 128 layers operating on a sequence shortened by 32 and still fit into the memory of a single GPU. A weak correlation between true Hourglass' computational cost and its number of parameters can be observed in Table 1.
Hourglass achieves the biggest speedup with the standard  $\mathcal{O}(l^2)$  attention. In that case, a single shortening by a shorten factor  $k$  reduces the complexity to  $\mathcal{O}\left(\frac{l^2}{k^2}\right)$  so by a factor of  $k^2$ . For more recent linear-time attention mechanisms (Katharopoulos et al., 2020; Choromanski et al., 2021) the reduction would be smaller - but still by a factor of  $k$ . Feed-forward layers also have linear complexity so shortening reduces it by a factor of  $k$ .
In Table 1 we show an empirical efficiency comparison between Hourglass and Transformer-XL.
1github.com/google/trax/blob/master/trax/models/research/hourglass.py
Hierarchy	BPC	GB	Speed	#Param
6@1 (Baseline)	1.182	4.53	0.95	21M
2@1 1@3 2@1	1.163	4.41	1.11	24M
2@1 4@4 2@1	1.143	4.41	1.10	34M
8@1 (Baseline)	1.151	5.75	0.73	28M
2@1 4@3 2@1	1.128	4.88	1.00	34M
2@1 8@4 2@1	1.128	4.98	0.99	48M
2@1 1@2 4@4 1@2 2@1	1.115	4.69	0.86	48M
2@1 8@3 2@1	1.111	5.50	0.88	48M
10@1 (Baseline)	1.128	6.99	0.56	34M
3@1 8@4 3@1	1.109	6.14	0.76	55M
12@1 (Baseline)	1.115	8.12	0.47	41M
4@1 8@4 4@1	1.098	7.20	0.62	62M
2@1 16@3 2@1	1.096	5.89	0.71	75M
14@1 (Baseline)	1.102	9.35	0.40	48M
5@1 8@2 5@1	1.079	9.57	0.45	69M
Table 1: Efficiency comparison between Hourglass variants and Transformer-XL baseline on enwik8 - we report validation set perplexity (BPC), running memory (GB) and number of training steps per second (Speed). We observe significant perplexity gains over the baseline for a matching computation cost. It is also visible that for Hourglass the number of model parameters (#Param) correlates poorly with true computational cost.
3.2 Impact of Hourglass
To demonstrate the efficiency of Hourglass, we measured how computational cost decreases and perplexity improves, purely adding the technique to Transformer-XL (Dai et al., 2019) and Reformer (Kitaev et al., 2020) backbones (results depicted in Figures 6 and 7, respectively).
In both cases, models are implemented under the same codebase and the only difference between Hourglass and its corresponding baseline is the usage of shortening and upsampling layers. We show that by incorporating a single shortening of the input, we can train larger models with the same memory requirements and training speed and achieve better perplexity than baselines.
3.3 Enwik8
Enwik8 (Mahoney, 2011) is a byte-level language modeling benchmark containing the first 100M bytes of unprocessed English Wikipedia text, split into 90M train, 5M valid, and 5M test sets.
Similarly to (Dai et al., 2019) and (Beltagy et al., 2020), we evaluate our model on the test set, splitting it into overlapping sequences of size  $l = 4096$  with a step size of 128 and calculate the test loss only over the last 128 tokens. With a  $(4@1, 8@3, 4@1)$  hierarchy,  $d_{model} = 768$ ,  $d_{ff} = 3072$  and 8 heads, we reach 0.98 test bits-per-character.
Figure 6: Comparison between Transformer-XL baseline and Hourglass on Enwik8 valid set w.r.t. maximum memory used during training. All models are trained for 200k steps with the same hyperparameters.
Figure 7: Comparison between Reformer baseline and Hourglass, both with LSH attention, on Enwik8 valid set w.r.t. cost of one training step in seconds.
3.4 Image Generation
We use datasets introduced in (van den Oord et al., 2016a) which are downsampled versions of the popular ImageNet. In the autoregressive image generation setup, they consist of respectively  $32 \times 32 \times 3$  and  $64 \times 64 \times 3$  tokens, corresponding to RGB channels, per image. As the only preprocessing step we flatten the images.
3.4.1 ImageNet32
For our main result the following hierarchy is used: (3@1, 24@3, 3@1). We use  $d_{\mathrm{model}} = 512$ ,  $d_{\mathrm{ff}} = 2048$ , 8 attention heads and 0.01 dropout rate. With this configuration we achieve 3.741 bits/dim, yielding the new state-of-the-art among autoregressive (Transformer-based) models on this dataset, compared to the previous state-of-the-art of 3.758 bpd by (Ho et al., 2019).
Enwik8	#Param	BPC
Transformer-XL (2019) 24L	277M	0.99
Hourglass	146M	0.98
Adaptive-Span (2019) 24L	209M	0.98
Transformer-LS (2021)	110M	0.97
Feedback Transformer (2021)	77M	0.96
Expire-Span (2021) 24L	277M	0.95
Figure 8: Examples of our model completions, where bottom half of each image was generated by our model, prompted by the upper half.
3.4.2 ImageNet64
The sequence length that our model can handle is limited mainly by the computational complexity of used attention module. We replace relative attention in vanilla layers by LSH attention (Kitaev et al., 2020), which allows us to handle 12288-long sequences. To achieve relative attention parametrization, the LSH attention is combined with rotary positional embeddings (Su et al., 2021). In shortened layers, standard relative attention is used. For LSH attention, we set chunk length to 128 and use 2 hashes, which results in small memory complexity in our full-size layers. In this setup, we reach a score of 3.443 bpd with a  $(3@1, 12@3, 3@1)$  architecture. All attention layers had  $d_{model} = 768$ ,  $d_{ff} = 3072$  and 8 heads. No dropout was used.
3.4.3 CIFAR-10
CIFAR-10 (Krizhevsky, 2009) is an image dataset consisting of 60000 images of size  $32 \times 32$ . We use this dataset primarily for our ablations (Section 4). Due to the relatively small number of examples compared to ImageNet, models reach convergence after 100k steps.
4 Ablations
In this section, we start by introducing a training technique called shorten factor dropout (Section 4.1), and then analyze Hourglass's components de
Table 2: Enwik8 Results. We report bits-per-character (BPC) on the test set and number of model parameters. Hourglass applied to Transformer-XL significantly outperforms its baseline. Our technique could be also used with other more performant attention methods which we leave for future work.
ImageNet32	BPD
PixelCNN (van den Oord et al., 2016b)	3.83
Image Transformer (Parmar et al., 2018)	3.77
Axial Transformer (Ho et al., 2019)	3.76
Hourglass	3.74
VDM (Kingma et al., 2021)	3.72
DenseFlow (Grcić et al., 2021)	3.63
ImageNet64	BPD
Reformer (Kitaev et al., 2020)	3.65
Performer (Choromanski et al., 2021)	3.64
Hourglass	3.44
Sparse Transformer (Child et al., 2019)	3.44
Routing Transformer (Roy et al., 2020)	3.43
Combiner (Ren et al., 2021)	3.42
VDM (2021)	3.40
DenseFlow (2021)	3.35
Table 3: Bits per Dimension (BPD) on downsampled imagenet. Autoregressive models are separated by a horizontal line from non-autoregressive ones. On ImageNet32, our model yields new state-of-the-art for autoregressive models.
scribed above. We show that shortened layers behave similarly to full token-level layers in terms of scalability (Section 4.2). Then we study the effect of different distributions of  $(pre, post)$  vanilla layers on Hourglass' accuracy (Section 4.3). We further analyze the performance of various upsampling and downsampling methods (Sections 4.4 and 4.5). Finally, we discuss different shorten factors and multi-stage shortening in Section 4.6.
We conduct the ablations on both text and image generation to show applicability across different domains. We report bits per character (BPC) on the enwik8 validation (dev) set evaluated without context (sequence length 2048) and bits per dim (BPD) on the CIFAR-10 test set. For the exact hyperparameter setup refer to the Appendix.
4.1 Shorten factor dropout
Different shorten factors can be used for the same model when using parameterless pooling methods. We propose a training procedure where the shorten factor is randomly sampled with uniform distribution from a predefined set in each step. We observe that such a training regime improves validation loss compared to a baseline trained with a single, fixed shorten factor. For example, a model trained with shorten factor randomly sampled from  $\{2,3\}$  performs better when evaluated with any of these shorten factors, compared to models trained with a corresponding fixed shorten factor (Tab. 4).
We hypothesise that such a technique promotes a more uniform distribution of information over the sequence of tokens. It may be essential for fixed-size pooling techniques as they do not ac
count for variable length constituents like words. By spreading information uniformly, we prevent a situation where we lose content by shortening three information-dense tokens or lose available capacity by merging three low information ones.
Shorten factor dropout is not limited to our architecture and can be applied to any model that utilizes shortening, particularly (Dai et al., 2020).
Hierarchy	Train k	Val k = 2	Val k = 3
2@1 8@k 2@1	{2,3}	1.104	1.116
2	1.116
3	1.124
4@1 12@k 4@1	{2,3}	1.086	1.094
2	1.098
3	1.101
5@1 10@k 5@1	{2,3}	1.082	1.087
2	1.096
3	1.095
4.2 Scaling shortened layers
In this study, we show that layers operating on the shortened sequence contribute significantly to Hourglass's accuracy. In Table 5 we measure the impact of scaling the depth of the shortened part of the model with a fixed number of vanilla layers.
We also check if scaling laws of Transformers, described in (Kaplan et al., 2020), hold by comparing a regression line fitted to various Hourglass configurations and one fitted to Transformer-XL baseline. We observe in Figure 1 that the slopes are very similar, which indicates that the laws hold.
Table 4: Comparison between models trained with shorten factor dropout (Train  $k = \{ 2,3\}$  ,Section 4.1) and fixed shorten factor baselines on enwik8.
Number of shortened layers	enwik8	CIFAR-10
Baseline (n = 1)	1.164	3.28
n = 4	1.134	3.16
n = 8	1.111	3.07
n = 16	1.096	3.03
4.3 Impact of vanilla layers