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

TL;DR

This work addresses interpretability and efficiency gaps in Transformer attention by decoupling the token-subspace membership from the learned subspaces within the MCR^2 objective, enabling independent learning of the membership matrix and sparse subspaces derived from a shared full space. By unrolling gradient steps on the variational MCR^2 objective, the authors derive Decoupled Membership-Subspaces Attention (DMSA), which yields a sparse, linear-attention operator with enhanced interpretability. Replacing the ToST attention with DMSA to form DMST yields consistent Top-1 accuracy gains on ImageNet-1K and substantial memory savings, while preserving linear computational scaling. The work bridges white-box optimization principles with gated attention concepts, and provides extensive analyses including ablations, layer-wise compression trends, and visualizations of learned token-subspace memberships, demonstrating practical impact for high-efficiency, interpretable vision models.

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.

Figures (6)

• Figure 1: Top-1 accuracy on ImageNet-1K (left axis) and peak active memory under 16 K-token inputs (right axis). Dark-blue solid line denotes TSSA memory; dark-purple solid line denotes DMSA memory. Bars indicate validation accuracy across model scales. Our DMST model achieves better performance and consumes less memory.
• Figure 2: One layer ${l}$ of the proposed Decoupled Membership and Subspaces Transformer(DMST). Notably, the self-attention of DMST transforms tokens $\boldsymbol{Z}^l$ efficiently to $\boldsymbol{Z}^{l+1}$, via multiplying each row of the projected token by a scalar. The projected subspace is selected via an independent sparse activate membership $\boldsymbol{\Pi}$.
• Figure 3: DMSA variants. The structure is theoretically equivalent to a form of gated channel attention.
• Figure 4: DMSA achieves lower memory usage than MHSA in ViT and linear attention TSSA.
• Figure 5: The variational compression term ${R}_{\mathrm{cf}}^{\mathrm{var}}(\boldsymbol{Z} \mid (\boldsymbol{\pi}_k, \boldsymbol{U}_k^S))$ of the DMSA outputs $\boldsymbol{Z}$.

Theorems & Definitions (1)

• Theorem 3.1