arXiv 2025 | 多尺度线性注意的混合CNN-Transformer医学分割网络，即插即用，涨点起飞！
1. 基本信息
-
标题: MSLAU-Net: A Hybird CNN-Transformer Network for Medical Image Segmentation
-
论文来源: https://arxiv.org/pdf/2505.18823
2. 核心创新点
-
提出多尺度线性注意力 (MSLA) 模块: 设计了一种并行结构，融合了 CNN 的多尺度特征提取能力和线性注意力的全局依赖建模能力，同时保持了低计算复杂度。
-
构建混合 CNN-Transformer 编码器: 在编码器的浅层（前两阶段）使用基于 CNN 的局部特征提取（LFE）模块，在深层（后两阶段）使用基于 MSLA 的全局特征提取（GFE）模块，有效结合了两种范式的优势。
-
设计轻量级自顶向下解码器: 采用非对称的解码器结构，通过轻量级卷积和上采样操作实现自顶向下的多级特征聚合，高效地融合语义信息并恢复空间细节。
➔➔➔➔点击查看原文，获取本文及其他精选即插即用模块集合
https://mp.weixin.qq.com/s/4NTp6pvVSR97DKG2nGDMYg
3. 方法详解
整体结构概述： MSLAU-Net 是一个混合了 CNN 和 Transformer 的编码器-解码器架构。其编码器采用四阶段金字塔结构，前两阶段使用 LFE 模块（CNN-based）捕获局部细节，后两阶段使用 GFE 模块（Transformer-based）捕获全局上下文。解码器则采用一种自顶向下的特征聚合机制，将编码器不同层级的特征有效融合，并逐步上采样恢复至原始图像分辨率，以生成最终的分割掩码。
MSLAU-Net 整体架构图
- 步骤分解:
-
编码器 (Encoder):
-
分层结构: 编码器包含四个阶段，通过 Patch Embedding 层（使用步长卷积）逐步降低空间分辨率并增加通道维度。
-
局部特征提取 (LFE) 模块: 在前两个阶段，模型采用 LFE 模块（详见图2）。该模块主要由深度可分离卷积和前馈网络（FFN）构成，专注于高效提取图像的局部纹理和结构化细节。
-
全局特征提取 (GFE) 模块: 在后两个阶段，模型转而使用 GFE 模块（详见图3）。此模块的核心是新提出的 MSLA 模块，用于在更深的层级上建模长距离依赖关系。
-
-
多尺度线性注意力 (MSLA) 模块:
-
多尺度特征提取: 输入特征图在通道维度上被分为四份，分别送入四个并行的分支。每个分支使用不同大小的卷积核（3x3, 5x5, 7x7, 9x9）进行深度可分离卷积，以捕获不同尺度的特征。
-
线性注意力计算: 每个分支提取的多尺度特征随后被送入高效注意力（Efficient Attention）模块中。该模块利用矩阵乘法的结合律将计算复杂度从 O(N²) 降低到 O(N)，高效地建模全局依赖。
-
特征融合: 四个分支的输出通过可学习权重进行加权，并使用一个 1x1 卷积进行融合，最终得到兼具多尺度信息和全局上下文的特征表示。
-
-
解码器 (Decoder):
-
特征对齐: 解码器从编码器的第二、三、四阶段获取特征图，通过级联的 1x1 卷积和双线性上采样操作，将它们的通道数和分辨率与第一阶段的输出对齐。
-
自顶向下聚合: 对齐后的特征图通过逐元素相加的方式，自顶向下（从最深层到最浅层）进行聚合，增强了不同层级特征之间的交互。
-
分辨率恢复: 聚合后的特征图经过拼接和进一步的卷积、上采样操作，最终恢复到输入图像的原始分辨率，用于像素级预测。
-
-
损失函数:
-
论文针对不同数据集采用了混合损失函数，结合了 Dice Loss 和交叉熵（Cross-Entropy）损失，以应对类别不平衡问题。公式如下：
-
-
4. 即插即用模块作用
本文的核心是 MSLA (Multi-Scale Linear Attention) 模块，它可作为一个更高效、更强大的 Transformer 构建块。
适用场景
-
核心任务: 医学图像分割（如器官分割、病灶分割）。
-
具体模态: CT 扫描、MRI 扫描、结肠镜图像等。
-
扩展应用: 可推广至其他需要密集预测的计算机视觉任务，如通用语义分割、目标检测等，尤其是在高分辨率图像处理场景下。
主要作用
-
替代标准自注意力: MSLA 可作为标准 Softmax 自注意力或稀疏注意力（如Swin Transformer中的窗口注意力）的替代品。
-
大幅降低计算复杂度: 将注意力机制的计算复杂度从二次方 O(N²) 降低到线性 O(N)，显著提升了模型处理高分辨率图像时的效率和速度。
-
增强多尺度特征捕获能力: 通过并行卷积分支，显式地建模了不同尺度的局部上下文，弥补了传统注意力机制在捕捉精细局部结构上的不足。
-
保持全局感受野: 与窗口注意力等稀疏机制不同，线性注意力能够建模输入特征图上任意两个像素之间的依赖关系，提供了真正的全局感受野。
总结
MSLA 模块是一个兼具 CNN 多尺度局部感知能力和 Transformer 线性复杂度全局建模能力的“高效混合注意力单元”，旨在用更少的计算资源实现更精准的密集预测。
➔➔➔➔点击查看原文，获取本文及其他精选即插即用模块集合https://mp.weixin.qq.com/s/4NTp6pvVSR97DKG2nGDMYg
5. 即插即用模块
import torch.nn as nn
import torch.nn.functional as F
import torchclass LinearAttention(nn.Module):def __init__(self, dim, num_heads):super().__init__()self.dim = dimself.num_heads = num_headsself.qkv = nn.Linear(dim, 3 * dim, bias=False)self.proj = nn.Linear(dim, dim)def forward(self, x):b, c, h, w = x.shapex = x.view(b, c, h * w).permute(0, 2, 1) # (b, h*w, c)qkv = self.qkv(x).reshape(b, h * w, 3, self.num_heads, self.dim // self.num_heads).permute(2, 0, 3, 1, 4) # 微信公众号:AI缝合术q, k, v = qkv[0], qkv[1], qkv[2]key = F.softmax(k, dim=-1)query = F.softmax(q, dim=-2)context = key.transpose(-2, -1) @ vx = (query @ context).reshape(b, h * w, c)x = self.proj(x)x = x.permute(0, 2, 1).view(b, c, h, w)return xclass DepthwiseConv(nn.Module):def __init__(self, in_channels, kernel_size):super(DepthwiseConv, self).__init__()self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, groups=in_channels, padding=kernel_size // 2)self.relu = nn.ReLU()def forward(self, x):residual = xx = self.depthwise(x)x = x + residualx = self.relu(x)return xclass MSLA(nn.Module):def __init__(self, dim, num_heads):super().__init__()self.dim = dimself.num_heads = num_headsself.dw_conv_3x3 = DepthwiseConv(dim // 4, kernel_size=3)self.dw_conv_5x5 = DepthwiseConv(dim // 4, kernel_size=5)self.dw_conv_7x7 = DepthwiseConv(dim // 4, kernel_size=7)self.dw_conv_9x9 = DepthwiseConv(dim // 4, kernel_size=9)self.linear_attention = LinearAttention(dim = dim // 4, num_heads = num_heads) self.final_conv = nn.Conv2d(dim, dim, 1)self.scale_weights = nn.Parameter(torch.ones(4), requires_grad=True)def forward(self, input_):b, c, h, w = input_.shapeinput_reshaped = input_.view(b, c, h, w)split_size = c // 4x_3x3 = input_reshaped[:, :split_size, :, :]x_5x5 = input_reshaped[:, split_size:2 * split_size, :, :]x_7x7 = input_reshaped[:, 2 * split_size:3 * split_size:, :, :]x_9x9 = input_reshaped[:, 3 * split_size:, :, :]x_3x3 = self.dw_conv_3x3(x_3x3)x_5x5 = self.dw_conv_5x5(x_5x5)x_7x7 = self.dw_conv_7x7(x_7x7)x_9x9 = self.dw_conv_9x9(x_9x9)att_3x3 = self.linear_attention(x_3x3)att_5x5 = self.linear_attention(x_5x5)att_7x7 = self.linear_attention(x_7x7)att_9x9 = self.linear_attention(x_9x9)processed_input = torch.cat([att_3x3 * self.scale_weights[0],att_5x5 * self.scale_weights[1],att_7x7 * self.scale_weights[2],att_9x9 * self.scale_weights[3]], dim=1)final_output = self.final_conv(processed_input)return final_outputif __name__ == "__main__":batch_size = 1height, width = 128, 128 # 输入图像大小channels = 32 # 输入通道数num_heads = 8 # 注意力头数# 创建输入张量：形状为 (B, C, H, W)，四维张量x = torch.randn(batch_size, channels, height, width)# 初始化 MSLA 模块msla = MSLA(dim=channels, num_heads=num_heads)# 前向传播测试output = msla(x)# 输出结果print(msla)print("输入张量形状:", x.shape) # [B, C, H, W]print("输出张量形状:", output.shape) # [B, n, dim]