Elamur Research⌕
☰
home / papers / 2022 / Transformer Quality in Linear Time

Transformer Quality in Linear Time

Paper · 2022 · Weizhe Hua, Zihang Dai, Hanxiao Liu et al.

Summary

Standard Transformers scale quadratically with sequence length. The authors argue that existing efficient-attention alternatives have at least one of three problems: they lose quality against well-tuned baselines, they run slower on GPUs/TPUs than their theoretical cost suggests, or they train slowly in the auto-regressive setting.

The work first redesigns the layer itself. The Gated Attention Unit (GAU) fuses a GLU-style gating block with a lightweight attention module, so a single small attention head that uses squared ReLU instead of softmax is enough. On top of GAU, the authors add mixed chunk attention. It splits the sequence into chunks, applies exact attention within each chunk and cheap linear attention between chunks. This gives linear cost that runs efficiently on accelerators.

The resulting model, FLASH, is tested on masked and auto-regressive language modeling (C4, Wiki-40B, PG-19) at 110M–500M parameters, with contexts from 512 to 8K tokens. It matches the perplexity of an enhanced Transformer++ while training several times faster. This shows that efficient attention can deliver real speed gains, not only theoretical ones.

Key points

  • Over Transformer++, FLASH trains up to 4.9× faster on Wiki-40B, up to 4.8× faster for masked LM on C4, and up to 12.1× cheaper on PG-19. Measurements were made on TPU-v4.
  • GAU's attention adds only one small projection (shared dimension s = 128 by default). Two GAU layers replace one Transformer block (MHSA + MLP/GLU) at similar size and speed.
  • Ablations show the local quadratic and global linear parts are complementary; removing either clearly hurts quality. Applying the same chunking to Transformer++ (MC-TFM++) costs much more perplexity, e.g., +0.87 vs. +0.11 for FLASH on Wiki-40B at 8K context.
  • On TriviaQA (4096-token context, ~110M models), FLASH s=512 with first-to-all attention reaches F1 73.9, versus 74.2 for Transformer+. It cuts pretraining and fine-tuning latency by about 2.8× and 2.7×.
  • The default chunk size is 256; larger chunks such as 512 tend to work better once the context exceeds 1024 tokens.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🏢GoogleMountain View · United States492 works · 440,463 citations🎓Cornell UniversityIthaca · United States134 works · 65,159 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 worksComputer Vision and Pattern Recognition · 3,486 works

Cite

Cite as Weizhe Hua, Zihang Dai, Hanxiao Liu and Quoc V. Le. 2022. Transformer Quality in Linear Time. arXiv:2202.10447.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2202.10447 · arXiv 2202.10447

@misc{hua2022transformer,
  title = {{Transformer Quality in Linear Time}},
  author = {Weizhe Hua, Zihang Dai, Hanxiao Liu and Quoc V. Le},
  year = {2022},
  doi = {10.48550/arxiv.2202.10447},
  eprint = {2202.10447},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2202.10447}
}

References

