Fast Multipole Attention: a physics-inspired multilevel attention that cuts attention cost to O(n log n) or O(n)

October 18, 20238 min

Overview

Decision SnapshotReady For Pilot

FMA shows consistent gains across language and vision benchmarks and includes a practical CUDA/TVM implementation, but full production hardening and broader community replication are pending.

Citations0

Evidence Strength0.80

Confidence0.85

Risk Signals7

Trust Signals

Findings with numeric evidence: 6/6

Findings with evidence refs: 6/6

Results with explicit delta: 5/5

Reproducibility

Status: Code + data available

Open source: Partial

At A Glance

Cost impact: 80%

Production readiness: 70%

Novelty: 70%

Authors

Yanming Kang, Giang Tran, Hans De Sterck

Links

Abstract / PDF / Code

Why It Matters For Business

FMA lowers GPU memory and inference latency for long text and high-resolution images, letting teams train bigger models or use longer contexts without buying more hardware.

Who Should Care

Summary TLDR

This paper introduces Fast Multipole Attention (FMA), a learned multilevel attention that keeps exact attention in a small local neighborhood and summarizes distant context with learned group summaries. FMA reduces self-attention cost from quadratic to O(n log n) and, with query downsampling, to O(n). The authors provide 1D and 2D implementations and show better accuracy-than-other-efficient methods on long-context language (enwik8, WikiText-103) and higher accuracy + lower memory on image tasks (ImageNet, ADE20K). Code is promised on GitHub.

Problem Statement

Full self-attention gives a global receptive field but costs O(n^2) time and memory, which blocks training or inference on long text and high-resolution images. The paper aims to keep global context while reducing cost to practical levels.

Main Contribution

Fast Multipole Attention (FMA): a learned multilevel attention that keeps exact attention in a fixed local window and aggregates distant tokens via learned summaries.

Theoretical complexity: O(n log n) time/memory and O(n) with query downsampling, for both 1D (text) and 2D (images).

Key Findings

FMA changes attention complexity from quadratic to log-linear or linear.

NumbersComplexity reduced from O(n^2) to O(n log n); O(n) with query downsampling

Practical UseYou can process much longer documents or larger images without quadratic memory blowup; enable contexts that were previously infeasible on the same hardware.

Evidence RefAbstract; Sec.3

On character-level enwik8 with context 4096, FMA achieves best efficient-model bpc.

NumbersFMA bpc=1.133 at n=4096; Reformer 1.180

Practical UseFor long-context autoregressive language models, FMA gives lower test loss than prior efficient attention methods while using less or comparable memory.

Evidence RefTable 1

Results

MetricValueBaselineDeltaSplit / DatasetEvidenceEvidence Ref
bits-per-character (bpc)1.133 (FMA, n=4096)Reformer 1.180 (n=4096)-0.047enwik8 (context 4096)Table 1 enwik8 resultsTable 1
perplexity (ppl)8.95 (FMA, n=2048)Full attention 8.70 (n=2048)+0.25WikiText-103 (context 2048)Table 2 masked LM resultsTable 2

What To Try In 7 Days

Replace standard self-attention with FMA in one Transformer layer to measure memory and accuracy changes.

Run a short language eval (e.g., WikiText-103 subset) with FMA and FMA-linear to compare perplexity and peak memory.

For vision, swap attention blocks in a ViT backbone with FMA2D on a small ImageNet subset to check top-1 and GPU footprint.

Optimization Features

Token Efficiency
Multilevel summarization preserves local detail while compressing distant context
Infra Optimization

Lower peak GPU memory enables larger batch sizes or higher resolution inputs on the same hardware

Model Optimization
Learned low-rank group summaries (coarse rank p)Separable learned aggregation kernels for 2D groups
System Optimization
Custom TVM schedule and optimized CUDA kernels to realize theoretical speedups
Training Optimization
End-to-end learning of aggregation kernels; same training recipes as baselines
Inference Optimization
Avoids materializing off-diagonal blocks; on-the-fly generation of block contributionsQuery downsampling option to reach strict O(n) inference

Reproducibility

Code AvailableYes
Data AvailableYes
Open Source StatusPartial
LicenseUnknown

Risks & Boundaries

Limitations

Two hyperparameters require tuning (base cell size r and coarse rank p) and affect memory/accuracy trade-offs (paper studies r and p).

FMA-linear trades a modest accuracy loss for memory savings compared to the log-linear variant.

When Not To Use

For short sequences or small images where full quadratic attention fits easily and yields slightly better accuracy.

When very tight, deterministic latency bounds are required and the implementation has not been production-optimized.

Failure Modes

Too-aggressive downsampling (small p or large r) can hurt accuracy by over-compressing long-range detail.

FMA-linear may lose some accuracy compared to FMA (log-linear) as reported in language experiments.

Core Entities

Models

FMAFMA-linearFMA2D-BFMA2D-LFMA2D-B-SegFormerFMA2D-L-SegFormerH-Transformer-1DReformerLinear TransformerSwin TransformerViTSegFormer

Metrics

bits-per-character (bpc)perplexity (ppl)AccuracymIoU (%)peak GPU memory (GB)

Datasets

enwik8WikiText-103ImageNet-1KImageNet-22KADE20K

Benchmarks

character-level language modelingmasked language modelingImageNet classificationADE20K semantic segmentation