PaperlayerView exact revision on arXiv

Source-linked AI summary

Hierarchical Transformers Are More Efficient Language Models

Piotr Nawrot, Szymon Tworkowski, Michał Tyrolski, Łukasz Kaiser, Yuhuai Wu, Christian Szegedy, Henryk Michalewski

arXiv:2110.13711v2cs.LGcs.CL

TL;DR

Long-context Transformers incur high costs, and generative models are forced to operate at input granularity through much of the stack. Hourglass shortens internal sequences and expands them before generation; it improves results across language modeling and image generation, including a Transformer-model state of the art on ImageNet32.

  • Problem

    Self-attention has quadratic complexity in context length, while generative Transformers are forced to retain input-scale sequences through much of the layer stack.

    Source: §1 · pp. 1–2

  • Method

    Hourglass shortens internal activation sequences deeper in the stack and upsamples them before generation, combining them with representations from earlier layers.

    Source: §1 · p. 2

  • Results

    Hourglass outperforms the baseline in perplexity at a given computation cost and running-memory metrics, achieves Transformer-model state of the art on ImageNet32, and is competitive on other image-generation and language-modeling tasks.

    Source: §6 · p. 8

  • Takeaways & Limitations

    Hourglass works with any attention type, opening directions for research on Transformers capable of processing longer sequences.

    Source: §6 · p. 8

    Shortening can make a token independent of preceding tokens in its group; including at least one vanilla layer solves this issue.

    Source: §2.2 · p. 3

Abstract

from arXiv · show

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: full 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 downsampling 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.

1 Introduction

Transformers handle long contexts but face high computational costs and are constrained by operating at the input sequence’s granularity throughout the stack. Hourglass introduces hierarchical processing, shortening sequences deeper in the stack and expanding them before generation; it improves enwik8 results and sets a Transformer-model state of the art on ImageNet32.

1 Introduction

  • Transformers handle long contexts across tasks including summarization and image generation, where relevant context may be far from the next output.The paper attributes this capability to self-attention.Source: §1 · p. 1
  • Original self-attention scales quadratically with context length, and even linear-complexity layers can be slow and memory-intensive on very long sequences.The paper describes vanilla Transformers as infeasible for tasks with very long input sequences, such as high-resolution images.Source: §1 · pp. 1–2
  • Generative models need input-scale processing early and output-scale processing at the end, but high-level words, entities, or events differ in granularity from input letters.The authors identify forcing the same granularity throughout the layer stack as a fundamental and practical shortcoming.Source: §1 · p. 2
  • Hourglass shortens internal activations deeper in the stack and upsamples before generation, combining with earlier-layer sequences; unlike Funnel-Transformer and U-Nets, it is autoregressive.The shortening stage is analogous to Funnel-Transformer, while the full architecture takes inspiration from U-Nets.Source: §1 · p. 2
  • Hourglass improves enwik8 perplexity with Transformer-XL and Reformer blocks, trains faster with less memory, and sets a Transformer-model state of the art on ImageNet32.The paper also reports competitive results on other image-generation and language-modeling tasks.Source: §1 · p. 2

2 Model

Hourglass shortens internal sequences for Transformer layers and restores them to token-level resolution for language modeling. Its design combines multiple shortening and upsampling methods with a right shift that preserves autoregressive prediction.

2 Model

  • Hourglass uses shortening to give Transformer layers significantly shorter sequences, then upsamples activations to full token-level resolution for language modeling.The architecture includes pre-vanilla layers at full resolution, recursively shorter Hourglass blocks, and token-level post-vanilla layers.Source: §2 · p. 2, §2.3 · p. 3, Figure 3 · p. 4

2.1 Methods of shortening the input sequence

  • Shortening maps a sequence tensor to a shorter sequence with the same feature dimension; options include average pooling, linear pooling, or attention-based shortening.Attention shortening combines a shortened query with attention over the original sequence; linear pooling is reported as more effective on image tasks.Source: §2.1 · p. 2

2.2 Shortening and autoregressive property

  • A shortened representation can contribute to predicting up to k finest-scale tokens, so the sequence is shifted right by k−1 tokens before shortening to prevent information leakage.Without finest-scale vanilla layers, a token is generated independently of preceding tokens within its shortened group; including at least one vanilla layer solves this issue.Source: §2.2 · pp. 2–3

