Parameter-Efficient Transformer with Hybrid Axial-Attention for Medical Image Segmentation

TL;DR

This paper introduces a parameter-efficient transformer with hybrid axial-attention and a gating mechanism, improving medical image segmentation on small datasets by incorporating spatial and relative position biases.

Contribution

The paper proposes a novel Hybrid Axial-Attention mechanism and a gating strategy to enhance transformer efficiency and effectiveness in medical image segmentation, especially on small datasets.

Findings

Outperforms baseline and previous methods on BraTS and Covid19 datasets.

Incorporating position information improves ROI prediction accuracy.

Gating mechanism reduces training complexity and enhances feature selection.

Abstract

Transformers have achieved remarkable success in medical image analysis owing to their powerful capability to use flexible self-attention mechanism. However, due to lacking intrinsic inductive bias in modeling visual structural information, they generally require a large-scale pre-training schedule, limiting the clinical applications over expensive small-scale medical data. To this end, we propose a parameter-efficient transformer to explore intrinsic inductive bias via position information for medical image segmentation. Specifically, we empirically investigate how different position encoding strategies affect the prediction quality of the region of interest (ROI), and observe that ROIs are sensitive to the position encoding strategies. Motivated by this, we present a novel Hybrid Axial-Attention (HAA), a form of position self-attention that can be equipped with spatial pixel-wise information and relative position information as inductive bias. Moreover, we introduce a gating mechanism to alleviate the burden of training schedule, resulting in efficient feature selection over small-scale datasets. Experiments on the BraTS and Covid19 datasets prove the superiority of our method over the baseline and previous works. Internal workflow visualization with interpretability is conducted to better validate our success.