Graph-Augmented Sparse Attention for Medical Image Segmentation
摘要
Medical image segmentation is essential for identifying lesions and organs, aiding in disease diagnosis. Traditional deep learning models, such as convolutional neural networks and transformers, typically used for this task, can only operate on Euclidean grids. However, high-dimensional medical image features often reside on non-Euclidean manifolds, and current dense self-attention methods treat the image as a complete graph by using all image patches to compute the attention matrix, which can include irrelevant noisy tokens. In this study, we improve TransUNet’s performance by integrating sparse attention mechanisms. Our first model, named \( k \) -NN Attention TransUNet (KATU), simplifies the Vision Transformer component of TransUNet by applying top- \( k \) pruning to produce a sparse attention weight matrix. The second model, Vision-Graph TransUNet (VGTU), replaces vanilla attention in TransUNet with graph attention, representing image patches as nodes connected by spatial relationships using \( k \) -NN. Our models were tested on both Synapse and SynapsePlus datasets, the latter expanded from Synapse to evaluate generalization in segmenting smaller organs. The results showed a significant improvement over various strong baselines, this not only highlights the effectiveness of sparse attention but also suggests that Transformer can be viewed as a specific instance of Graph Attention Network (GAT) applied to complete graphs, indicating that GAT can perform tasks typically handled by Transformer.