Plug-and-Play Attention Linearization for Pretrained Transformers
Abstract
We introduce LATTE (Linearized Attention with Tunable Taylor Series Expansion), a training-free attention linearization method that selectively applies linearization to a subset of self-attention blocks in a pretrained vision transformer. LATTE approximates the softmax function via a second-order Taylor expansion centered on tunable anchor points, thereby enabling a hybrid architecture that seamlessly integrates standard self-attention with linear attention. This method preserves model performance without requiring any fine-tuning or retraining but only minimal calibration on the target dataset. LATTE is highly scalable, linearizing more than half of the attention layers retains more than 95% of the original model’s performance and reduces attention GFLOPs by up to 15× (e.g., from 16.384 to 1.082 at T=1000 for the CLIP attention block). We evaluate LATTE across diverse benchmarks, including image classification, object detection, and vision–language multimodal tasks, demonstrating consistent superiority over current state-of-the-art training-free linearization methods. Code will be released here.