2.3 Upsampling methods

  • Upsampling restores shortened representations to token-level resolution; repeating vectors is efficient but ignores within-group position, while linear and attention upsampling offer alternatives.Linear upsampling projects shortened vectors into position-specific representations; attention upsampling can retrieve relevant information from shortened embeddings.Source: §2.3 · p. 3
  • Attention upsampling performs best across datasets by retrieving content-based information from shortened embeddings and combining it with group-level information.The method uses queries formed from the pre-shortening embeddings plus linearly upsampled shortened embeddings.Source: §2.3 · p. 3

3 Experiments

Experiments evaluate Hourglass across language and image generation, comparing efficiency and performance with Transformer baselines and testing multiple model configurations.

3.1 Computational cost analysis

  • Shortening reduces standard quadratic-attention complexity by k^2 and linear-attention and feed-forward costs by k; parameter count correlates poorly with Hourglass's true computational cost.Table 1 compares validation BPC, running memory, and speed, and reports perplexity gains over the baseline at matching computation cost.Source: §3.1 · p. 4, Table 1 · p. 5

3.2 Impact of Hourglass

  • Shortening and upsampling alone let models train faster and with less memory while improving perplexity over Transformer-XL and Reformer baselines.The comparisons use the same codebase and otherwise corresponding baseline models; Figures 6 and 7 compare memory and training-step cost, respectively.Source: §3.2 · p. 5, Figure 6 · p. 5, Figure 7 · p. 5

3.3 Enwik8

  • 0.98 test bits-per-character is achieved on enwik8 using overlapping 4,096-token sequences stepped by 128 and scoring the final 128 tokens.Table 2 reports that Hourglass applied to Transformer-XL significantly outperforms its baseline.Source: §3.3 · p. 5, Table 2 · p. 6
  • Hourglass reaches 0.98 test bits-per-character on enwik8, a byte-level benchmark built from the first 100M bytes of unprocessed English Wikipedia text.The benchmark is split into 90M training, 5M validation, and 5M test bytes.Source: §3.3 · p. 5

3.4.1 ImageNet32

  • With the (3@1, 24@3, 3@1) hierarchy, the model achieves 3.741 bits/dim, a new state of the art among autoregressive Transformer-based models on this dataset, versus the prior 3.758 bpd result by (Ho et al., 2019).Source: §3.4.1 · p. 5

3.4 Image Generation

  • Image-generation examples prompt on each image's upper half and generate the lower half; the ImageNet data are flattened RGB-token sequences at 32×32×3 and 64×64×3.The experiments also include CIFAR-10, a 32×32 image dataset used primarily for ablations, with models reaching convergence after 100k steps.Source: §3.4 · p. 5, Figure 8 · p. 6, §3.4.3 · p. 6

3.4.2 ImageNet64

  • The model reaches 3.443 bpd with a (3@1, 12@3, 3@1) architecture, using LSH attention to handle sequences of length 12,288 with chunk length 128 and two hashes.Source: §3.4.2 · p. 6

4 Ablations

Ablations test shorten-factor dropout, layer allocation, upsampling and pooling choices, and shortening strategies on enwik8 and CIFAR-10.

4.1 Shorten factor dropout

  • Uniformly sampling shorten factors from {2, 3} during training improves validation loss over fixed-factor training when evaluated with either factor.Table 4 compares the sampled-factor model with fixed-factor baselines on enwik8.Source: §4.1 · p. 6, Table 4 · p. 7

4.2 Scaling shortened layers

  • Hourglass scaling-law regression slopes are very similar to those of the Transformer-XL baseline.The section also reports that shortened layers contribute significantly to accuracy as their depth increases with vanilla-layer count fixed.Source: §4.2 · p. 7, Table 5 · p. 7

4.3 Impact of vanilla layers

  • Equal numbers of vanilla layers before and after shortening yield the best perplexity; adding vanilla layers also significantly improves performance.Vanilla layers process the sequence at token level without shortening.Source: Table 6 · p. 7, §4.3 · p. 7

