UTANet：用任务自适应专家混合重构 U-Net 跳跃连接

Date 2026-07-25 · Category tech · Status finished · Confidence likely

论文解读, 医学影像, 图像分割

论文信息

论文：Rethinking U-Net: Task-Adaptive Mixture of Skip Connections for Enhanced Medical Image Segmentation

作者：Zichen Luo、Xinshan Zhu、Lan Zhang、Biao Sun

会议：AAAI Conference on Artificial Intelligence (AAAI-25)，2025

这篇论文讨论医学图像分割中一个基础却经常被默认接受的设计：U-Net 的跳跃连接通常遵循固定的一对一规则，即第 $i$ 层编码器特征直接交给对应的第 $i$ 个解码阶段。UTANet 提出的不是新的编码器或解码器主干，而是对这条规则本身进行改写：每个解码阶段应从多尺度特征中选择何种补充信息，应由任务和输入特征共同决定。

问题：固定跳跃连接的隐含假设

在标准 U-Net 中，编码器逐层下采样，浅层特征保留边缘、纹理和精确位置，深层特征保留类别和结构语义。跳跃连接将前者传回解码器，以补偿下采样造成的空间细节损失。

传统写法可概括为：

S_{i} = E_{i}

其中， $E_{i}$ 是第 $i$ 层编码器特征， $S_{i}$ 是送入对应解码阶段的跳跃连接特征。这种写法隐含两个假设：不同解码阶段只需要同尺度编码特征；该匹配关系可在所有数据集上保持不变。

作者通过改变四条跳跃连接的排列顺序进行实验，发现最优排列会随数据集变化，且往往偏离原始 U-Net 的连接顺序。该现象说明，解码器对补充特征存在阶段性和任务相关的偏好。细胞核、腺体、皮肤病灶与腹部多器官分割对局部边界、组织纹理和全局语义的依赖并不相同，固定连接可能同时带来信息不足和信息干扰。

不同跳跃连接范式与本文提出的稀疏专家连接。图源：Luo et al., AAAI 2025。

不同跳跃连接排列在四个数据集上的 Dice 结果。图源：Luo et al., AAAI 2025。

因此，本文将跳跃连接重述为一个信息分配问题：不再预先规定 $E_{i} \to D_{i}$ ，而是让每个解码阶段从多尺度信息的候选集合中选择合适的专家组合。

UTANet：在编码器与解码器之间插入 TA-MoSC

UTANet 保持 U 型主干，只将原始跳跃连接替换为 Task-Adaptive Mixture of Skip Connections，简称 TA-MoSC。该模块由三部分组成：多尺度特征聚合带、router bank 与共享 expert bank；每个输出还经由 Docker 模块适配到目标解码器的尺寸和通道数。

UTANet 总体结构。TA-MoSC 位于编码器和解码器之间，取代原始的直接跳跃连接。图源：Luo et al., AAAI 2025。

整体信息流为：

E1, E2, E3, E4
  -> 尺寸对齐与通道融合
  -> 公共多尺度特征 E
  -> 共享专家 EP1...EPN
  -> 各 decoder 的专属 router 选择 TopK 专家
  -> Docker 对齐
  -> Ô1...Ô4 送入解码器

最深层编码器特征 $E_{5}$ 仍用于初始化最深解码阶段；需要重新分配的是由 $E_{1}$ 到 $E_{4}$ 提供的补充信息。

多尺度特征聚合：先建立公共候选池

设编码器输出为 ${E_{1}, E_{2}, E_{3}, E_{4}, E_{5}}$ 。TA-MoSC 先将前四层特征通过双线性插值调整到统一空间尺寸：

\hat{E}_{i} = f_{bi} (E_{i}), i = 1, 2, 3, 4

随后按通道维拼接，并用 $1 \times 1$ 卷积完成通道混合与压缩：

E = F_{f} (Cat (\hat{E}_{1}, \hat{E}_{2}, \hat{E}_{3}, \hat{E}_{4}))

F_{f} (\cdot) = ReLU (BN (Conv_{1 \times 1} (\cdot)))

这里的 $E$ 不代表某个单独尺度，而是同时含有浅层细节和中深层语义的公共信息池。后续每个解码阶段都从这一个池中获得候选信息，因此系统可以建立跨尺度连接，但不会为每一种显式连接排列单独训练一套网络。

