Elamur Research⌕
☰
home / papers / 2024 / Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models

Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models

Paper · 2024 · Zhen Qin, Weigao Sun, Dong Li et al.

Summary

Linear attention should cost O(n) instead of the O(n²) of softmax attention. In causal language modeling, however, existing implementations lose this advantage: the kernel trick needs a cumulative summation (cumsum) that parallelizes poorly on GPUs. Lightning Attention-1 had already reduced the memory-I/O bottleneck, but its complexity stayed at O(n²d). Lightning Attention-2 splits the sequence into blocks. Within a block, it computes attention the conventional way (left product). Contributions from earlier blocks come through the linear-attention right product, using a KV state that is accumulated in SRAM. Tiling covers both the forward and backward passes, and the kernel is written in Triton to be IO-aware. The authors tested it inside TransNormerLLM (TNL) at 400M, 1B and 3B parameters, plus a 15B model. Training throughput stays nearly constant as sequence length grows, beating LLaMA with FlashAttention-2 and TNL with Lightning Attention-1. It also uses less memory and reaches comparable or lower loss. This makes training long-context LLMs from scratch practical without a speed penalty, with GPU memory as the main remaining limit.

Key points

  • Core idea: intra-block computation uses the standard left product, while inter-block computation uses the right product (K⊤V) via the kernel trick. This avoids the cumsum bottleneck in the causal setting.
  • Speed (tokens per GPU per second, TGS): TNL-LA2 stays roughly flat from 1K to 92K tokens, and up to 131,072 in the 400M/1B/3B plots. LLaMA-FA2 and TNL-LA1 slow down sharply as length grows.
  • Loss comparisons against HGRN, TNN and LLaMA-FA2 were run at a 2K context length. The 1B model used 16×A800 80G GPUs and the 3B model used 32×A800 80G GPUs.
  • TransNormerLLM-15B has 42 layers, 40 heads and an embedding dimension of 5120. It is planned for training on more than 1.3T tokens at sequence length 6,144 and runs at about 1,620 TGS.
  • The authors evaluated an interim 15B checkpoint with lm-evaluation-harness:
    • Commonsense reasoning: the checkpoint beats Pythia-12B by about 2% on average across BoolQ, PIQA, SIQA, HellaSwag, WinoGrande, ARC and OBQA. The CSR average is 56.76 vs 54.58 at about 100B tokens.
    • C-Eval: it scores about 2% above Pythia-12B.
    • MMLU and C-Eval overall: results are only slightly above the 25% random baseline.

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

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksHardware and Architecture · 81 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works

Cite

Cite as Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun and Yiran Zhong. 2024. Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models. arXiv:2401.04658.

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

@misc{qin2024lightning,
  title = {{Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models}},
  author = {Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun and Yiran Zhong},
  year = {2024},
  doi = {10.48550/arxiv.2401.04658},
  eprint = {2401.04658},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2401.04658}
}

