SOFT: Softmax-free Transformer with Linear Complexity

Jiachen Lu; Jinghan Yao; Junge Zhang; Xiatian Zhu; Hang Xu; Weiguo Gao; Chunjing Xu; Tao Xiang; Li Zhang

Source-linked AI summary

SOFT: Softmax-free Transformer with Linear Complexity

Jiachen Lu, Jinghan Yao, Junge Zhang, Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, Li Zhang

arXiv:2110.11945v3Oct 22, 2021cs.CVcs.AIcs.LG

TL;DR

Vision Transformers face quadratic computation and memory costs as image-token sequences grow, while existing linear approximations retain softmax-related limitations. SOFT replaces softmax attention with Gaussian-kernel similarity and low-rank decomposition, achieving linear complexity and a superior accuracy–complexity trade-off on ImageNet.

Problem
Vision Transformer self-attention has quadratic computation and memory complexity, while prior linear approximations can be theoretically flawed or ineffective for visual recognition.
Method
SOFT uses unnormalized Gaussian-kernel attention, low-rank matrix decomposition, and a Newton–Raphson computation of the Moore–Penrose inverse.
Results
SOFT achieves linear space and time complexity and a superior accuracy–complexity trade-off in extensive ImageNet experiments.
Takeaways & Limitations
SOFT permits much longer image-token sequences while maintaining an improved accuracy–complexity trade-off for visual recognition.
Takeaways & Limitations
SOFT sets the query and key projection matrices identically to preserve a symmetric, positive-semidefinite attention matrix required for its Nyström approximation.

Abstract

from arXiv · show

Vision transformers (ViTs) have pushed the state-of-the-art for various visual recognition tasks by patch-wise image tokenization followed by self-attention. However, the employment of self-attention modules results in a quadratic complexity in both computation and memory usage. Various attempts on approximating the self-attention computation with linear complexity have been made in Natural Language Processing. However, an in-depth analysis in this work shows that they are either theoretically flawed or empirically ineffective for visual recognition. We further identify that their limitations are rooted in keeping the softmax self-attention during approximations. Specifically, conventional self-attention is computed by normalizing the scaled dot-product between token feature vectors. Keeping this softmax operation challenges any subsequent linearization efforts. Based on this insight, for the first time, a softmax-free transformer or SOFT is proposed. To remove softmax in self-attention, Gaussian kernel function is used to replace the dot-product similarity without further normalization. This enables a full self-attention matrix to be approximated via a low-rank matrix decomposition. The robustness of the approximation is achieved by calculating its Moore-Penrose inverse using a Newton-Raphson method. Extensive experiments on ImageNet show that our SOFT significantly improves the computational efficiency of existing ViT variants. Crucially, with a linear complexity, much longer token sequences are permitted in SOFT, resulting in superior trade-off between accuracy and complexity.

1 Introduction

Vision Transformers improve visual recognition but self-attention incurs quadratic computation and memory costs, especially for long image-token sequences. SOFT removes softmax and uses Gaussian-kernel attention with low-rank approximation to achieve linear complexity and improved accuracy–complexity trade-offs.

1 Introduction: Self-attention requires O(n^2) computation and memory because it constructs an n × n attention matrix over token pairs.The cost is particularly acute for vision because image tokenization produces longer sequences than typical NLP inputs.
1 Introduction: The reported ImageNet comparison evaluates Top1-Accuracy against parameter count and practical memory usage for CNN and Transformer models.Memory is measured with batch size 1 on a 16GB Tesla V100.
1 Introduction: Existing efficient-Transformer approaches approximate self-attention through projection, kernel, hashing, or matrix-decomposition strategies.These approaches are introduced as attempts to reduce self-attention complexity in NLP and vision Transformers.
1 Introduction: SOFT replaces softmax-normalized dot-product attention with Gaussian-kernel similarity and approximates the resulting attention matrix using low-rank decomposition.The method computes the Moore–Penrose inverse with Newton–Raphson iteration to provide robust approximation.
1 Introduction: SOFT provides linear O(n) space and time complexity, allowing much longer image-token sequences than conventional approaches.The paper reports improved accuracy/complexity trade-offs at comparable model sizes.

