DCMM-Transformer: Degree-Corrected Mixed-Membership Attention for Medical Imaging

TL;DR

DCMM-Transformer addresses the challenge that standard Vision Transformers overlook latent anatomical groupings in medical images by injecting a differentiable, DCMM-based additive bias into self-attention. The method learns soft community memberships and degree centrality from queries to construct a connection-probability matrix that biases attention, preserving gradient flow while guiding the model toward anatomically coherent patch interactions. Empirically, it delivers state-of-the-art performance across five medical imaging datasets and yields attention maps that align with clinically meaningful regions, improving interpretability. The approach offers a practical, scalable way to leverage structural priors in medical imaging, with broad implications for diagnosis and prognosis across modalities.

Abstract

Medical images exhibit latent anatomical groupings, such as organs, tissues, and pathological regions, that standard Vision Transformers (ViTs) fail to exploit. While recent work like SBM-Transformer attempts to incorporate such structures through stochastic binary masking, they suffer from non-differentiability, training instability, and the inability to model complex community structure. We present DCMM-Transformer, a novel ViT architecture for medical image analysis that incorporates a Degree-Corrected Mixed-Membership (DCMM) model as an additive bias in self-attention. Unlike prior approaches that rely on multiplicative masking and binary sampling, our method introduces community structure and degree heterogeneity in a fully differentiable and interpretable manner. Comprehensive experiments across diverse medical imaging datasets, including brain, chest, breast, and ocular modalities, demonstrate the superior performance and generalizability of the proposed approach. Furthermore, the learned group structure and structured attention modulation substantially enhance interpretability by yielding attention maps that are anatomically meaningful and semantically coherent.

Figures (4)

• Figure 1: Overview of the $\hbox{DCMM-Transformer}$. The DCMM model imposes latent community structure among image patches, grouping them into blocks (e.g., Blocks 1–3 in blue, red, and green), and assigns probabilities to the interactions within and between these communities. This community-based bias is then added to the standard attention logits.
• Figure 2: Overview of the proposed DCMM-Transformer mechanism integrated into the self-attention module of Vision Transformers. For each token, a soft group membership vector and a degree scalar are constructed from the query representation. These are used to construct a structured attention bias based on learnable inter-group affinities. The bias is added to the standard attention scores to modulate them before softmax.
• Figure 3: Visualization of attention maps from a standard Transformer and the proposed DCMM-Transformer (abbreviated as DCMM-TF in the figure) on the five datasets. The left column shows the original input images, the middle column displays attention maps from the standard Transformer, and the right column presents attention maps from $\hbox{DCMM-Transformer}$. For each category, two random subjects are displayed.
• Figure 4: Influence of the number of communities on classification performance across five medical image datasets. The mean classification accuracy and standard deviation over three independent runs for each setting are reported.