References

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

  1. Biderman, S., Schoelkopf, H., Anthony, Q., Bradley, H., O’Brien, K., Hallahan, E., Khan, M. A., Purohit, S., Prashanth, U. S., Raff, E., Skowron, A., Sutawika, L., and van der Wal, O. Pythia: A suite for analyzing large language models across training and scaling, 2023. find
  2. Bisk, Y., Zellers, R., Bras, R. L., Gao, J., and Choi, Y. Piqa: Reasoning about physical commonsense in natural language, 2019.
  3. Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877–1901, 2020.
  4. Chen, S., Wong, S., Chen, L., and Tian, Y. Extending context window of large language models via positional interpolation, 2023. find
  5. Chi, T.-C., Fan, T.-H., Ramadge, P. J., and Rudnicky, A. I. Kerple: Kernelized relative positional embedding for length extrapolation, 2022. find
  6. Chi, T.-C., Fan, T.-H., Rudnicky, A. I., and Ramadge, P. J. Dissecting transformer length extrapolation via the lens of receptive field analysis, 2023. find
  7. Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlós, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., Belanger, D., Colwell, L. J., and Weller, A. Rethinking attention with performers. ArXiv, abs/2009.14794, 2020.
  8. Clark, C., Lee, K., Chang, M.-W., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no questions, 2019.
  9. Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018.
  10. Dao, T. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023.
  11. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, 2022. find
  12. Gao, L., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi, A., Foster, C., Golding, L., Hsu, J., Le Noac’h, A., Li, H., McDonell, K., Muennighoff, N., Ociepa, C., Phang, J., Reynolds, L., Schoelkopf, H., Skowron, A., Sutawika, L., Tang, E., Thite, A., Wang, B., Wang, K., and Zou, A. A framework… find
  13. Hao, D., Mao, Y., He, B., Han, X., Dai, Y., and Zhong, Y. Improving audio-visual segmentation with bidirectional generation. In Proceedings of the AAAI Conference on Artificial Intelligence, 2024. find
  14. Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding, 2021.
  15. Hua, W., Dai, Z., Liu, H., and Le, Q. V. Transformer quality in linear time. arXiv preprint arXiv:2202.10447, 2022. find
  16. Huang, Y., Bai, Y., Zhu, Z., Zhang, J., Zhang, J., Su, T., Liu, J., Lv, C., Zhang, Y., Lei, J., Fu, Y., Sun, M., and He, J. C-eval: A multi-level multi-discipline chinese evaluation suite for foundation models, 2023.
  17. 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, 2020a. find
  18. Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Re… find
  19. PMLR, 2020b. URL http://proceedings.mlr.press/v119/katharopoulos20a.html. ↗
  20. Ke, G., He, D., and Liu, T.-Y. Rethinking positional encoding in language pre-training. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=09-528y2Fgf. ↗
  21. Li, J., Li, D., Xiong, C., and Hoi, S. BLIP: Bootstrapping language-image pre-training for unified vision-language understanding and generation. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), Proceedings of the 39th International Conference on Machine Learni…
  22. Li, J., Li, D., Savarese, S., and Hoi, S. Blip-2: Bootstrapping language-image pre-training with frozen image encoders and large language models. In arXiv, 2023a.
  23. Li, W., Li, D., Li, W., Wang, Y., Jie, H., and Zhong, Y. MAP: Low-data regime multimodal learning with adapter-based pre-training and prompting. In Breitholtz, E., Lappin, S., Loaiciga, S., Ilinykh, N., and Dobnik, S. (eds.), Proceedings of the 2023 CLASP Conference on Learning with Small Data (LSD)… find
  24. Liu, H., Li, C., Wu, Q., and Lee, Y. J. Visual instruction tuning. In arXiv, 2023.
  25. Lu, K., Liu, Z., Wang, J., Sun, W., Qin, Z., Li, D., Shen, X., Deng, H., Han, X., Dai, Y., and Zhong, Y. Linear video transformer with feature fixation, 2022. find
  26. Mao, Y., Zhang, J., Xiang, M., Zhong, Y., and Dai, Y. Multimodal variational auto-encoder based audio-visual segmentation. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 954–965, October 2023. find
  27. Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can a suit of armor conduct electricity? a new dataset for open book question answering, 2018. find
  28. Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
  29. Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Biderman, S., Cao, H., Cheng, X., Chung, M., Derczynski, L., Du, X., Grella, M., Gv, K., He, X., Hou, H., Kazienko, P., Kocon, J., Kong, J., Koptyra, B., Lau, H., Lin, J., Mantri, K. S. I., Mom, F., Saito, A., Song, G., Tang, X., Wind,… find
  30. Peng, H., Pappas, N., Yogatama, D., Schwartz, R., Smith, N. A., and Kong, L. Random feature attention. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. find
  31. OpenReview.net, 2021. URL https://openreview.net/forum?id=QtTKTdVrFBB. ↗
  32. Press, O., Smith, N., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=R8sQPpGCv0. ↗
  33. Qin, Z., Han, X., Sun, W., Li, D., Kong, L., Barnes, N., and Zhong, Y. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 7025–7041, Abu Dhabi, United Arab Emirates, December 2022a. Association for Computational Linguistics…
  34. Qin, Z., Sun, W., Deng, H., Li, D., Wei, Y., Lv, B., Yan, J., Kong, L., and Zhong, Y. cosformer: Rethinking softmax in attention. In International Conference on Learning Representations, 2022b. URL https://openreview.net/forum?id=Bl8CQrx2Up4. ↗
  35. Qin, Z., Han, X., Sun, W., He, B., Li, D., Li, D., Dai, Y., Kong, L., and Zhong, Y. Toeplitz neural network for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023a. URL https://openreview.net/forum?id=IxmWsm4xrua. ↗
  36. Qin, Z., Li, D., Sun, W., Sun, W., Shen, X., Han, X., Wei, Y., Lv, B., Yuan, F., Luo, X., et al. Scaling transnormer to 175 billion parameters. arXiv preprint arXiv:2307.14995, 2023b.
  37. Qin, Z., Sun, W., Lu, K., Deng, H., Li, D., Han, X., Dai, Y., Kong, L., and Zhong, Y. Linearized relative positional encoding. Transactions on Machine Learning Research, 2023c. find
  38. Qin, Z., Yang, S., and Zhong, Y. Hierarchically gated recurrent neural network for sequence modeling. In NeurIPS, 2023d. find
  39. Qin, Z., Zhong, Y., and Deng, H. Exploring transformer extrapolation. In Proceedings of the AAAI Conference on Artificial Intelligence, 2024. find
  40. Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, I. Learning transferable visual models from natural language supervision. In arXiv, 2021.
  41. Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale, 2019. find
  42. Sap, M., Rashkin, H., Chen, D., LeBras, R., and Choi, Y. Socialiqa: Commonsense reasoning about social interactions, 2019.
  43. Shen, X., Li, D., Zhou, J., Qin, Z., He, B., Han, X., Li, A., Dai, Y., Kong, L., Wang, M., Qiao, Y., and Zhong, Y. Fine-grained audible video description. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10585–10596, June 2023. find
  44. Su, J., Lu, Y., Pan, S., Wen, B., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864, 2021.
  45. Sun, W., Qin, Z., Deng, H., Wang, J., Zhang, Y., Zhang, K., Barnes, N., Birchfield, S., Kong, L., and Zhong, Y. Vicinity vision transformer. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(10):12635–12649, 2023a. doi: 10.1109/TPAMI.2023.3285569. — [doi:10.1109/tpami.2023.3285569](https://doi.org/10.1109/tpami.2023.3285569) find
  46. Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F. Retentive network: A successor to transformer for large language models, 2023b.
  47. Tillet, P., Kung, H.-T., and Cox, D. D. Triton: an intermediate language and compiler for tiled neural network computations. Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, 2019.
  48. Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lample, G. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023a.
  49. Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshorn, A., Hossei… — Llama 2: Open Foundation and Fine-Tuned Chat Models
  50. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems, 30, 2017.
  51. Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Efficient streaming language models with attention sinks, 2023. find
  52. Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. Gated linear attention transformers with hardware-efficient training, 2023. find
  53. Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can a machine really finish your sentence?, 2019. find
  54. Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., Mihaylov, T., Ott, M., Shleifer, S., Shuster, K., Simig, D., Koura, P. S., Sridhar, A., Wang, T., and Zettlemoyer, L. Opt: Open pre-trained transformer language models, 2022.
  55. Zheng, L., Wang, C., and Kong, L. Linear complexity randomized self-attention mechanism. In International Conference on Machine Learning, pp. 27011–27041. PMLR, 2022. find
  56. Zheng, L., Yuan, J., Wang, C., and Kong, L. Efficient attention via control variates. In International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=G-uNfHKrj46. ↗
  57. Zhou, J., Shen, X., Wang, J., Zhang, J., Sun, W., Zhang, J., Birchfield, S., Guo, D., Kong, L., Wang, M., and Zhong, Y. Audio-visual segmentation with semantics, 2023. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/qin-2024-lightning-attention-2-free/