Dodging the Quadratic Tax on Long Transformer Sequences

What it does The library provides drop-in attention replacements for PyTorch transformer encoders. It bundles the authors’ own research implementations—linear attention and clustered attention—alongside methods from other papers like Linformer and Reformer. You pick an attention_type such as "full" or "linear" via a builder and get a model that replaces the standard softmax attention with a cheaper approximation.

The interesting bit The README shows a concrete payoff: for a 1,000-token sequence on a GTX 1080 Ti, linear attention clocks in at 68 ms versus 144 ms for full softmax. That is not just a theoretical complexity reduction; it is a measurable speedup on mid-range hardware. The library essentially acts as a research testbed where proven papers become toggleable builder arguments.

Key highlights

• Implements both the authors’ own methods (linear and clustered attention) and third-party alternatives (Reformer, Linformer, Efficient Attention).
• Exposes a single TransformerEncoderBuilder that swaps attention mechanisms with a string argument.
• Ships with custom CUDA kernels, so GPU execution is first-class rather than emulated.
• MIT licensed, which is about as permissive as it gets.

Caveats

• Compiling for GPU requires a C++ and CUDA toolchain, and macOS builds specifically depend on llvm and libomp, so it is not a guaranteed one-click install on every platform.
• The documentation and tone are research-first; long-term maintenance and production-hardening are not explicitly discussed.

Verdict Worth a look if you are training or experimenting with long-sequence transformers in PyTorch and want to benchmark attention approximations without reimplementing NeurIPS papers yourself. Skip it if you are after a polished, high-level production framework rather than a research testbed.