Interpretable and Sparse Linear Attention with Decoupled Membership-Subspace Modeling via MCR2 Objective

TL;DR

This paper introduces Decoupled Membership-Subspace Attention (DMSA), a sparse linear attention mechanism derived from MCR2 that enhances interpretability and efficiency in visual transformer models, outperforming existing methods on ImageNet-1K.

Contribution

It decouples membership and subspace matrices in MCR2, deriving an interpretable sparse attention operator, and demonstrates improved accuracy and efficiency in visual transformers.

Findings

DMSA achieves faster coding reduction rate.

DMST outperforms ToST by 1.08%-1.45% in top-1 accuracy.

DMST shows higher computational efficiency and interpretability.

Abstract

Maximal Coding Rate Reduction (MCR2)-driven white-box transformer, grounded in structured representation learning, unifies interpretability and efficiency, providing a reliable white-box solution for visual modeling. However, in existing designs, tight coupling between "membership matrix" and "subspace matrix U" in MCR2 causes redundant coding under incorrect token projection. To this end, we decouple the functional relationship between the "membership matrix" and "subspaces U" in the MCR2 objective and derive an interpretable sparse linear attention operator from unrolled gradient descent of the optimized objective. Specifically, we propose to directly learn the membership matrix from inputs and subsequently derive sparse subspaces from the fullspace S. Consequently, gradient unrolling of the optimized MCR2 objective yields an interpretable sparse linear attention operator: Decoupled Membership-Subspace Attention (DMSA). Experimental results on visual tasks show that simply replacing the attention module in Token Statistics Transformer (ToST) with DMSA (we refer to as DMST) not only achieves a faster coding reduction rate but also outperforms ToST by 1.08%-1.45% in top-1 accuracy on the ImageNet-1K dataset. Compared with vanilla Transformer architectures, DMST exhibits significantly higher computational efficiency and interpretability.