Discrete Cosine Transform Based Decorrelated Attention for Vision Transformers

TL;DR

This work tackles the challenge of training efficiency and accuracy for self-attention in Vision Transformers by introducing two Discrete Cosine Transform (DCT)–based strategies. First, it initializes one of the projection weight matrices ($\mathbf{W}_Q$, $\mathbf{W}_K$, or $\mathbf{W}_V$) with a DCT basis to ensure full-spectrum representation, improving convergence and accuracy. Second, it implements a DCT-based compression that truncates high-frequency channel components (controlled by $\tau$) to reduce parameter counts and FLOPs while maintaining performance. Across CIFAR-10, ImageNet-1K, and COCO, the methods yield accuracy gains with notable efficiency improvements (parameter and compute reductions) and minimal or no loss in accuracy, demonstrating practical benefits for Swin and ViT backbones. The results underscore the value of frequency-domain techniques for improving both the learning dynamics and the efficiency of vision transformers.

Abstract

Central to the Transformer architectures' effectiveness is the self-attention mechanism, a function that maps queries, keys, and values into a high-dimensional vector space. However, training the attention weights of queries, keys, and values is non-trivial from a state of random initialization. In this paper, we propose two methods. (i) We first address the initialization problem of Vision Transformers by introducing a simple, yet highly innovative, initialization approach utilizing discrete cosine transform (DCT) coefficients. Our proposed DCT-based \textit{attention} initialization marks a significant gain compared to traditional initialization strategies; offering a robust foundation for the attention mechanism. Our experiments reveal that the DCT-based initialization enhances the accuracy of Vision Transformers in classification tasks. (ii) We also recognize that since DCT effectively decorrelates image information in the frequency domain, this decorrelation is useful for compression because it allows the quantization step to discard many of the higher-frequency components. Based on this observation, we propose a novel DCT-based compression technique for the attention function of Vision Transformers. Since high-frequency DCT coefficients usually correspond to noise, we truncate the high-frequency DCT components of the input patches. Our DCT-based compression reduces the size of weight matrices for queries, keys, and values. While maintaining the same level of accuracy, our DCT compressed Swin Transformers obtain a considerable decrease in the computational overhead.

Figures (3)

• Figure 1: Multi-head Self-Attention. In (b), to reduce the computational overhead, IDCT and the final linear layer $\mathcal{L}_O$ are implemented as a single linear layer as no non-linearity exists in between.
• Figure 2: Basis vectors of an $8\times 8$ DCT matrix where $D_i$ refers to the $i$-th column. The DCT basis vectors provide a good approximation to the eigenvectors of the Teoplitz matrix with $\rho=0.9$ (KLT).
• Figure 3: Frequency response of the DCT basis vectors. Each frequency response is shown in a solid curve in different colors. The frequency is normalized to $[0, \pi]$. The black dashed line denotes the root of the squared sum of these frequency responses, showing the entire spectrum is covered.