LaplacianFormer:Rethinking Linear Attention with Laplacian Kernel

1. Problem Statement

The standard Softmax attention mechanism in Transformers suffers from quadratic time and space complexity ($O(N^2)$) relative to sequence length, creating a bottleneck for high-resolution vision tasks and long sequences. While Linear Attention variants have been proposed to reduce complexity to $O(N)$ by replacing Softmax with kernel-based approximations, existing methods predominantly rely on Gaussian kernels.

The authors identify two critical flaws in the widespread use of Gaussian kernels:

1. Theoretical Mismatch: Gaussian kernels assume similarity decays rapidly with squared $\ell_2$ distance. Empirical analysis of Q-K (Query-Key) distances in models like DeiT, PVT, and Swin reveals that these distances follow a heavy-tailed distribution. The exponential decay of Gaussian kernels amplifies outliers, causing them to dominate attention maps while overly suppressing moderately relevant tokens.
2. Optimization Instability: The gradient of the Gaussian kernel vanishes linearly as the distance between vectors approaches zero ($\|x-y\|_2 \to 0$). This leads to vanishing updates, hindering convergence, especially in early training stages.

2. Methodology: LaplacianFormer

The paper proposes LaplacianFormer, a linear attention framework that replaces the Gaussian kernel with a Laplacian kernel to address the issues above.

A. The Laplacian Kernel

Instead of squared $\ell_2$ distance, the model uses the $\ell_1$ distance (Manhattan distance) within the kernel function:
$$k(x, y) = \exp\left(-\frac{\|x - y\|_1}{\lambda}\right)$$

• Slower Decay: The Laplacian kernel decays more slowly than the Gaussian kernel, preventing the oversuppression of mid-range token interactions.
• Non-Vanishing Gradients: Due to the piecewise linear nature of the $\ell_1$ norm, the gradient remains non-zero even when vectors are nearly identical, ensuring stable optimization and faster convergence.
• Robustness: The $\ell_1$ metric is less sensitive to outliers compared to $\ell_2$, leading to more faithful token relevance measurements.

B. Injective Feature Mapping

To prevent expressiveness degradation often seen in low-rank linear attention approximations, the authors introduce a provably injective feature map.

• They construct a normalized kernel representation $z_i$ for each query by centering and whitening the similarity vectors.
• Instead of computing the full inverse square root of the covariance matrix (which is $O(N^3)$), they use a diagonal estimator to normalize feature dimensions independently across the batch. This preserves fine-grained token distinctions while maintaining linear complexity.

C. Efficient Computation & Solvers

To handle the computational cost of kernel inversion and matrix operations efficiently:

1. Nyström Approximation: The full kernel matrix is approximated using a small set of landmark keys selected via average pooling. This reduces the rank of the matrix significantly.
2. Newton–Schulz Iteration: To compute the inverse of the landmark kernel matrix (required for the attention mechanism), the authors avoid costly SVD or direct inversion. Instead, they use Newton–Schulz iteration, a method relying solely on matrix multiplications and additions, which is highly GPU-friendly.
   • A small diagonal perturbation ($\epsilon I$) is added to ensure the matrix is strictly positive definite for convergence.
3. CUDA Acceleration: Custom CUDA kernels are implemented for both the Laplacian distance computation and the Newton–Schulz solver, optimizing memory access and register reuse for high-throughput forward and backward passes.

3. Key Contributions

1. Theoretical Insight: Demonstrated that the standard Gaussian kernel is suboptimal for vision Transformers due to heavy-tailed Q-K distance distributions and vanishing gradients, proposing the Laplacian kernel as a principled alternative.
2. Novel Architecture: Introduced LaplacianFormer, a scalable linear attention model that integrates an injective feature map with Laplacian similarity, ensuring fine-grained token discrimination.
3. Algorithmic Efficiency: Developed a Newton–Schulz-based inverse solver combined with Nyström approximation, avoiding expensive matrix inversions while guaranteeing convergence.
4. System Optimization: Implemented custom CUDA kernels for both the kernel computation and the solver, enabling efficient deployment on edge devices and high-resolution inputs.

4. Experimental Results

The model was evaluated on ImageNet-1K, Object Detection (COCO), and Instance Segmentation.

• ImageNet Classification:

  • LaplacianFormer achieves state-of-the-art (SOTA) Top-1 accuracy across all FLOPs ranges (from Tiny to Huge models).
  • Example: The LaplacianFormer-Medium (46.3M params, 7.43G FLOPs) achieves 85.3% Top-1 accuracy, outperforming SOFT++-Large (84.1%) and BiFormer-B (84.3%) despite having fewer parameters and FLOPs.
  • Convergence: Ablation studies show models with Laplacian kernels converge faster and achieve higher final accuracy than their Gaussian counterparts (e.g., in SOFT++ and Skyformer baselines).

• Downstream Tasks (COCO):

  • Object Detection & Segmentation: Using Mask R-CNN and RetinaNet backbones, LaplacianFormer consistently outperforms existing backbones (Swin, PVT, SOFT++).
  • Example: LaplacianFormer-Tiny achieves 43.2 AP (bounding box) and 40.3 AP (mask) on Mask R-CNN, surpassing SOFT++-Tiny and FL-PVT-T.

• Efficiency Analysis:

  • Memory Scaling: Shows linear memory scaling with sequence length, matching efficient Transformers like Performer and SOFT, and far outperforming vanilla Transformers.
  • Solver Robustness: Newton–Schulz iteration demonstrates superior stability and convergence compared to Conjugate Gradient (CG) solvers, especially under ill-conditioned scenarios.

5. Significance

This work challenges the "default" convention of using Gaussian kernels in linear attention. By grounding the design in empirical distribution analysis and theoretical gradient behavior, LaplacianFormer offers a more robust and expressive alternative. It successfully bridges the gap between theoretical efficiency (linear complexity) and practical performance (high accuracy and stable training), making it a strong candidate for deploying high-resolution vision Transformers on resource-constrained edge devices. The custom CUDA implementation further ensures that these theoretical gains translate into real-world inference speedups.