LaplacianFormer:Rethinking Linear Attention with Laplacian Kernel

TL;DR

LaplacianFormer introduces a Laplacian kernel-based Transformer that improves efficiency and expressiveness for high-resolution vision tasks, with theoretical grounding and practical CUDA implementations.

Contribution

It proposes a Laplacian kernel as a principled alternative to softmax in Transformers, along with a provably injective feature map and efficient Nyström approximation for better performance.

Findings

Achieves strong performance-efficiency trade-offs on ImageNet.

Retains fine-grained token information with a new injective feature map.

Utilizes efficient CUDA implementations for high-throughput processing.

Abstract

The quadratic complexity of softmax attention presents a major obstacle for scaling Transformers to high-resolution vision tasks. Existing linear attention variants often replace the softmax with Gaussian kernels to reduce complexity, but such approximations lack theoretical grounding and tend to oversuppress mid-range token interactions. We propose LaplacianFormer, a Transformer variant that employs a Laplacian kernel as a principled alternative to softmax, motivated by empirical observations and theoretical analysis. To address expressiveness degradation under low-rank approximations, we introduce a provably injective feature map that retains fine-grained token information. For efficient computation, we adopt a Nystr\"om approximation of the kernel matrix and solve the resulting system using Newton--Schulz iteration, avoiding costly matrix inversion and SVD. We further develop custom CUDA implementations for both the kernel and solver, enabling high-throughput forward and backward passes suitable for edge deployment. Experiments on ImageNet show that LaplacianFormer achieves strong performance-efficiency trade-offs while improving attention expressiveness.