Token Sparsification for Faster Medical Image Segmentation
Token Sparsification for Faster Medical Image Segmentation
复制标题
DOI：
10.48550/arxiv.2303.06522
复制
发表时间：
2023-03
期刊：
影响因子：
--
通讯作者：
Lei Zhou;Huidong Liu;Joseph Bae;Junjun He;D. Samaras;P. Prasanna
中科院分区：
文献类型：
--
作者：
Lei Zhou;Huidong Liu;Joseph Bae;Junjun He;D. Samaras;P. Prasanna
我们可以使用稀疏标记进行密集预测，例如分割吗？尽管令牌稀疏化已应用于视觉变换器（ViT）以加速分类，但仍然未知如何从稀疏令牌执行分割。为此，我们将分段重新表述为稀疏编码→令牌完成→密集解码（SCD）管道。我们首先凭经验表明，天真地应用分类标记修剪和掩模图像建模（MIM）的现有方法会导致由于不适当的采样算法和恢复的密集特征的低质量而导致的失败和低效训练。在本文中，我们提出 Soft-topK 令牌修剪（STP）和多层令牌组装（MTA）来解决这些问题。在稀疏编码中，STP 使用轻量级子网络预测 token 重要性分数并对 topK token 进行采样。棘手的 topK 梯度通过连续的扰动分数分布来近似。在令牌完成中，MTA 通过组装稀疏输出令牌和修剪后的多层中间令牌来恢复完整的令牌序列。最后的密集解码阶段与现有的分段解码器兼容，例如 UNETR。实验表明，配备 STP 和 MTA 的 SCD 管道在训练（吞吐量提高高达 120%）和推理（吞吐量提高高达 60.6%）方面比没有令牌修剪的基线快得多，同时保持分段质量。代码可在此处获取：https://github.com/cvlab-stonybrook/TokenSparse-for-MedSeg。
Can we use sparse tokens for dense prediction, e.g., segmentation? Although token sparsification has been applied to Vision Transformers (ViT) to accelerate classification, it is still unknown how to perform segmentation from sparse tokens. To this end, we reformulate segmentation as a sparse encoding → token completion → dense decoding (SCD) pipeline. We first empirically show that naïvely applying existing approaches from classification token pruning and masked image modeling (MIM) leads to failure and inefficient training caused by inappropriate sampling algorithms and the low quality of the restored dense features. In this paper, we propose Soft-topK Token Pruning (STP) and Multi-layer Token Assembly (MTA) to address these problems. In sparse encoding, STP predicts token importance scores with a lightweight sub-network and samples the topK tokens. The intractable topK gradients are approximated through a continuous perturbed score distribution. In token completion, MTA restores a full token sequence by assembling both sparse output tokens and pruned multi-layer intermediate ones. The last dense decoding stage is compatible with existing segmentation decoders, e.g., UNETR. Experiments show SCD pipelines equipped with STP and MTA are much faster than baselines without token pruning in both training (up to 120% higher throughput) and inference (up to 60.6% higher throughput) while maintaining segmentation quality. Code is available here: https://github.com/cvlab-stonybrook/TokenSparse-for-MedSeg.