37 references, 21 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
  2. Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
  3. Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. arXiv preprint arXiv:2005.14165, 2020.
  4. Ho, J., Kalchbrenner, N., Weissenborn, D., and Salimans, T. Axial attention in multidimensional transformers. arXiv preprint arXiv:1912.12180, 2019. — [arXiv:1912.12180](https://arxiv.org/abs/1912.12180) find
  5. Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
  6. Huang, Z., Wang, X., Huang, L., Huang, C., Wei, Y., and Liu, W. Ccnet: Criss-cross attention for semantic segmentation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 603–612, 2019. find
  7. Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794, 2020.
  8. Jaegle, A., Gimeno, F., Brock, A., Vinyals, O., Zisserman, A., and Carreira, J. Perceiver: General perception with iterative attention. In International Conference on Machine Learning, pp. 4651–4664. PMLR, 2021.
  9. Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q. V., and Salakhutdinov, R. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.02860, 2019. find
  10. Joshi, M., Choi, E., Weld, D., and Zettlemoyer, L. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1601–1611, Vancouver, Canada, July 2017…
  11. Dauphin, Y. N., Fan, A., Auli, M., and Grangier, D. Language modeling with gated convolutional networks. In Proceedings of the 34th International Conference on Machine Learning - Volume 70, ICML’17, pp. 933–941. JMLR.org, 2017.
  12. Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
  13. Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
  14. Du, N., Huang, Y., Dai, A. M., Tong, S., Lepikhin, D., Xu, Y., Krikun, M., Zhou, Y., Yu, A. W., Firat, O., et al. Glam: Efficient scaling of language models with mixture-of-experts. arXiv preprint arXiv:2112.06905, 2021.
  15. Elfwing, S., Uchibe, E., and Doya, K. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 107:3–11, 2018. find
  16. Gehring, J., Auli, M., Grangier, D., Yarats, D., and Dauphin, Y. N. Convolutional sequence to sequence learning. In Proceedings of the 34th International Conference on Machine Learning - Volume 70, ICML’17, pp. 1243–1252. JMLR.org, 2017.
  17. Guo, M., Dai, Z., Vrandecic, D., and Al-Rfou, R. Wiki-40b: Multilingual language model dataset. In LREC 2020, 2020. find
  18. Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415, 2016. find
  19. Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pp. 5156–5165. PMLR, 2020. find
  20. Kitaev, N., Kaiser, Ł., and Levskaya, A. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451, 2020.
  21. Li, S., Jin, X., Xuan, Y., Zhou, X., Chen, W., Wang, Y.-X., and Yan, X. Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting. Advances in Neural Information Processing Systems, 32:5243–5253, 2019. find
  22. Liu, H., Dai, Z., So, D. R., and Le, Q. V. Pay attention to mlps. NeurIPS, 2021. find
  23. Narang, S., Chung, H. W., Tay, Y., Fedus, W., Fevry, T., Matena, M., Malkan, K., Fiedel, N., Shazeer, N., Lan, Z., et al. Do transformer modifications transfer across implementations and applications? arXiv preprint arXiv:2102.11972, 2021. find
  24. Nguyen, T. Q. and Salazar, J. Transformers without tears: Improving the normalization of self-attention. CoRR, abs/1910.05895, 2019. URL http://arxiv.org/ abs/1910.05895. ↗
  25. Peng, H. et al. Random feature attention. In ICLR, 2021. find
  26. Rae, J. W., Potapenko, A., Jayakumar, S. M., and Lillicrap, T. P. Compressive transformers for long-range sequence modelling. arXiv preprint arXiv:1911.05507, 2019.
  27. Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1–67, 2020. URL http://jmlr.org/papers/v21/20-074.html. ↗
  28. Ramachandran, P., Zoph, B., and Le, Q. V. Searching for activation functions. CoRR, abs/1710.05941, 2017. URL http://arxiv.org/abs/1710.05941. ↗
  29. Ren, H., Dai, H., Dai, Z., Yang, M., Leskovec, J., Schuurmans, D., and Dai, B. Combiner: Full attention transformer with sparse computation cost. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, 2021. URL https://openreview.net… ↗
  30. Roy, A., Saffar, M., Vaswani, A., and Grangier, D. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53–68, 2021.
  31. Shazeer, N. GLU variants improve transformer. CoRR, abs/2002.05202, 2020. URL https://arxiv.org/ abs/2002.05202. ↗
  32. So, D. R., Mańke, W., Liu, H., Dai, Z., Shazeer, N., and Le, Q. V. Primer: Searching for efficient transformers for language modeling. NeurIPS, 2021. find
  33. Su, J., Lu, Y., Pan, S., Wen, B., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding, 2021.
  34. Thoppilan, R., De Freitas, D., Hall, J., Shazeer, N., Kulshreshtha, A., Cheng, H.-T., Jin, A., Bos, T., Baker, L., Du, Y., et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.08239, 2022.
  35. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. In Advances in neural information processing systems, pp. 5998–6008, 2017.
  36. Wang, S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020. find
  37. Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q., Yang, L., et al. Big bird: Transformers for longer sequences. In NeurIPS, 2020.

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/hua-2022-transformer-quality-linear-time/