Token Sparsification for Faster Medical Image Segmentation笔记

Token稀疏使得更快得医学影像分割

摘要

• 利用稀疏标记进行密集预测
• 分割重新定义：稀疏编码→Token补全→密集解码(SCD)
• 提出soft-topk token修剪(STP)和多层token组装(MTA)
• STP使用轻量级子网络预测token重要性分数，并对topK token进行采样
• MTA通过组装稀疏输出token和修剪的多层中间token来恢复完整的token序列
• 密集解码阶段与现有的分割解码器兼容
• 在保持分割质量的同时，配备STP和MTA的SCD在训练(高达120%的吞吐量)和推断(高达60.6%的吞吐量)方面都比没有token修剪的基线快得多

引言

在本文之前没有Vit token 稀疏方法来分割
作者设计思路：
稀疏编码→Token补全→密集解码（SCD）
采用现有办法，EVIT和Dynamic Vit去稀疏编码，MIM用来token 补全
存在的问题：

1. token分数估计：EVIT的分数存在CLS里面，不便于分割解码器
2. token采样：DynamicViT将token采样框架为一系列独立的二进制决定，以保留或删除令牌。这并不能保证每个训练输入都有固定数量的采样令牌。为了适应批量训练，DynamicViT将所有标记保存在内存中，并mask自我注意条目，导致训练效率低下。
3. token补全：MIM对预训练有用，但它不能准确地恢复详细信息，导致分割结果较差

本文提出的方法：

1. Soft-topK Token Pruning（STP）：预测token重要性通过子网络，采用topK的分数的token，在内存和计算中只保留采样的标记来加速训练
2. MTA通过组装来自多层的稀疏输出token和修剪的中间令牌来恢复完整的令牌序列

方法

分割模型由编码器和解码器组成
目标是加速ViT分割编码器。
稀疏编码学习了一个用于加速的稀疏标记表示
令牌补全恢复密集预测的完整token
密集解码从密集特征中预测分割掩码
如图所示：

稀疏编码STP

TF是transformer模块
STP选择效果最佳的模块：由分数预测和选择组成
token的分数估计：

topk采样：
将离散topK掩模的采样放宽到一个连续的近似，即Gumbel-Softmax分布（方便求导），左边是前向，右边是反向:

token补全

输入为：稀疏编码加上被丢掉的token
网络为多层感知机
位置编码：

解码器

UNETR，需要来自编码器的多层输出的输入，这导致了问题，因为中间特征仍然稀疏
受最近研究的非层次特征金字塔的激励，使用补全网络的输出zcompl来替换分割头所需的所有中间特征

实验结果：