4.4 Upsampling method

  • Attention upsampling has the lowest perplexity across the tested functions; among attention-free methods, linear works better on images and repeat works well for text.Table 7 compares upsampling methods using the stated enwik8 and CIFAR-10 baseline configurations.Source: Table 7 · p. 7, §4.4 · p. 7

4.5 Pooling method

  • Attention pooling yields the lowest perplexity on both datasets; among attention-free methods, average pooling favors text while linear pooling favors images.These modality preferences persist with attention pooling despite small differences across shortening functions.Source: Table 8 · p. 7, §4.5 · p. 7

4.6 Shortening strategies

  • At least one shortening by a factor of 3 or more, with 2–3 vanilla layers, is consistently beneficial; similar-computation configurations perform similarly.The experiments use sequence length 2048, and the authors hypothesize that more hierarchy may help on longer sequences.Source: §4.6 · p. 8

5 Related Work

Related work places Hourglass among approaches to sequence shortening, positional encoding, sparse attention, and image generation.

Shortening in Transformers

  • Unlike Funnel-Transformer's text-classification encoder, Hourglass targets autoregressive language modeling and uses attention upsampling rather than repeat upsampling.Other character-level shortening work cited here focuses mainly on the encoder, whereas Hourglass applies shortening in the decoder.Source: §5 · p. 8
  • Hourglass uses a symmetric, U-Net-like shortening and upsampling design, with an upsampling operation at each scale change.The cited Subramanian et al. approaches focus on downsampling or upsampling; Hourglass also preprocesses tokens at the finest scale.Source: §5 · p. 8

Relative positional encoding

  • Hourglass replaces Transformer-XL's segment-level recurrence with shortening and can combine with different attention types, including LSH attention.Its relative-attention parametrization follows Transformer-XL; rotary positional embeddings are associated with RoFormer.Source: §5 · p. 8

Sparse Attention

  • Unlike sparse-attention approaches, Hourglass retains full attention over shortened sequences and treats the attention module as a drop-in choice.The related-work discussion names Routing and Sparse Transformer as sparse-pattern approaches and Combiner as using structured factorization.Source: §5 · p. 8

Image generation on downsampled ImageNet

  • The cited VDM and DenseFlow image-generation methods are non-autoregressive, unlike the Transformer-based methods discussed here.The comparison includes Parmar et al., Ho et al., and this work.Source: §5 · p. 8

6 Conclusion

The conclusion reports hierarchy-based efficiency gains and Hourglass’s benchmark standing, while leaving task-specific hierarchy selection open and pointing to longer-context experiments for guidance.

6 Conclusion

  • Hourglass improves Transformer efficiency in language modeling and achieves state-of-the-art autoregressive results on ImageNet32, with competitive results on other image-generation and language-modeling tasks.The conclusion also reports improvements in perplexity at a given computation cost and in empirical running-memory metrics.Source: §6 · p. 8
  • Dynamic pooling is proposed as a possible response to fixed-size groups in multi-stage shortening, and choosing the best hierarchy for a task remains open.The authors conjecture that experiments with much longer contexts will provide better guidance and benefit Hourglass more.Source: §6 · pp. 8–9

A Autoregressive shortening

A.1 Motivation behind using vanilla layers

The appendix motivates vanilla layers by showing that autoregressive prediction can require fine-scale token access, which shortening alone may remove.

A.1 Motivation behind using vanilla layers

  • 66% sequence accuracy is achievable on a repeated-chunk task because a vanilla Transformer can copy the token two positions earlier to predict each chunk’s last token.It still cannot predict the token at the beginning of the chunk.Source: §A.1 · p. 11
  • With shorten factor 3 and no vanilla layers, the example task cannot exceed 50% accuracy when the vocabulary has at least two tokens.Tokens at positions divisible by 3 cannot be predicted better than randomly from the available shortened and residual embeddings.Source: §A.1 · p. 11
  • At least one vanilla layer solves this issue by allowing attention within the neighborhood of k previous tokens; local attention with context size k suffices for the example.Source: §A.1 · p. 11

A.2 Information leaks – analysis

The analysis defines autoregressive dependence and information leakage, then describes how shifting and shortening can violate the autoregressive property.

