Less is More: AMBER-AFNO -- a New Benchmark for Lightweight 3D Medical Image Segmentation

TL;DR

This paper introduces AMBER-AFNO, a lightweight 3D medical image segmentation model that replaces self-attention with spectral operations for improved efficiency and competitive accuracy.

Contribution

The paper proposes a novel AFNO-based architecture for 3D segmentation, reducing computational complexity and maintaining high performance compared to transformer-based models.

Findings

Achieves state-of-the-art or near-state-of-the-art results on three datasets.

Offers quasi-linear computational complexity and linear memory scaling.

Outperforms recent compact CNN and Transformer architectures in Dice scores.

Abstract

We adapt the remote sensing-inspired AMBER model from multi-band image segmentation to 3D medical datacube segmentation. To address the computational bottleneck of the volumetric transformer, we propose the AMBER-AFNO architecture. This approach uses Adaptive Fourier Neural Operators (AFNO) instead of the multi-head self-attention mechanism. Unlike spatial pairwise interactions between tokens, global token mixing in the frequency domain avoids $\mathcal{O}(N^2)$ attention-weight calculations. As a result, AMBER-AFNO achieves quasi-linear computational complexity and linear memory scaling. This new way to model global context reduces reliance on dense transformers while preserving global contextual modeling capability. By using attention-free spectral operations, our design offers a compact parameterization and maintains a competitive computational complexity. We evaluate AMBER-AFNO on three public datasets: ACDC, Synapse, and BraTS. On these datasets, the model achieves state-of-the-art or near-state-of-the-art results for DSC and HD95. Compared with recent compact CNN and Transformer architectures, our approach yields higher Dice scores while maintaining a compact model size. Overall, our results show that frequency-domain token mixing with AFNO provides a fast and efficient alternative to self-attention mechanisms for 3D medical image segmentation.