用记忆机制提升医学图像分割效率,速度更快精度不降。
TM-UNet: Token-Memory Enhanced Sequential Modeling for Efficient Medical Image Segmentation
- 将图像转为令牌序列,用记忆单元动态存取关键上下文信息。
- 在多个数据集上超越现有方法,计算量减少超30%。
- 适合医疗影像实时分析,尤其对算力受限场景友好。
医学图像分割对临床诊断与治疗规划至关重要。尽管基于Transformer的方法表现优异,但其高计算成本限制了临床应用。为此,我们提出TM-UNet,一种融合令牌序列建模与高效记忆机制的轻量级框架。具体而言,引入多尺度令牌记忆(MSTM)模块,通过策略性空间扫描将2D空间特征转化为令牌序列,利用矩阵记忆单元选择性保留并传播判别性上下文信息。该机制作为动态知识库,以线性复杂度捕捉长距离依赖,实现高效全局推理且无冗余计算。MSTM模块进一步结合指数门控识别令牌有效性,并通过并行池化操作实现多尺度上下文提取,支持层次化表征学习而无额外开销。大量实验表明,TM-UNet在多种医学分割任务中优于当前最优方法,同时显著降低计算成本。代码已公开于https://github.com/xq141839/TM-UNet。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。