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

2025-11-15

AI Summary

Medical image segmentation is critical for clinical applications, yet existing Transformer-based methods suffer from high computational overhead, hindering deployment. To address this, we propose TM-UNet, a lightweight framework centered on a multi-scale token-memory mechanism. It generates token sequences via spatial scanning and employs parallel pooling coupled with an exponential-gated dynamic matrix memory unit to model long-range dependencies with linear complexity—enabling redundancy-free, hierarchical global context extraction. This design achieves an optimal balance between efficiency and representational capacity. On multiple medical segmentation benchmarks, TM-UNet significantly outperforms leading lightweight models (e.g., TransUNet, MedT), reducing parameter count by 37–52% and accelerating inference by 2.1–3.4×, while matching or exceeding state-of-the-art accuracy—achieving a Dice score of 89.7% on the ACDC dataset.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Segmentation
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Web Mining and Content Analysis: Large pretrained models with web data
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Efficient medical image segmentation with reduced computational cost
• Capturing long-range dependencies through token-memory enhanced modeling
• Achieving hierarchical representation learning without computational overhead

Innovation

• Lightweight token-memory mechanism for efficient segmentation
• Multi-scale token-memory block with linear complexity
• Exponential gating and parallel pooling for hierarchical learning