Router 与 Expert：为每个解码阶段生成不同的跳跃连接

阶段专属 router

对第 $i$ 个解码阶段，router 对聚合特征 $E$ 做全局最大池化，并使用其专属参数 $W_{i}$ 生成专家权重：

g_{i} = Softmax (Pool_{ma x} (E) W_{i})

若共有 $N$ 个专家，则：

g_{i} = [g_{i, 1}, g_{i, 2}, \dots, g_{i, N}], n = 1 \sum N g_{i, n} = 1

论文图示中使用四个专家，因此存在四个 router 参数组 $W_{1}, W_{2}, W_{3}, W_{4}$ 。它们分别控制四条送往解码器的新的跳跃连接。权重并非人工设定，而是由当前输入的特征 $E$ 计算，并通过最终分割损失反向传播学习。

共享专家库

第 $n$ 个专家接收同一个 $E$ ，但有独立的可学习参数。专家本身是轻量的两层 $1 \times 1$ 卷积子网络：

E P_{n} (E) = [ReLU \circ BN \circ Conv_{1 \times 1}]^{2} (E)

记专家输出为：

q_{n} = E P_{n} (E), n = 1, \dots, N

共享指的是四个解码阶段共用同一组 ${E P_{1}, \dots, E P_{N}}$ ，而不是所有专家共享同一组参数。该设计避免为每条跳跃连接配置一套独立专家库，同时允许专家在训练中形成不同的特征变换偏好。

TA-MoSC 的路由过程。每个 router 为其对应解码阶段生成一组专家权重，并从共享专家库中进行稀疏选择。图源：Luo et al., AAAI 2025。

TopK 稀疏混合

router 不直接使用全部专家，而是保留权重最大的 $K$ 个。令：

I_{i} = TopK (g_{i}, K)

则第 $i$ 条跳跃连接的中间结果可展开为：

O_{i} = BN (n \in I_{i} \sum g_{i, n} \cdot E P_{n} (E))

这表明新跳跃连接不再等于某一层编码器输出，而是由多尺度聚合特征经过若干专家变换后、按当前解码阶段偏好加权得到：

\hat{O}_{i} \neq = T (E_{i}), \hat{O}_{i} = T_{i} (E_{1}, E_{2}, E_{3}, E_{4})

最后，Docker 使用插值和 $1 \times 1$ 卷积完成空间尺寸、通道数适配：

\hat{O}_{i} = ReLU (Conv_{1 \times 1} (f_{bi} (O_{i})))

论文给出的解码递推形式为：

D_{4} = D E_{4} (E_{5})

D_{i} = D E_{i} (Cat (\hat{O}_{i + 1}, D_{i + 1})), i = 0, 1, 2, 3

Pred = Seg (D_{0})

BEU：防止稀疏路由使专家闲置

TopK 路由带来效率，但也可能使少数专家长期占据高权重，其他专家因缺少梯度而无法形成有效分工。论文为此提出 Balanced Expert Utilization，简称 BEU。

第一部分是 Expert Variance Loss。它惩罚不同专家使用量之间过大的方差，使 router 不会持续将概率集中在少数专家上。总损失为：

L = L_{D B} + β L_{E V}, β = 1 0^{- 3}

其中 $L_{D B}$ 是 Dice-BCE 分割损失， $L_{E V}$ 是专家使用均衡的正则项。该项并不要求每个样本平均使用所有专家，而是在整体训练中保持专家库的可训练性和选择空间。

第二部分是 Unused Experts Handling。当某些专家没有进入当前 router 的 TopK 集合时，论文使用随机训练样本额外驱动这些未选专家计算并更新参数。其目标是避免未选专家在稀疏路由下长期没有有效训练信号。

训练采用两阶段策略：首先以原始 U-Net 跳跃连接训练编码器和解码器；随后冻结这两个主干，仅训练 TA-MoSC。该策略将主干表征学习与跳跃连接重分配分离，有助于降低联合优化的难度。

GlaS 上的专家特征与不同三专家组合。不同组合产生不同的中间特征和分割输出。图源：Luo et al., AAAI 2025。

