# PolySketchFormer: Fast Transformers via Sketching Polynomial Kernels

arXiv: https://arxiv.org/abs/2310.01655 · HTML: https://mlsanity.com/paper/2310.01655/

## Key facts

- Cited by 36 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 51 references that are in the map.
- First posted on arXiv 2 Oct 2023.

## Builds on (most cited first, 30 of 34)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [PIQA: Reasoning about Physical Commonsense in Natural Language](https://mlsanity.com/paper/1911.11641/) (2019)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [PaLM 2 Technical Report](https://mlsanity.com/paper/2305.10403/) (2023)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Compressive Transformers for Long-Range Sequence Modelling](https://mlsanity.com/paper/1911.05507/) (2019)
- [Language Modeling with Gated Convolutional Networks](https://mlsanity.com/paper/1612.08083/) (2016)
- [Efficient Content-Based Sparse Attention with Routing Transformers](https://mlsanity.com/paper/2003.05997/) (2020)
- [Random Feature Attention](https://mlsanity.com/paper/2103.02143/) (2021)
- [Sketching as a Tool for Numerical Linear Algebra](https://mlsanity.com/paper/1411.4357/) (2014)
- [Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel](https://mlsanity.com/paper/1908.11775/) (2019)
- [LongNet: Scaling Transformers to 1,000,000,000 Tokens](https://mlsanity.com/paper/2307.02486/) (2023)
- [Transformer Quality in Linear Time](https://mlsanity.com/paper/2202.10447/) (2022)
- [Fast Attention Requires Bounded Entries](https://mlsanity.com/paper/2302.13214/) (2023)
- [Finetuning Pretrained Transformers into RNNs](https://mlsanity.com/paper/2103.13076/) (2021)

## Built on by (most cited first, 30 of 36)

- [Titans: Learning to Memorize at Test Time](https://mlsanity.com/paper/2501.00663/) (2024)
- [A Survey of Resource-efficient LLM and Multimodal Foundation Models](https://mlsanity.com/paper/2401.08092/) (2024)
- [Transformers, parallel computation, and logarithmic depth](https://mlsanity.com/paper/2402.09268/) (2024)
- [Nested Learning: The Illusion of Deep Learning Architectures](https://mlsanity.com/paper/2512.24695/) (2025)
- [ATLAS: Learning to Optimally Memorize the Context at Test Time](https://mlsanity.com/paper/2505.23735/) (2025)
- [Log-Linear Attention](https://mlsanity.com/paper/2506.04761/) (2025)
- [ThunderKittens: Simple, Fast, and Adorable AI Kernels](https://mlsanity.com/paper/2410.20399/) (2024)
- [Conv-Basis: A New Paradigm for Efficient Attention Inference and Gradient Computation in Transformers](https://mlsanity.com/paper/2405.05219/) (2024)
- [PolaFormer: Polarity-aware Linear Attention for Vision Transformers](https://mlsanity.com/paper/2501.15061/) (2025)
- [A Survey: Collaborative Hardware and Software Design in the Era of Large Language Models](https://mlsanity.com/paper/2410.07265/) (2024)
- [Multi-Layer Transformers Gradient Can be Approximated in Almost Linear Time](https://mlsanity.com/paper/2408.13233/) (2024)
- [Lattice: Learning to Efficiently Compress the Memory](https://mlsanity.com/paper/2504.05646/) (2025)
- [VGG-T3: Offline Feed-Forward 3D Reconstruction at Scale](https://mlsanity.com/paper/2602.23361/) (2026)
- [Exploring the Frontiers of Softmax: Provable Optimization, Applications in Diffusion Model, and Beyond](https://mlsanity.com/paper/2405.03251/) (2024)
- [Unmasking Transformers: A Theoretical Approach to Data Recovery via Attention Weights](https://mlsanity.com/paper/2310.12462/) (2023)
- [Towards Infinite-Long Prefix in Transformer](https://mlsanity.com/paper/2406.14036/) (2024)
- [Selective Rotary Position Embedding](https://mlsanity.com/paper/2511.17388/) (2025)
- [A Theoretical Insight into Attack and Defense of Gradient Leakage in Transformer](https://mlsanity.com/paper/2311.13624/) (2023)
- [Best of Both Worlds: Advantages of Hybrid Graph Sequence Models](https://mlsanity.com/paper/2411.15671/) (2024)
- [Subquadratic Algorithms and Hardness for Attention with Any Temperature](https://mlsanity.com/paper/2505.14840/) (2025)
- [Trellis: Learning to Compress Key-Value Memory in Attention Models](https://mlsanity.com/paper/2512.23852/) (2025)
- [Unlocking the Theory Behind Scaling 1-Bit Neural Networks](https://mlsanity.com/paper/2411.01663/) (2024)
- [Auto-Regressive Masked Diffusion Models](https://mlsanity.com/paper/2601.16971/) (2026)
- [LevAttention: Time, Space, and Streaming Efficient Algorithm for Heavy Attentions](https://mlsanity.com/paper/2410.05462/) (2024)
- [Fast attention mechanisms: a tale of parallelism](https://mlsanity.com/paper/2509.09001/) (2025)
- [Two Heads Are Better than One: Simulating Large Transformers with Small Ones](https://mlsanity.com/paper/2506.12220/) (2025)
- [Sublinear Time Quantum Algorithm for Attention Approximation](https://mlsanity.com/paper/2602.00874/) (2026)
- [Every Bit Counts: A Theoretical Study of Precision-Expressivity Tradeoffs in Quantized Transformers](https://mlsanity.com/paper/2602.02707/) (2026)
- [SMat-Attention: Structured Long-Context Sequence Modeling](https://mlsanity.com/paper/2609.36062/) (2026)
- [Data-Aware Random Feature Kernel for Transformers](https://mlsanity.com/paper/2603.04127/) (2026)

## Abstract (arXiv)

The quadratic time and memory complexity inherent to self-attention mechanisms, with respect to sequence length, presents a critical computational bottleneck in the training and deployment of large-scale Transformer-based language models. Recent theoretical results indicate the intractability of sub-quadratic softmax attention approximation under reasonable complexity assumptions. This paper addresses this challenge by first demonstrating that polynomial attention with high degree can effectively replace softmax without sacrificing model quality. Next, we develop polynomial sketching techniques from numerical linear algebra to achieve linear-time polynomial attention with approximation guarantees. Crucially, our approach achieves this speedup without requiring the sparsification of attention matrices. We also present a block-based algorithm to apply causal masking efficiently. Combining these techniques, we provide PolySketchFormer, a practical linear-time Transformer architecture for language modeling that offers provable guarantees. We validate PolySketchFormer empirically by training language models capable of handling long contexts. These experiments utilize both synthetic and real-world datasets (PG19, Wikipedia and C4) on Google Cloud TPUs. For context lengths of 32k and GPT-2 style models, our model achieves a 2.5-4x speedup in training compared to FlashAttention, with no observed degradation in quality across our experiments.