拓冰网站定制 
【医学图像分割模块】RMT —— 保留注意力 —— 给视觉 Transformer 加一个显式记忆一、论文出处论文全名RMT: Retentive Networks Meet Vision Transformers会议/年份CVPR 2024论文链接https://arxiv.org/abs/2309.11523官方代码https://github.com/qhfan/RMT说明RMT 把 NLP 里 Retentive NetworkRetNet的「显式衰减」思想搬进视觉骨干提出Manhattan Self-AttentionMaSA曼哈顿自注意力——用与曼哈顿距离挂钩的二维空间衰减矩阵给 ViT 补上一个显式的空间先验。本合集只取其中可插拔的RMT Block / MaSA来讲。二、模块图截自论文原文 Figure 3图注RMT 整体四阶段架构上下为 RMT Block 展开——DWConv 3×3 → LN → Manhattan Self-Attention → LN → FFN两条残差。MaSA 用一个按曼哈顿距离衰减的矩阵D调制注意力权重让空间先验「显式」地进入模型。三、核心思想与作用一句话总括MaSA 在标准自注意力Softmax(QKᵀ)V外面乘上一个按「曼哈顿距离」指数衰减的矩阵D让离得近的 token 权重大、离得远的 token 权重指数级变小从而给模型一个显式的二维空间先验——注意力算的还是全局的但先验告诉它「谁更重要」。拆解成几步从 RetNet 的「衰减」出发。RetNet 在语言模型里用D_nm γ^(n-m)n≥m因果、一维给历史 token 降权带来显式的时序先验。RMT 把这个思想迁移到视觉。从一维到二维曼哈顿距离。图像里每个 token 有 (x, y) 坐标RMT 把衰减改成按曼哈顿距离计算D_nm γ^(|xₙ-xₘ| |yₙ-yₘ|)。距离越远衰减越狠且横竖两个方向分别累积——这就是「Manhattan」的由来。注意力公式MaSA(X) (Softmax(QKᵀ) ⊙ D) · V。注意一个反直觉的点RMT 试验后发现去掉 RetNet 的门控gating、保留 Softmax反而在视觉上更好——所以 MaSA 并没有像很多线性注意力那样把 softmax 干掉而是「Softmax 注意力 × 空间衰减」。可分解Decomposed MaSA。前三个阶段用分解版把注意力得分与衰减矩阵沿横、纵两个轴分别拆开算Attn_H、Attn_W再接DWConv的局部上下文增强LCE在保留显式空间先验的同时把早期阶段的开销压下来最后一个阶段用原始 MaSA。整个 Block 就是「DWConv → LN → MaSA → LN → FFN」。为什么好显式空间先验。普通 ViT 的自注意力天生没有任何空间归纳偏置RMT 用衰减矩阵把「距离近的更相关」直接写进注意力比纯数据学出来的位置关系更稳。保留 Softmax 的非线性。不像线性注意力为了省算力丢掉 softmaxMaSA 保留 softmax实验上精度更好。全局感受野 局部增强。曼哈顿衰减负责长程的「软」加权DWConv(LCE) 补局部纹理粗细结合。即插即用。Block 只依赖特征张量可整块替换进现有 ViT / U-Net 的注意力位置。四、在 U-Net 里的插入位置推荐放在瓶颈层bottleneck其次可放在解码器的高层低分辨率层。理由瓶颈层特征图最小、通道最多全局建模需求最强而这里正是把「显式空间先验」发挥出来的地方——医学图像里器官/病灶的空间关系是先验知识用曼哈顿衰减天然贴合。解码器低分辨率层同样适合能在恢复空间细节前先整合全局上下文。不建议放在编码器浅层那里分辨率高、token 多自注意力本身开销大分解版可缓解且浅层更依赖局部纹理。一句话瓶颈层放一个解码器低分辨率层可选放编码器浅层别放。五、复现代码PyTorch逐行中文注释说明以下是简化教学版只保留 MaSA 最核心的「曼哈顿衰减矩阵 空间调制」两个组件去掉了官方实现里的多头分组、分轴向分解、DWConv(LCE) 等工程细节方便理解结构。importtorchimporttorch.nnasnnclassManhattanSelfAttention(nn.Module):简化版曼哈顿自注意力标准 softmax 注意力 × 曼哈顿距离衰减矩阵def__init__(self,dim,num_heads4,gamma0.9):super().__init__()self.dimdim# 输入特征维度self.num_headsnum_heads# 头数self.head_dimdim//num_heads# 每头维度self.gammagamma# 衰减因子越接近 1 空间衰减越慢self.q_projnn.Linear(dim,dim)# Q 投影self.k_projnn.Linear(dim,dim)# K 投影self.v_projnn.Linear(dim,dim)# V 投影self.out_projnn.Linear(dim,dim)# 输出投影self.scaleself.head_dim**-0.5# softmax 缩放defforward(self,x,hwNone):# x: (B, N, C)N H*Whw(H, W) 用于计算 token 的 (x, y) 坐标B,N,Cx.shape Hself.num_heads Dself.head_dimifhwisNone:Hhint(N**0.5);hw(Hh,Hh)h,whwasserth*wN,token 数需等于 H*W# 多头投影: (B, H, N, D)qself.q_proj(x).view(B,N,H,D).transpose(1,2)kself.k_proj(x).view(B,N,H,D).transpose(1,2)vself.v_proj(x).view(B,N,H,D).transpose(1,2)# 标准自注意力得分: (B, H, N, N)attn(q k.transpose(-2,-1))*self.scale attnattn.softmax(dim-1)# 构造曼哈顿距离衰减矩阵 D_nm gamma^(|xn-xm| |yn-ym|)ystorch.arange(h,devicex.device).repeat_interleave(w)# (N,)xstorch.arange(w,devicex.device).repeat(h)# (N,)dist(ys[:,None]-ys[None,:]).abs()(xs[:,None]-xs[None,:]).abs()# (N, N)decay(self.gamma**dist.float()).to(x.dtype)# (N, N)# 用空间衰减调制注意力权重再做归一化attnattn*decay# (B, H, N, N) * (N, N)attnattn/(attn.sum(dim-1,keepdimTrue)1e-6)# 加权聚合: (B, H, N, D)outattn v# 合并多头 输出投影outout.transpose(1,2).reshape(B,N,C)returnself.out_proj(out)注官方实现把上面的dist拆成横、纵两个一维衰减矩阵分别作用Decomposed MaSA并加了一个DWConv的局部上下文增强LCE早期阶段用它降开销。教学版把两步合成一步显式计算重点是理解**「softmax 注意力 × 曼哈顿距离衰减」**这个核心设计。六、插入示例几行塞进你的网络fromtorchimportnnclassBottleneckWithRMT(nn.Module):def__init__(self,channels):super().__init__()self.normnn.LayerNorm(channels)self.masaManhattanSelfAttention(dimchannels,num_heads4,gamma0.9)defforward(self,x):B,C,H,Wx.shape# 展平成 token 序列: (B, H*W, C)tokensx.flatten(2).transpose(1,2)tokenstokensself.masa(self.norm(tokens),hw(H,W))# 残差# 还原回特征图returntokens.transpose(1,2).view(B,C,H,W)# 用法替换掉 U-Net 瓶颈层的普通卷积或自注意力# bottleneck BottleneckWithRMT(512)七、实测经验与注意点γ衰减因子是关键超参越接近 1空间衰减越慢、感受野越大但远距离噪声也越多建议从 0.9 附近试起分割任务可略小0.8~0.9。token 数别太大原始 MaSA 是 O(N²) 的毕竟保留了 softmax瓶颈层展平后 token 数控制在 1k~4k 比较舒服token 再多就上分解版 MaSA。不要盲目去掉 SoftmaxRMT 的实验显示视觉任务上保留 softmax 反而比换成 RetNet 的门控更好——这点和很多「线性注意力」教程的直觉相反别照搬。和卷积互补曼哈顿衰减负责长程加权DWConv(LCE) 负责局部俩一起用比只换注意力更稳。别在浅层硬塞编码器浅层分辨率高、token 多收益低开销大优先放瓶颈和解码器低分辨率层。八、完整工程完整可运行代码与更多即插即用模块已整理在仓库https://github.com/CaiCy6/med-modules下一篇预告拆解另一个可插拔子模块——「通道-空间双路注意力块」看它如何在不增加太多参数的前提下同时建模通道与空间依赖。 【医学图像分割模块】RMT —— 保留注意力 —— 给视觉 Transformer 加一个“显式记忆“
【医学图像分割模块】RMT —— 保留注意力 —— 给视觉 Transformer 加一个显式记忆一、论文出处论文全名RMT: Retentive Networks Meet Vision Transformers会议/年份CVPR 2024论文链接https://arxiv.org/abs/2309.11523官方代码https://github.com/qhfan/RMT说明RMT 把 NLP 里 Retentive NetworkRetNet的「显式衰减」思想搬进视觉骨干提出Manhattan Self-AttentionMaSA曼哈顿自注意力——用与曼哈顿距离挂钩的二维空间衰减矩阵给 ViT 补上一个显式的空间先验。本合集只取其中可插拔的RMT Block / MaSA来讲。二、模块图截自论文原文 Figure 3图注RMT 整体四阶段架构上下为 RMT Block 展开——DWConv 3×3 → LN → Manhattan Self-Attention → LN → FFN两条残差。MaSA 用一个按曼哈顿距离衰减的矩阵D调制注意力权重让空间先验「显式」地进入模型。三、核心思想与作用一句话总括MaSA 在标准自注意力Softmax(QKᵀ)V外面乘上一个按「曼哈顿距离」指数衰减的矩阵D让离得近的 token 权重大、离得远的 token 权重指数级变小从而给模型一个显式的二维空间先验——注意力算的还是全局的但先验告诉它「谁更重要」。拆解成几步从 RetNet 的「衰减」出发。RetNet 在语言模型里用D_nm γ^(n-m)n≥m因果、一维给历史 token 降权带来显式的时序先验。RMT 把这个思想迁移到视觉。从一维到二维曼哈顿距离。图像里每个 token 有 (x, y) 坐标RMT 把衰减改成按曼哈顿距离计算D_nm γ^(|xₙ-xₘ| |yₙ-yₘ|)。距离越远衰减越狠且横竖两个方向分别累积——这就是「Manhattan」的由来。注意力公式MaSA(X) (Softmax(QKᵀ) ⊙ D) · V。注意一个反直觉的点RMT 试验后发现去掉 RetNet 的门控gating、保留 Softmax反而在视觉上更好——所以 MaSA 并没有像很多线性注意力那样把 softmax 干掉而是「Softmax 注意力 × 空间衰减」。可分解Decomposed MaSA。前三个阶段用分解版把注意力得分与衰减矩阵沿横、纵两个轴分别拆开算Attn_H、Attn_W再接DWConv的局部上下文增强LCE在保留显式空间先验的同时把早期阶段的开销压下来最后一个阶段用原始 MaSA。整个 Block 就是「DWConv → LN → MaSA → LN → FFN」。为什么好显式空间先验。普通 ViT 的自注意力天生没有任何空间归纳偏置RMT 用衰减矩阵把「距离近的更相关」直接写进注意力比纯数据学出来的位置关系更稳。保留 Softmax 的非线性。不像线性注意力为了省算力丢掉 softmaxMaSA 保留 softmax实验上精度更好。全局感受野 局部增强。曼哈顿衰减负责长程的「软」加权DWConv(LCE) 补局部纹理粗细结合。即插即用。Block 只依赖特征张量可整块替换进现有 ViT / U-Net 的注意力位置。四、在 U-Net 里的插入位置推荐放在瓶颈层bottleneck其次可放在解码器的高层低分辨率层。理由瓶颈层特征图最小、通道最多全局建模需求最强而这里正是把「显式空间先验」发挥出来的地方——医学图像里器官/病灶的空间关系是先验知识用曼哈顿衰减天然贴合。解码器低分辨率层同样适合能在恢复空间细节前先整合全局上下文。不建议放在编码器浅层那里分辨率高、token 多自注意力本身开销大分解版可缓解且浅层更依赖局部纹理。一句话瓶颈层放一个解码器低分辨率层可选放编码器浅层别放。五、复现代码PyTorch逐行中文注释说明以下是简化教学版只保留 MaSA 最核心的「曼哈顿衰减矩阵 空间调制」两个组件去掉了官方实现里的多头分组、分轴向分解、DWConv(LCE) 等工程细节方便理解结构。importtorchimporttorch.nnasnnclassManhattanSelfAttention(nn.Module):简化版曼哈顿自注意力标准 softmax 注意力 × 曼哈顿距离衰减矩阵def__init__(self,dim,num_heads4,gamma0.9):super().__init__()self.dimdim# 输入特征维度self.num_headsnum_heads# 头数self.head_dimdim//num_heads# 每头维度self.gammagamma# 衰减因子越接近 1 空间衰减越慢self.q_projnn.Linear(dim,dim)# Q 投影self.k_projnn.Linear(dim,dim)# K 投影self.v_projnn.Linear(dim,dim)# V 投影self.out_projnn.Linear(dim,dim)# 输出投影self.scaleself.head_dim**-0.5# softmax 缩放defforward(self,x,hwNone):# x: (B, N, C)N H*Whw(H, W) 用于计算 token 的 (x, y) 坐标B,N,Cx.shape Hself.num_heads Dself.head_dimifhwisNone:Hhint(N**0.5);hw(Hh,Hh)h,whwasserth*wN,token 数需等于 H*W# 多头投影: (B, H, N, D)qself.q_proj(x).view(B,N,H,D).transpose(1,2)kself.k_proj(x).view(B,N,H,D).transpose(1,2)vself.v_proj(x).view(B,N,H,D).transpose(1,2)# 标准自注意力得分: (B, H, N, N)attn(q k.transpose(-2,-1))*self.scale attnattn.softmax(dim-1)# 构造曼哈顿距离衰减矩阵 D_nm gamma^(|xn-xm| |yn-ym|)ystorch.arange(h,devicex.device).repeat_interleave(w)# (N,)xstorch.arange(w,devicex.device).repeat(h)# (N,)dist(ys[:,None]-ys[None,:]).abs()(xs[:,None]-xs[None,:]).abs()# (N, N)decay(self.gamma**dist.float()).to(x.dtype)# (N, N)# 用空间衰减调制注意力权重再做归一化attnattn*decay# (B, H, N, N) * (N, N)attnattn/(attn.sum(dim-1,keepdimTrue)1e-6)# 加权聚合: (B, H, N, D)outattn v# 合并多头 输出投影outout.transpose(1,2).reshape(B,N,C)returnself.out_proj(out)注官方实现把上面的dist拆成横、纵两个一维衰减矩阵分别作用Decomposed MaSA并加了一个DWConv的局部上下文增强LCE早期阶段用它降开销。教学版把两步合成一步显式计算重点是理解**「softmax 注意力 × 曼哈顿距离衰减」**这个核心设计。六、插入示例几行塞进你的网络fromtorchimportnnclassBottleneckWithRMT(nn.Module):def__init__(self,channels):super().__init__()self.normnn.LayerNorm(channels)self.masaManhattanSelfAttention(dimchannels,num_heads4,gamma0.9)defforward(self,x):B,C,H,Wx.shape# 展平成 token 序列: (B, H*W, C)tokensx.flatten(2).transpose(1,2)tokenstokensself.masa(self.norm(tokens),hw(H,W))# 残差# 还原回特征图returntokens.transpose(1,2).view(B,C,H,W)# 用法替换掉 U-Net 瓶颈层的普通卷积或自注意力# bottleneck BottleneckWithRMT(512)七、实测经验与注意点γ衰减因子是关键超参越接近 1空间衰减越慢、感受野越大但远距离噪声也越多建议从 0.9 附近试起分割任务可略小0.8~0.9。token 数别太大原始 MaSA 是 O(N²) 的毕竟保留了 softmax瓶颈层展平后 token 数控制在 1k~4k 比较舒服token 再多就上分解版 MaSA。不要盲目去掉 SoftmaxRMT 的实验显示视觉任务上保留 softmax 反而比换成 RetNet 的门控更好——这点和很多「线性注意力」教程的直觉相反别照搬。和卷积互补曼哈顿衰减负责长程加权DWConv(LCE) 负责局部俩一起用比只换注意力更稳。别在浅层硬塞编码器浅层分辨率高、token 多收益低开销大优先放瓶颈和解码器低分辨率层。八、完整工程完整可运行代码与更多即插即用模块已整理在仓库https://github.com/CaiCy6/med-modules下一篇预告拆解另一个可插拔子模块——「通道-空间双路注意力块」看它如何在不增加太多参数的前提下同时建模通道与空间依赖。