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

TL;DR

DCMM-Transformer introduces a fully differentiable, interpretable attention mechanism that models anatomical groupings in medical images, outperforming prior methods and enhancing interpretability across diverse datasets.

Contribution

It proposes a novel ViT architecture incorporating a degree-corrected mixed-membership model as an additive bias in self-attention, addressing limitations of previous stochastic masking approaches.

Findings

Superior performance across multiple medical imaging modalities.

Enhanced interpretability with anatomically meaningful attention maps.

Demonstrated generalizability and robustness of the approach.

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.