实验结果：跨模态任务上的性能与代价

实验覆盖四类医学分割数据：GlaS 腺体分割、MoNuSeg 细胞核分割、ISIC16 皮肤病灶分割，以及 Synapse 腹部 CT 多器官分割。前三类任务报告 Dice 和 IoU，Synapse 额外使用 95% Hausdorff Distance，简称 HD95。

GlaS、MoNuSeg 与 ISIC16 的五折交叉验证结果。图源：Luo et al., AAAI 2025。

在三个二维任务上，UTANet 均取得最高 Dice：

数据集	UTANet Dice	次优方法 Dice	Dice 差值
MoNuSeg	79.10	78.94，R34UNet	+0.16
GlaS	92.00	91.14，R34UNet	+0.86
ISIC16	91.63	91.32，U-Net v2	+0.31

Synapse 上，UTANet 获得 88.23 ± 0.87 Dice 和 28.13 ± 0.21 HD95；对比的最佳已有 Dice 为 87.26 ± 0.74，最佳已有 HD95 为 31.28 ± 0.87。因此该方法不仅提高区域重叠，也降低了边界距离误差。

参数与计算量也需要同时解读。UTANet 的参数量为 24.17M、FLOPs 为 64.41G；原始 U-Net 为 14.8M、50.3G。它显著轻于 UNet++、Attention U-Net 和 TransUNet，但并非没有额外成本。更准确的结论是：该方法以中等额外计算换取在多个基线之上的稳定增益。

TopK 的任务依赖性

论文在推理阶段比较了不同 $K$ 。对 MoNuSeg、GlaS 和 ISIC16 三个单标签任务， $K = 3$ 最好；例如 GlaS 的 Dice 从 $K = 1$ 时的 89.05 上升到 $K = 3$ 时的 92.00，当 $K = 4$ 时又回落至 91.18。

其含义是： $K$ 过小时，解码阶段获得的补充信息不足； $K$ 过大时，低相关专家输出也会混入，增加后续筛选难度。Synapse 是多器官多标签任务， $K = 4$ 取得最佳 Dice 88.09，说明更复杂的多类别场景可能需要保留更完整的专家信息。

这一结果支持本文的核心判断：跳跃连接的最佳稀疏程度不是固定超参数，而与任务的语义复杂度有关。论文将不同 $K$ 描述为推理阶段设置，但未充分说明每个 $K$ 是否以独立训练模型获得，因此该消融更适合解释为路由输出对稀疏程度的敏感性，而非完整的训练策略比较。

局限与适用边界

第一，方法的主要证据来自四个医学分割数据集，尚未充分验证在自然图像分割、目标检测等其他密集预测任务上的可迁移性。论文指出这些方向可以扩展，但尚未提供对应实验。

第二，专家的语义分工主要由特征可视化和性能变化间接支持。论文展示了不同专家组合产生不同特征，但没有建立每个专家与特定医学结构、边缘类型或尺度模式之间的可量化对应关系。

第三，UTANet 相比原始 U-Net 增加了参数量和 FLOPs。TopK 的稀疏机制在理论上有助于控制专家计算，但论文没有单独报告不同 $K$ 下的端到端延迟、显存占用和吞吐量。

第四，四个 router 是按解码阶段独立学习的，而非跨任务联合训练的多任务系统。因此，本文所称的 task-adaptive 主要表现为模型在不同数据集上可学习到不同路由偏好，以及 router 随输入特征动态输出权重；并不表示单一模型已经在多数据集之间完成统一任务路由。

总结

UTANet 的贡献在于将 U-Net 的跳跃连接从固定通路改写为条件化的信息分配机制：

多尺度编码特征
  -> 统一特征池
  -> 共享轻量专家
  -> 解码阶段专属 router
  -> TopK 稀疏加权融合
  -> 尺寸与通道对齐后送入 decoder

该设计没有否定跳跃连接对空间细节恢复的作用，而是指出“传递细节”本身仍需根据解码阶段和任务需求进行选择。实验结果表明，该策略在病理、皮肤镜和 CT 分割任务上获得了稳定提升；同时，专家均衡训练、TopK 设置和额外计算成本也是实际复现与部署时不可忽略的条件。