2 Related work

Related work establishes that Transformer self-attention scales quadratically with sequence length, limiting high-resolution vision models. Efficient variants reduce this cost through projections, hashing, kernels, or decomposition, but SOFT targets softmax-related theoretical limitations.

2 Related work: Vision Transformers inherit quadratic space and time complexity from exhaustive pairwise token comparisons in self-attention.For higher-resolution images, controlling this cost requires larger patches, which sacrifices spatial resolution and fine-grained representations.
2 Related work: Efficient Transformer methods use low-rank Key–Value projections, hashing, kernel approximations, random feature mappings, or reduced attention matrices.These strategies are presented as alternatives for lowering sequence-length-dependent computation and memory.
2 Related work: Nyströmformer decomposes the attention matrix but applies softmax to ingredient matrices, which lacks a theoretical guarantee for approximating standard self-attention.Its polynomial iteration also lacks a guarantee for inversion when the matrix is nearly singular.
2 Related work: SOFT removes softmax normalization so that matrix decomposition can be applied with theoretical guarantees.This distinguishes its approach from efficient Transformers that retain softmax-based attention during approximation.

3 Method

SOFT replaces softmax dot-product attention with an unnormalized Gaussian-kernel matrix, then uses low-rank Nyström decomposition and Newton–Raphson inversion to obtain linear-complexity attention.

Softmax-free self-attention formulation: SOFT replaces softmax dot-product similarity with an unnormalized Gaussian kernel, producing a symmetric self-attention matrix.The query and key projections are tied so Q = K, preserving symmetry.
Softmax-free self-attention formulation: The Gaussian-kernel attention matrix is symmetric, lies in [0, 1], has unit diagonal entries, and is positive semidefinite as a Gram matrix.Without linearization, however, training with this matrix fails to converge.
Low-rank regularization via matrix decomposition with linear complexity: SOFT approximates the full attention matrix using Nyström decomposition with m ≪ n bottleneck tokens, avoiding explicit computation of the n × n matrix.The approximation uses sampled-token correlations and the Moore–Penrose inverse of their correlation matrix.
Low-rank regularization via matrix decomposition with linear complexity: Newton–Raphson iteration computes the Moore–Penrose inverse, with Ak+1 = 2Ak − AkAAk converging to A† when α is sufficiently small.The paper states that the relevant residuals decrease monotonically under this condition.
Low-rank regularization via matrix decomposition with linear complexity: SOFT has linear space and time complexity in the token count, with total time O((de + 4mde + m2)n + Tm3 + dem2).The reported complexity includes sampling, decomposed-matrix construction, iterative inversion, and matrix multiplication.
SOFT architecture: The SOFT block replaces self-attention in Transformer layers and is integrated into a pyramidal visual backbone with modified overlapping-convolution patch embedding.The resulting architecture is designed for general image-recognition tasks.

4 Experiments

Experiments evaluate SOFT against efficient Transformers and vision backbones on ImageNet and Long Range Arena, including architectural, sampling, sequence-length, and convergence ablations. SOFT generally improves accuracy–efficiency trade-offs while supporting longer token sequences and robust attention modeling.

Comparison with existing linear Transformers: SOFT achieves the best classification accuracy among the compared linearization methods while substantially reducing memory and FLOPs relative to the standard Transformer.Inference speed is on-par with other linear Transformers, although training is slightly slower than Nyströmformer and slower than Performer and Linformer.
Comparison with CNNs and ViTs: SOFT outperforms PVT across all variants and beats Twins with fewer parameters and fewer floating-point operations on ImageNet-1K.The comparison reports top-1 accuracy on the ImageNet-1K validation set and FLOPs at batch size 1.
Attention visualization: SOFT exhibits robustness and versatility in capturing both local and long-distance pixel relations, including shared semantic concepts and instance-specific features.Figure 3 compares attention heatmaps from Transformer, Performer, Nyströmformer, and SOFT.
Ablation studies: m = 49 provides the best trade-off between performance and computational overhead, while longer bottleneck sequences increase memory use and computation.The ablation varies the bottleneck length across 36, 49, 64, and 81 tokens; memory is measured with batch size 128.
Ablation studies: Average pooling yields the best performance with less computational overhead than convolution, supporting its use for token sampling in SOFT.Random sampling is sensitive to m, while biased sampling may miss salient samples and random sampling lacks a uniformity guarantee.
Ablation studies: The Moore-Penrose inverse approximation converges within 20 Newton-Raphson iterations across all stages of SOFT-Tiny.Convergence is measured using the p = 2 norm metric described in Proposition 1.

