Towards Interpretable and Efficient Attention: Compressing All by Contracting a Few

TL;DR

Attention mechanisms excel but are often opaque and quadratically costly. This work formulates a unified optimization objective based on the $MCR^2$ coding-rate criterion that compresses input tokens toward a small set of representatives, and unrolls its gradient to yield Contract-and-Broadcast Self-Attention (CBSA), which scales linearly in $N$ for fixed $m$. By selecting different sets of representatives, CBSA recovers softmax attention, linear attention, and channel attention, presenting a principled, interpretable unification. Empirical results on visual tasks show comparable accuracy to black-box attention with significantly reduced parameters and FLOPs, plus interpretable attention maps and improved segmentation behavior, validating the interpretability-efficiency trade-off.

Abstract

Attention mechanisms have achieved significant empirical success in multiple fields, but their underlying optimization objectives remain unclear yet. Moreover, the quadratic complexity of self-attention has become increasingly prohibitive. Although interpretability and efficiency are two mutually reinforcing pursuits, prior work typically investigates them separately. In this paper, we propose a unified optimization objective that derives inherently interpretable and efficient attention mechanisms through algorithm unrolling. Precisely, we construct a gradient step of the proposed objective with a set of forward-pass operations of our \emph{Contract-and-Broadcast Self-Attention} (CBSA), which compresses input tokens towards low-dimensional structures by contracting a few representatives of them. This novel mechanism can not only scale linearly by fixing the number of representatives, but also covers the instantiations of varied attention mechanisms when using different sets of representatives. We conduct extensive experiments to demonstrate comparable performance and superior advantages over black-box attention mechanisms on visual tasks. Our work sheds light on the integration of interpretability and efficiency, as well as the unified formula of attention mechanisms.

Figures (10)

• Figure 1: Overview of CBSA.Left panel: A detailed illustration of the operations in the proposed CBSA. Besides projecting tokens onto subspaces and back the ambient space, there are generally two stages in CBSA: 1) representative initialization and extraction; 2) representative contraction and contraction broadcast. The former extracts representatives satisfying the inequality constraints in \ref{['eq:CoCa']}, while the latter is a gradient step of the compression term in \ref{['eq:CoCa']}. Right panel: CBSA covers instantiations of varied attention mechanisms. Their compression patterns are distinct as illustrated above. Further analysis is elaborated in Section \ref{['sec:variants']}.
• Figure 2: Computation complexity.
• Figure 3: Different attention mechanisms.$\top$ stands for the matrix transpose. For simplicity, the agent bias in the diagram of Agent Attention Han:ECCV2024-AgentAttention are omitted, and the diagram of our CBSA is also simplified for comparison.
• Figure 4: Compact and structured representation. There are 10 classes indicated by colors. Points in each class are generated by sampling from a one-dimensional subspace and are then perturbed by adding noise.
• Figure 5: Evaluation on compression effect. We measure the compression term of \ref{['eq:MCR2']} as a function of layers.