Research map / Architectures and efficiency / Attention and transformers
PolySketchFormer: Fast Transformers via Sketching Polynomial Kernels
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PolySketchFormer, a linear-time Transformer that uses high-degree polynomial attention approximated by non-negative polynomial sketches (self-tensoring), a block-based causal lower-triangular multiplication, and optional learned sketches plus exact local attention, trained for language modeling.
Problem
Quadratic time and memory cost of softmax self-attention limits training of Transformer language models on long contexts, and existing linear attention approximations (e.g., Performer) lose quality or lack guarantees.
Approach
Replace softmax with normalized degree-p polynomial attention (p>=4 matches softmax quality), approximate it with a polynomial sketch followed by self-tensoring so attention weights are provably non-negative, make sketches learnable, and compute causal masking with a block-based prefix-sum scheme; local blocks use exact polynomial attention.
Key results
- At 32k context on GPT-2 small style models, Polysketch (learned + local, r=32) trains about 2x faster than the fastest FlashAttention configuration (block 512) on 32 TPUs.
- On PG-19 at 32k context, Polysketch (learned + local, r=64) reaches test perplexity 11.34 versus 11.55 for softmax (FlashAttention).
- Polynomial attention with degree 4 or 8 matches softmax quality on C4 perplexity and Q/A accuracies (e.g., degree 8 C4 perplexity 17.77 vs softmax 17.81 for GPT-2 small at 8k), but OOMs beyond 8k context.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use learned sketches with local exact polynomial attention instead of softmax attention | GPT-2 Small style language model · PG19 test, 32k context | perplexity: 11.55 -> 11.34 | Helped |
| use learned sketches with local exact polynomial attention instead of softmax attention | GPT-2 Small style language model · Wiki-40B test, 32k context | perplexity: 14.35 -> 14.29 | Helped |
| use Polysketch (learned + local, r=32) instead of FlashAttention | GPT-2 Small style language model · Training, 32k context | speedup: 2x | Helped |
| add local exact polynomial attention (ablation) | Polysketch (learned, r=64) · PG19 test, 32k context | perplexity: 12.19 -> 11.34 | Helped |
| replace random sketches with learned sketches, retaining local exact polynomial attention (ablation) | Polysketch, r=64 · PG19 test, 32k context | perplexity: 11.91 -> 11.34 | Helped |
| increase sketch size from r=32 to r=64 (ablation) | Polysketch (learned + local) · PG19 test, 32k context | perplexity: 11.47 -> 11.34 | Helped |
| remove local exact polynomial attention and use random sketches (ablation) | Polysketch (learned + local, r=64) · PG19 test, 32k context | perplexity: 11.34 -> 13.66 | Hurt |
| use a different learning rate schedule (ablation) | Polysketch (learned + local) · Selective Copying, 16k context | accuracy: 92.75% -> 99.44% | Helped |
| replace softmax attention with degree-4 polynomial attention | GPT-2 Small style language model · PG19 test, 8k context | perplexity: 11.65 -> 11.85 | No clear effect |
| increase polynomial attention degree from 4 to 8 (ablation) | GPT-2 Small style language model · PG19 test, 8k context | perplexity: 11.85 -> 11.64 | Helped |
| increase context length to 32k | Softmax and non-kernel polynomial attention · PG19 test, 32k context | perplexity | Mixed |
| use a 13-layer Polysketch model rather than 12 layers (ablation) | Polysketch (learned + local, r=64) · C4 validation, 8192 context, 125k steps | perplexity: 17.68 -> 17.98 | Hurt |
Surprising findings
- Polynomial attention with degree p>=4 performs comparably to softmax attention despite lacking softmax's normalization structure.
- Learned sketches (a small learnable dense network replacing random Gaussian projections) substantially improve perplexity over random sketches.
Limitations
- Speedups are measured on TPUs with a JAX implementation and without I/O-aware optimizations; gains may change with better kernels.
- Evaluated only for decoder-only language models; encoder or vision settings are unexplored.
- Synthetic selective copying results show Polysketch can be sensitive to learning-rate schedule at 16k context.
Experimental setup
Datasets: PG-19 Wiki-40B C4 HellaSwag PIQA Physics Selective Copying Induction heads
Benchmarks: PG-19 Wiki-40B C4 HellaSwag PIQA Physics
Metrics: perplexity 0-shot accuracy 5-shot accuracy training steps/sec selective copying accuracy
Baselines: Softmax attention (FlashAttention) Performer Polynomial attention FlashAttention block 256 FlashAttention block 512
Builds on: Performer FlashAttention Transformer++ polynomial sketching (Ahle et al.)
Closest papers by meaning
- Learning Linear Attention in Polynomial Time (NeurIPS 2025 · similarity 0.88)
- Beyond Linear Approximations: A Novel Pruning Approach for Attention Matrix (ICLR 2025 · similarity 0.88)
- MemoryFormer : Minimize Transformer Computation by Removing Fully-Connected Layers (NeurIPS 2024 · similarity 0.88)
- Polynomial Composition Activations: Unleashing the Dynamics of Large Language Models (ICLR 2025 · similarity 0.87)
- MonarchAttention: Zero-Shot Conversion to Fast, Hardware-Aware Structured Attention (NeurIPS 2025 · similarity 0.87)
- Gated Linear Attention Transformers with Hardware-Efficient Training (ICML 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.