Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yaxuan Jiao, Qing Xu, Yuxiang Luo, Xiangjian He, Zhen Chen, Wenting Duan

TL;DR

TM-UNet tackles the trade-off between segmentation accuracy and computational cost by integrating token-based sequence modeling with a memory-based mechanism in a U-Net architecture. It introduces the Multi-Scale Token-Memory Block (MSTM) built on xLSTM with matrix memory cells to capture long-range dependencies at linear complexity, aided by exponential gating and multi-scale pooling. Across four datasets, TM-UNet delivers state-of-the-art Dice scores with substantially lower FLOPs and higher FPS than competing methods, demonstrating practical viability for clinical deployment. The work advances efficient, globally-aware medical image segmentation by combining lightweight sequence modeling with hierarchical context.

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.

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

TL;DR

TM-UNet tackles the trade-off between segmentation accuracy and computational cost by integrating token-based sequence modeling with a memory-based mechanism in a U-Net architecture. It introduces the Multi-Scale Token-Memory Block (MSTM) built on xLSTM with matrix memory cells to capture long-range dependencies at linear complexity, aided by exponential gating and multi-scale pooling. Across four datasets, TM-UNet delivers state-of-the-art Dice scores with substantially lower FLOPs and higher FPS than competing methods, demonstrating practical viability for clinical deployment. The work advances efficient, globally-aware medical image segmentation by combining lightweight sequence modeling with hierarchical context.

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.

Paper Structure

This paper contains 11 sections, 9 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Methodology
  3. Preliminaries
  4. Extended Long Short-Term Memory Layer
  5. Multi-Scale Token-Memory Block
  6. The Optimization of TM-UNet Framework
  7. Experiments
  8. Datasets and Implementation Details
  9. Comparison with State-of-the-art Methods
  10. Ablation Study
  11. Conclusion

Figures (3)

  • Figure 1: Performance and computation efficiency comparison of state-of-the-art medical image segmentation methods. Results demonstrate the superior performance of TM-UNet with lower computation complexity and faster inference speed.
  • Figure 2: The overview of our TM-UNet framework for efficient medical image segmentation. For ease of understanding, we elaborate on the case of TM-UNet in polyp segmentation. Our TM-UNet fully exploits token sequence modeling to achieve efficient long-range dependency modeling with linear computation complexity.
  • Figure 3: Qualitative comparison on four medical datasets.