Towards Interpretable and Efficient Attention: Compressing All by Contracting a Few
Proposes Contract-and-Broadcast Self-Attention (CBSA), unifying interpretable and linear-complexity attention via few representatives.
Key Findings
Methodology
This work employs algorithm unrolling to convert the gradient step of a compression-based objective into a forward pass, creating CBSA. It uses a small set of representatives to compress input tokens into low-dimensional structures, enabling linear complexity. The approach unifies various attention types—softmax, linear, channel—within a single interpretable framework, grounded in information theory and optimization. Extensive experiments on vision datasets demonstrate that CBSA maintains competitive accuracy while significantly reducing computational costs, validating its theoretical and practical advantages.
Key Results
- CBSA achieves 79.2% top-1 accuracy on ImageNet-1K, surpassing softmax attention (77.8%) with fewer parameters and lower FLOPs. In semantic segmentation on COCO, it improves mIoU to 45.3%, outperforming baseline models. The method's parameter count decreases by 15%, with inference speed up by 20%. Ablation studies confirm fixed representative numbers still yield high performance, illustrating effective low-dimensional compression. Variants of CBSA cover multiple attention types, confirming its unifying potential.
Significance
This research bridges the gap between interpretability and efficiency in attention mechanisms, offering a unified theoretical framework that enhances transparency and scalability. It addresses longstanding challenges in understanding attention's inner workings and deploying large-scale models efficiently, paving the way for more trustworthy AI systems in vision and beyond. The approach's versatility and solid theoretical grounding mark a significant step forward in attention research, with broad implications for model design and interpretability.
Technical Contribution
The paper introduces a novel optimization objective based on information-theoretic coding rates, combined with algorithm unrolling, to derive CBSA. This mechanism employs few representatives to compress input tokens, achieving linear complexity and covering multiple attention variants. It provides a unified formula that explains the core differences among attention types through their distinct compression patterns. Theoretical guarantees and extensive empirical validation demonstrate its effectiveness and interpretability, setting a new direction for attention mechanism design.
Novelty
This is the first work to unify softmax, linear, and channel attention within a single, interpretable, low-complexity framework based on information theory and algorithm unrolling. Unlike prior methods that treat attention as a black box or rely on heuristic simplifications, CBSA offers a principled, unified understanding of attention mechanisms, revealing their fundamental connections and differences. Its ability to adaptively cover multiple variants through representative structures marks a significant innovation in the field.
Limitations
- The current validation is primarily on visual tasks; applicability to NLP and other modalities remains to be explored. Future work should extend the framework to sequence modeling beyond vision.
- Representatives are manually set or learned with fixed structures, which may limit flexibility; adaptive or dynamic schemes are needed.
- Handling extremely long sequences or ultra-high-resolution images may still pose challenges, requiring integration with sparse or multi-scale techniques. Additionally, training complexity and hyperparameter tuning could be further optimized.
Future Work
Future directions include extending CBSA to multi-modal and language tasks, integrating sparse and multi-scale methods for longer sequences, and developing automatic mechanisms for learning optimal representative structures. Further theoretical analysis on expressive capacity and robustness, as well as hardware-efficient implementations, are also promising avenues to enhance practical deployment.
AI Executive Summary
Attention mechanisms have revolutionized deep learning, especially in vision and NLP, but their quadratic complexity and opaque nature limit scalability and trustworthiness. Traditional softmax attention, while effective, becomes computationally prohibitive as data scales. Linear attention variants like Performer and LinFormer reduce complexity but often sacrifice interpretability and flexibility. This paper introduces Contract-and-Broadcast Self-Attention (CBSA), a novel mechanism grounded in information theory and algorithm unrolling, that compresses input tokens into a few representative vectors. These representatives serve as low-dimensional summaries, enabling the attention operation to scale linearly with input size. The core idea is to formulate a unified optimization objective based on coding rates, which encourages the input tokens to contract towards a small set of meaningful representatives. By unrolling the optimization process, the authors derive CBSA, which performs representative contraction and broadcasts the compressed information back to tokens, effectively mimicking various attention types—softmax, linear, and channel attention—within a single interpretable framework. Extensive experiments on image classification and segmentation tasks demonstrate that CBSA maintains competitive accuracy while significantly reducing computational costs, with a 20% speed-up and 15% parameter reduction compared to baseline models. The mechanism’s interpretability is validated through visualization of the low-dimensional structures learned at each layer, revealing how the model progressively compresses and organizes information. The unifying formula not only clarifies the relationships among different attention variants but also opens avenues for designing more transparent and efficient models. Future work will explore extending CBSA to NLP, multi-modal tasks, and dynamic representative learning, aiming to build scalable, interpretable AI systems for real-world applications.
Deep Analysis
Background
Attention mechanisms在深度学习中扮演着核心角色，尤其在Transformer架构中实现了信息的动态加权。早期如Vaswani等提出的Transformer，依赖softmax注意力，虽提升性能，但计算复杂度为二次级，限制了大规模应用。为解决这一瓶颈，线性注意力机制如Performer、LinFormer等被提出，显著降低了复杂度，但多为黑盒，缺乏可解释性。近年来，学界开始关注机制的内在理解，利用聚类、能量最小化等方法实现可解释性，但多依赖后置解释，缺乏统一的理论框架。与此同时，长序列和高分辨率图像的处理仍面临巨大挑战。本文基于信息论和算法展开思想，试图在效率与可解释性之间找到平衡，为注意力机制提供新的理论基础。
Core Problem
现有注意力机制在性能、效率和可解释性方面存在矛盾。softmax机制虽效果优异，但复杂度高，难以扩展。线性机制虽降低复杂度，但在表达能力和可解释性上不足。黑盒模型难以理解其决策过程，限制了可信度。如何设计一种既高效又具可解释性的注意力机制，成为核心难题。特别是在视觉任务中，长序列和高分辨率图像带来的计算瓶颈亟待突破。缺乏统一的理论框架，使得不同变体的本质关系难以理解。
Innovation
本文提出CBSA机制，通过代表收缩与广播实现注意力的统一表达。创新点包括：1）基于信息论的编码率思想，将输入逐层压缩到低维空间；2）利用算法展开，将目标函数中的梯度步骤转化为前向操作，确保机制的可解释性；3）设计多代表结构，兼容softmax、线性和通道注意力，提供统一公式。该机制实现线性复杂度，能在多任务中保持优异性能，突破黑盒限制，为理解和优化注意力机制提供新思路。
Methodology
- �� 定义信息论中的编码率，作为输入压缩的目标。• 通过代表提取和收缩，利用梯度展开实现机制设计。• 代表初始化采用平均池化，收缩通过梯度步骤实现。• 代表结构支持多变体，支持不同注意力类型。• 将目标函数转化为前向操作，确保机制的可解释性。• 设计代表收缩与广播流程，保证线性复杂度。• 多代表、多变体设置，验证机制的统一性与泛化能力。
Experiments
在ImageNet-1K和COCO数据集上，比较CBSA与softmax、MSSA、TSSA等方法的性能。采用标准训练策略，调节代表数和层数，进行消融分析。评估指标包括准确率、参数量、推理速度。验证模型逐层低维收缩的效果。结果显示，CBSA在保持性能的同时，显著降低复杂度，参数减少15%，推理速度提升20%。多代表设置下，机制表现出良好的泛化能力和稳定性。
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.