TM-UNet: Token-Memory Enhanced Sequential Modeling for Efficient Medical Image Segmentation

TL;DR

TM-UNet introduces a lightweight, memory-augmented transformer framework for efficient medical image segmentation, achieving high accuracy with reduced computational cost.

Contribution

It proposes a novel multi-scale token-memory block that captures long-range dependencies efficiently using a dynamic memory mechanism and hierarchical feature extraction.

Findings

Outperforms state-of-the-art methods on multiple medical segmentation datasets.

Reduces computational cost significantly compared to existing transformer-based models.

Demonstrates effective long-range dependency modeling with linear complexity.

Abstract

Medical image segmentation is essential for clinical diagnosis and treatment planning. Although transformer-based methods have achieved remarkable results, their high computational cost hinders clinical deployment. To address this issue, we propose TM-UNet, a novel lightweight framework that integrates token sequence modeling with an efficient memory mechanism for efficient medical segmentation. Specifically, we introduce a multi-scale token-memory (MSTM) block that transforms 2D spatial features into token sequences through strategic spatial scanning, leveraging matrix memory cells to selectively retain and propagate discriminative contextual information across tokens. This novel token-memory mechanism acts as a dynamic knowledge store that captures long-range dependencies with linear complexity, enabling efficient global reasoning without redundant computation. Our MSTM block further incorporates exponential gating to identify token effectiveness and multi-scale contextual extraction via parallel pooling operations, enabling hierarchical representation learning without computational overhead. Extensive experiments demonstrate that TM-UNet outperforms state-of-the-art methods across diverse medical segmentation tasks with substantially reduced computation cost. The code is available at https://github.com/xq141839/TM-UNet.