【医学图像分割模块】DG-GSS / 方向-组图选择性扫描（arXiv 2026）—— 让轻量分割网络“扫“得更聪明

【医学图像分割模块】DG-GSS / 方向-组图选择性扫描arXiv 2026—— 让轻量分割网络扫得更聪明医学图像分割里长距离依赖和局部细节天生打架卷积擅长局部、但感受野有限Transformer 能建模全局、但注意力随分辨率平方级涨放进高分辨率 3D/2D 医学图很吃力。状态空间模型Mamba / Selective Scan给了第三条路——用线性复杂度逼近长程建模于是「VSS / VMamba / U-Mamba」这类工作火了起来。但直接用现成的 2D 选择性扫描SS2D有个隐忧多方向扫描之间、通道分组之间的信息是各扫各的方向与分组之间缺少结构化的交互。这篇工作提出DG-GSSDirection-Group Graph Selective Scan方向-组图选择性扫描把「扫描方向」和「通道分组」的响应当作图的节点在多方向融合之前先做一次图上的结构化信息交换——让不同方向、不同分组之间先商量再合并。一句话扫描不是各扫各的——先让方向与分组在图上互相通气再融合。一、论文出处论文GAD-MambaUNet: Direction-Group Mamba with Gradient-Adaptive DINOv3 Distillation for Lightweight Medical Image Segmentation来源arXiv 20262609.26729论文链接https://arxiv.org/abs/2609.26729核心模块DG-GSSDirection-Group Graph Selective Scan备注该论文完整框架还包含训练期的DINOv3 蒸馏 梯度自适应蒸馏GAD本文只拆解可即插即用的DG-GSS 模块。二、模块图截自论文原文 Figure 2图自 GAD-MambaUNet 原文 Figure 2。左侧 (a) 是整体非对称 U 形网络编码器浅层用MKIR Block局部建模深层 E4–E5 用DG-GSS Block长程建模训练时用冻结 DINOv3 教师 GAD 蒸馏监督右上 (b) 是DG-GSS Block的组成DWConv → ShuffleFN → GroupGAT 与 SS2D → DWConv → ShuffleFN右下 © 展开 DG-GSS 内部通道分组 G1–G4、多方向扫描以及把方向/组当图节点做消息传递的图结构。一句话主干是轻量 U-Net深层塞DG-GSS——用方向-组图组织多方向选择性扫描。三、核心思想与作用一句话把扫描方向和通道分组当图节点先在图上交换信息再做多方向融合。拆解三步通道分组 多方向扫描把特征通道分成若干组G1–G4每组在多个扫描方向上跑选择性扫描SS2D 思路四方向/多方向遍历得到方向的序列响应。方向-组图构建关键创新把「方向节点」和「分组节点」建成一张图边连接same direction同方向、same group同分组与self-connection。于是不同方向、不同分组之间能通过**图消息传递GAT 式注意力聚合**互相交换信息。融合输出把图上聚合后的响应合并回特征图再进后续卷积/归一化。为什么好补齐交互盲区传统多方向选择性扫描里方向之间是简单叠加DG-GSS 让方向/分组结构化成图、显式交互信息交换更充分。仍然线性核心仍是选择性扫描 轻量图聚合复杂度对分辨率近似线性适合高分辨率医学图。即插即用它是一个 block可替换 U-Net 深层/瓶颈的注意力或 SSM 模块也可嵌进任意骨干。在分割里的作用给深层低分辨率、高通道特征做长程建模同时用分组控制算力兼顾看得远和够轻。四、在 U-Net 里的插入位置DG-GSS 是长程建模 block放在需要全局感受野的位置编码器深层E4–E5/ 瓶颈层替换 Self-Attention 或普通卷积块解码器低分辨率阶段在高通道、低分辨率处收益最大序列短、通道多跳连融合处让多尺度特征做一次跨方向/跨分组的结构化交互。分辨率越低、通道越多DG-GSS 相对全注意力的效率优势越明显。五、复现代码PyTorch逐行中文注释下面是简化教学版保留 DG-GSS 的两个核心——多方向选择性扫描方向-组图消息传递用最朴素的实现表达结构便于理解与二次开发。生产使用建议对照官方实现含扫描核实现、ShuffleFN、MKIR 等细节。importtorchimporttorch.nnasnnclassSelectiveScan2D(nn.Module):最简版 2D 选择性扫描对 4 个方向做逐位置的状态递推再合并。 真实实现VSS/VMamba会做并行扫描与门控这里用可读的递推表达思路。 def__init__(self,dim):super().__init__()self.projnn.Linear(dim,dim)# 逐位置输入投影self.dwnn.Conv2d(dim,dim,3,padding1,groupsdim)# 局部前处理def_scan(self,x,order):# order: lr | rl | tb | bt按行/列正反序做状态递推B,C,H,Wx.shapeiforderin(lr,rl):seqx.permute(0,1,2,3).reshape(B,C,H*W)# 逐行else:seqx.permute(0,1,3,2).reshape(B,C,H*W)# 逐列iforderin(rl,bt):seqseq.flip(-1)# 反向htorch.zeros(B,C,devicex.device,dtypex.dtype)outs[]fortinrange(seq.shape[-1]):# 逐位置递推教学用慢h0.9*h0.1*seq[...,t]# 简化的状态更新真实为 A/B/C/Δ 参数化outs.append(h)ytorch.stack(outs,dim-1)iforderin(rl,bt):yy.flip(-1)returnydefforward(self,x):# x: [B, C, H, W]xself.dw(x)dirs[lr,rl,tb,bt]ys[self._scan(x,d)fordindirs]# 四个方向 [B, C, HW]ytorch.stack(ys,dim1)# [B, 4, C, HW]把方向当一维yy.reshape(x.shape[0],4,x.shape[1],x.shape[2],x.shape[3])returny# 返回每方向的响应交给图模块classDirectionGroupGraph(nn.Module):方向-组图消息传递节点 (方向, 分组)边 同方向 / 同分组 / 自连接。 用轻量注意力在图上聚合实现方向与分组之间的结构化信息交换。 def__init__(self,dim,n_groups4,n_dirs4):super().__init__()self.n_dirs,self.n_groupsn_dirs,n_groups self.group_dimdim//n_groups self.node_projnn.Linear(dim,dim)self.gatnn.MultiheadAttention(dim,num_headsn_dirs,batch_firstTrue)defbuild_adj(self,dev):构造邻接矩阵i、j 节点相连若 同方向 或 同分组 或 ij。nself.n_dirs*self.n_groups Atorch.eye(n,devicedev)foriinrange(n):di,gidivmod(i,self.n_groups)# 节点 i 的方向, 分组forjinrange(n):dj,gjdivmod(j,self.n_groups)ifdidjorgigj:# 同方向 或 同分组 → 连边A[i,j]1returnAdefforward(self,per_dir):# per_dir: [B, 4, C, H, W]B,D,C,H,Wper_dir.shape# 每个方向先按通道分组得到 方向×分组 的节点特征xper_dir.reshape(B,D,self.n_groups,self.group_dim,H,W)nodesx.mean(dim(-1,-2)).reshape(B,D*self.n_groups,self.group_dim)# [B, N, gd]nodesself.node_proj(torch.nn.functional.pad(nodes,(0,C-self.group_dim)))Aself.build_adj(per_dir.device)# [N, N]mask(A0)# 用作 attention 掩码不相连则屏蔽agg,_self.gat(nodes,nodes,nodes,attn_maskmask)# 图上注意力聚合returnagg.reshape(B,D,self.n_groups,self.group_dim)# 回到 (方向, 分组)classDGGSSBlock(nn.Module):DG-GSS Block选择性扫描 → 方向-组图交互 → 融合回特征图。def__init__(self,dim,n_groups4):super().__init__()self.normnn.GroupNorm(1,dim)self.scanSelectiveScan2D(dim)self.graphDirectionGroupGraph(dim,n_groupsn_groups)self.fusenn.Conv2d(dim,dim,1)# 方向/分组融合self.outnn.Conv2d(dim,dim,1)# 输出投影defforward(self,x):# x: [B, C, H, W]residualx xself.norm(x)per_dirself.scan(x)# [B, 4, C, H, W]aggself.graph(per_dir)# [B, 4, G, C/G] 图聚合结果# 把方向×分组的聚合结果加权回各方向特征wagg.reshape(agg.shape[0],agg.shape[1],-1)wtorch.softmax(w,dim-1)# 方向/分组权重y(per_dir.flatten(2)*w.mean(-1,keepdimTrue)).reshape_as(per_dir)yy.sum(dim1)# 合并方向 → [B, C, H, W]yself.fuse(y)returnresidualself.out(y)# 残差接住即插即用六、插入示例几行塞进你的网络# 例用 DG-GSS 替换 U-Net 深层编码器块低分辨率、高通道处self.deep_blockDGGSSBlock(dim256,n_groups4)defforward(self,x):# x: [B, C, H, W]xself.deep_block(x)# 替换原来的注意力块 / SSM 块无需改接口returnx七、实测经验与注意点分组数是主旋钮n_groups越大越省算力但每组的通道变少、表达受限n_groups4是论文的常用起点。图边设计同方向 / 同分组 / 自连接是核心增加边会变强也更贵建议先照搬再消融。复杂度核心仍是选择性扫描对分辨率近似线性图节点数只有方向数 × 分组数如 4×416图聚合开销很小。踩坑扫描的方向遍历实现是性能大头教学版逐位置递推只适合理解上生产务必换成并行扫描核输入建议用GroupNorm/BatchNorm预归一化深层高通道处更稳与卷积/注意力混用时注意通道对齐蒸馏DINOv3/GAD是训练技巧推理期移除教师不影响 DG-GSS 推理。医学分割场景适合需要长程依赖又要控显存的 2D/3D 分割小数据集上建议配合预训练或强增广。八、完整工程本文的完整可运行工程DG-GSS 的.py、U-Net 插入 demo、n_groups可调配置、逐行中文注释见开源仓库https://github.com/CaiCy6/med-modules下一篇预告继续拆「即插即用模块」——把高效设计塞进 U-Net越轻越准。