5 Conclusions

SOFT removes softmax normalization from self-attention by using Gaussian-kernel attention, enabling low-rank approximation of the full attention matrix and linear space-time complexity.

SOFT is a softmax-free self-attention mechanism for linearizing Transformer complexity in space and time.
Gaussian-kernel attention eliminates softmax normalization and enables low-rank approximation of the full self-attention matrix.
SOFT computes the Moore-Penrose inverse with Newton-Raphson iterations to improve approximation robustness.
SOFT yields a superior trade-off between accuracy and complexity in extensive experiments.

A.1 Nyström method

The Nyström construction approximates the Transformer attention Gram matrix through Gaussian-kernel eigenfunction approximation and a rank-m spectral decomposition.

Transformer self-attention can be viewed as a Gram matrix S formed with a Gaussian kernel applied to query features.
Gaussian kernel evaluations can be represented in a feature space using eigenvalues λ_i and eigenfunctions φ_i.
Eigenfunction approximations are obtained by sampling x_1, x_2, · · ·, x_q from the probability distribution p(x).
A submatrix S(m) containing m × m elements of S is spectrally decomposed using column-orthogonal U(m) and diagonal Λ(m).
The resulting approximation ˜S has rank m and uses approximations ˜U and ˜Λ for the spectral decomposition of S.

A.2.1 Proof of theorem 1

The proof establishes convergence of the Newton-Raphson iteration to the Moore-Penrose inverse under a sufficiently small initialization scale.

The iteration starts from A_0 = αA, where A is symmetric positive semi-definite with bounded entries and unit diagonal.
Newton-Raphson updates satisfy A_{k+1} = A_k(2I_n − AA_k) = (2I_n − A_kA)A_k.
Choosing α sufficiently small ensures the initial condition ∥AA† − AA_0∥ < 1.
Under this initialization, ∥AA† − AA_k∥ tends to zero as k tends to infinity, implying A_k tends to A†.

A.2.2 Proof of proposition 1

The proposition proof shows monotonic convergence for one approximation error, while explicitly noting that another residual lacks a corresponding monotonicity guarantee.

Multiplying the iteration relation by A yields the proof relation used to analyze convergence.
The error ∥A − AA_kA∥ decreases monotonically to zero, and ∥A_k − A†∥ also decreases monotonically to zero.
The residual ∥A_kAA_k − A_k∥ is not proved to decrease monotonically to zero.

A.3 Non-linearized gaussian kernel attention

SOFT approximates Gaussian-kernel attention through sampled bottleneck tokens, producing a low-rank attention matrix. Visualizations indicate more diverse attention heads and more precise capture of multiple objects at patch boundaries.

A.3 Non-linearized gaussian kernel attention: Sampled bottleneck tokens reconstruct token relations and produce a low-rank attention matrix when their count is much smaller than the sequence length.The bottleneck tokens also focus attentive learning on latent salient information.
A.3 Non-linearized gaussian kernel attention: The approximated Gram matrix performs better than a directly estimated matrix, although the supplied passage does not quantify the difference.
A.3 Non-linearized gaussian kernel attention: SOFT exhibits greater semantic diversity across attention heads than the other visualized Transformer methods.The comparison includes Transformer, Performer, and Nyströmformer.
A.3 Non-linearized gaussian kernel attention: At image patches spanning multiple object boundaries, SOFT more precisely captures all those objects than the other visualized methods.