A.2.1 Definition of autoregressive model

  • An autoregressive model predicts each token using only preceding tokens, never itself or later tokens.Source: §A.2.1 · p. 11

A.2.2 Definition of information leak

  • Information leak is defined as a transformation allowing a token to depend on a later token, contrary to autoregressive ordering.Source: §A.2.2 · p. 11

A.2.3 Model representation

  • The representation applies a right shift, shortening and upsampling functions, and decoder blocks; causal attention masking in the decoder prevents leaks caused by those blocks.The analysis also omits attention-mediated token flows because the mask preserves autoregressive behavior.Source: §A.2.3 · pp. 11–12, §A.2.4 · p. 12

A.2.4 Leak description

  • For a shorten factor k, the analysis states that any shift 0 < s < k − 1 causes a leak because the first output can depend on input tokens beyond the autoregressive boundary.The resulting dependency violates the autoregressive property.Source: §A.2.4 · p. 12

B Experimental setup

The experimental setup specifies shared model and optimizer choices, task-specific training configurations, hardware, and software.

B.1 Common parameters

  • Common experiments use Adam with β1 = 0.9, β2 = 0.98, and ϵ = 1e−9, without weight decay or gradient clipping, alongside Pre-Norm and FastGelu.Source: §B.1 · p. 12

B.2 Enwik8

  • Ablation-study models were trained for 200k steps with a cosine learning-rate schedule, 200k-step cycle, and 4,000-step linear warmup; the main 0.98 bpc result used 146M parameters.The main model used a (4@1, 8@3, 4@1) hierarchy; selected batch size, dropout, and learning rate were 8, 0.15, and 4e−4, respectively. Subsequent experiments used these parameters without additional searching.Source: §B.2 · p. 12

B.3 Downsampled ImageNet - common parameters

  • ImageNet32 and ImageNet64 experiments shared inverse-square-root learning-rate decay with 8000 warmup steps and total batch size 64.Source: §B.3 · p. 12

B.4 ImageNet32

  • ImageNet32 used input length 3072, 512-dimensional model width, 2048-dimensional feed-forward layers, eight heads, and 470k training steps.Training used linear warmup and inverse-square-root decay for 400k steps, followed by cosine decay for 70k steps.Source: §B.4 · p. 12

B.5 ImageNet64

  • ImageNet64 used 12288 tokens for 64 × 64 × 3 images, with model width 768, feed-forward width 3072, eight heads, zero dropout, and 300k training steps.Source: §B.5 · p. 12

B.6 CIFAR-10

  • Ablation studies generally ran for 100k training steps, unless specified otherwise, on length-3072 inputs with dmodel 512, dff 2048, 8 attention heads, zero dropout, and total batch size 8.Source: §B.6 · p. 12

C.1 Hardware

  • Ablations and short runs used four Titan V GPUs, longer runs used eight RTX 2080 Ti GPUs, and a few longest runs used 8×8 TPUv3 units.Source: §C.1 · p. 13

C.2 Software

  • Experiments used Trax 1.3.9, while shorten-factor-dropout experiments modified the Transformer-XL PyTorch codebase.Source: §C.2 · p. 13

D Reproducibility

The authors provide public code and Colab notebooks to reproduce the main Enwik8 and ImageNet32/64 results. Experiments use randomly chosen seeds, with reruns varying by study type.

D Reproducibility

  • Public code and Colab notebooks support reproducing the main Enwik8 and ImageNet32/64 results.The code is available at github.com/google/trax.Source: D Reproducibility · p. 13

D.1 Randomness

  • Experiment seeds were chosen randomly, but experiment histories allow retrieval of all randomly set parameters for reproduction.Source: §D.1 · p. 13
  • Baseline runs for ablations were repeated three times to calculate standard deviation, while other experiments ran once because of cost and minimal observed variance.Source: §D.1 · p. 13

D.2 Experiment representation

Each experiment is represented by a configuration file that unambiguously specifies its setup, including hyperparameters and training details.

D.2 Experiment representation

  • Configuration files unambiguously specify each experiment’s setup, including optimizers, data preprocessing functions, and batch size per device.Source: §D.2 · p. 13

Related papers

Loading 2110.13711v2…