← 返回 2026-09-11

UniH3：统一分层同质性与异质性的多合一医学图像修复 UniH^3: Unifying Hierarchical Homogeneity and Heterogeneity for All-in-One Medical Image Restoration

Zhiwen Yang, Jiayin Li, Chengyu Liu, Hui Zhang, Bingzheng Wei, Yan Xu 📅 2026-09-10 👍 14 2026-09-12 18:30

All-in-One 医学图像修复多任务学习注意力机制记忆网络

UniH3用层级同质记忆与异质平衡器统一七种医学图像修复任务并刷新SOTA

前置知识

All-in-One 图像修复

用一个统一模型同时处理多种退化类型和模态的图像修复任务，无需为每个任务单独训练模型。视觉领域的代表方法包括 AirNet（对比学习生成任务表征作为提示）、PromptIR/AdaIR（可学习提示注入）、TransWeather（多任务头）等，医学领域的先行者是采用混合专家（MoE）的 AMIR。

本文是医学领域的 All-in-One 工作，理解该范式才能明白它的基线选择（AMIR、PromptIR、AdaIR、AirNet）与实验设定的意义。

交叉注意力（Cross-Attention）

Query 来自一个序列、Key/Value 来自另一序列的注意力机制，形式为 $\text{CrossAttention}(Q,K,V)=\text{Softmax}(QK^{\mathsf{T}}/\sqrt{C'})V$，常用于以一组可学习查询去聚合另一组特征。

H2M 的蒸馏与检索两个阶段全部由交叉注意力实现：蒸馏时用可学习原型 $P$ 查询 HQ 特征，检索时用 LQ 特征查询记忆库 $M$，它是全文最核心的算子。

不确定性加权多任务学习

Kendall 等提出的损失平衡方法：为每个任务学习一个不确定性标量 $\sigma_t$，以 $\frac{1}{2\sigma_t^2}\mathcal{L}^{(t)}+\log\sigma_t$ 加权各任务损失——损失偏大的任务 $\sigma_t$ 自动升高从而降权，$\log$ 项防止 $\sigma_t$ 无限增大，从而缓解多任务梯度冲突。

本文的 H2B 正是把这一思想从任务级扩展到样本级（$\sigma_{t,s}=\sigma_t+\Delta\sigma_s$），读懂 H2B 的动机需要这个背景。

EMA 动量更新

指数滑动平均 $M\leftarrow\alpha M+(1-\alpha)V$，以动量系数 $\alpha$ 缓慢吸收新信息、抑制单批次噪声。本文取 $\alpha=0.99$，让记忆库在数十万批次中渐进积累干净、稳定的解剖先验。

H2M 的记忆库完全依赖 EMA 从 HQ 图像蒸馏更新，这是它能在测试时提供可靠先验、且测试期可关闭蒸馏分支的关键。

转置自注意力（通道注意力）

Restormer 提出的高效注意力变体：沿通道维度计算注意力（$Q,K,V\in\mathbb{R}^{C'\times H'W'}$），复杂度与空间分辨率呈线性而非平方关系，适合高分辨率修复任务。

UniH3 的 HGATB 采用 Restormer 式转置自注意力实现 HGA，理解它才能把握 HGA 的计算开销为何只有 26.33 GFLOPs 这么低。

研究动机

医学图像修复（MedIR）旨在从退化的低质量（LQ）采集中恢复高质量（HQ）图像，典型任务包括 PET/CT 去噪、MRI 超分辨率等。由于 PET、CT、MRI 各自成像物理原理不同，以往研究几乎都为单一模态训练专门模型。单任务范式在多模态并存的真实场景（如 PET/CT、PET/MRI）中暴露两大问题：为某任务训练的模型迁移到其他任务时性能大幅下降；为每个任务单独部署、维护模型成本高昂。医学领域的首个 All-in-One 工作 AMIR 用混合专家（MoE）适配三个任务，视觉领域则流行对比学习（AirNet）、退化分类、视觉提示、MoE 等区分任务的手段。但作者指出这些方法有两个盲区：一是完全忽略了医学图像固有的同质性——不同模态、不同患者间高度共享的解剖结构先验，放弃这些共享知识使得任务数增长时学习难度陡增；二是只处理任务间异质性，无视任务内异质性（不同扫描仪、医院、患者人群造成的分布偏移），粗粒度的任务级建模无法化解任务内部的优化冲突。

本文的目标是本文的目标是构建一个统一框架 UniH3，在任务间与任务内两个层级上同时建模同质性与异质性，实现医学图像 All-in-One 修复。具体拆成三件事：其一，设计能显式蒸馏、存储并按需检索高质量图像中解剖先验的机制，把跨任务共享知识变成模型可直接利用的资产；其二，设计能在优化层面同时化解任务间梯度冲突与任务内分布偏移的损失加权策略，保证多任务均衡收敛；其三，构建大规模基准支撑通用 MedIR 研究——MedIR-2D-500K 含 509,200 对 2D 图像、覆盖 PET/CT/MRI 去噪与超分、X-ray 去噪、OCT 去噪、超声去噪、病理超分七种任务，MedIR-3D-3K 含 3,522 对 3D 体数据、覆盖 PET/CT 去噪与 MRI 超分三种任务。最终目标是在 All-in-One 与 Single-Task 两种设定下均达到 SOTA，并验证单一通用模型的实用价值与可迁移性。

与已有工作不同的是，本文的独特切入是反转了 All-in-One 修复的主流叙事。现有方法（AMIR 的 MoE、PromptIR/AdaIR 的提示、AirNet 的对比学习）都把\u2018区分任务、施加特化处理\u2019作为主线索，本质只做减法——隔离异质性；本文则指出医学图像相对自然图像拥有更强、更可利用的同质性，应当做加法——把共享的解剖先验显式建成\u2018记忆\u2019并在推理时按需检索。这带来两个范式级转变：第一，先验注入方式从\u2018以退化 LQ 特征为主体、先验为补充\u2019（SFT、交叉注意力方案）反转为\u2018以 HQ 同质性先验为学习锚点\u2019，直接降低学习难度、加速收敛；第二，异质性建模从任务级标量细化到样本级——每个样本在共享任务级不确定性的同时获得各自的修正项，从而化解扫描仪、中心、人群差异引起的任务内冲突。同质性与异质性这两个看似对立的维度被统一进一个层级化框架，且都刻意区分\u2018任务间\u2019与\u2018任务内\u2019两个粒度，这是此前工作从未同时做到的。

核心方法

UniH3 的直觉是：修复网络不应从零学习\u2018干净图像长什么样\u2019，而应站在预先蒸馏好的解剖先验之上，只专注去除当前退化。技术路线上，给定 LQ 输入 $I^{LQ}\in\mathbb{R}^{H\times W\times 1}$，先经 3×3 卷积投影为浅层特征 $I^S\in\mathbb{R}^{H\times W\times C}$（$C=48$），再送入 4 级非对称 U 形编解码器，各级 HGATB 数量为 $N_1=2,\,N_2=N_3=3,\,N_4=4$。每个 HGATB 内部串联两个变体层：第一层用同质性引导注意力（HGA）替代标准自注意力建模全局交互，第二层用卷积加 Squeeze-and-Excitation（SE）捕获局部交互。深层特征 $I^D$ 经 3×3 卷积投影为残差图 $I^R$，最终输出 $\hat{I}^{HQ}=I^{LQ}+I^R$。H2M 模块在 U 形结构的 4 个尺度各布置一个（共 4 个），实现多分辨率先验引导；H2B 仅在训练期作用于损失加权。3D 版本 UniH3-3D 将各模块替换为 3D 对应物（$N_1=N_2=1,\,N_3=N_4=5$，$C=16$），其余设置不变。整体上参数量 28.96M、FLOPs 26.33G。

核心创新有三处。第一，H2M 把\u2018干净解剖结构\u2019物化为层级记忆库 $M\in\mathbb{R}^{(T+1)L\times C'}$：$T$ 个任务专属槽存任务内同质性（同模态不同患者的共同解剖），1 个共享槽存跨模态的人体结构表征，其中 $T=7$、$L=128$；训练中以 EMA（$\alpha=0.99$）从 HQ 图像蒸馏写入，测试时蒸馏分支关闭、仅用 LQ 特征检索。第二，HGA 公式 $V^{O}=A[(1-\lambda_1)V+\lambda_1 V^{H}]+\lambda_2(V^{H}-V)$（$A=\text{Softmax}(QK^{\mathsf{T}}/\sqrt{C'})$）：相比把先验直接加到 value 的对称做法（Eq. 6），HGA 用单位阵偏置 $(A-I)V+(A+I)V^H$ 提升高保真先验 $V^H$、压低退化特征 $V$ 的贡献，再以可学习通道权重 $\lambda_1,\lambda_2\in\mathbb{R}^{C'}$ 稳定训练；$\lambda_1=\lambda_2=0$ 时退化为普通自注意力。第三，H2B 将不确定性分解为 $\sigma_{t,s}=\sigma_t+\Delta\sigma_s$，$\Delta\sigma_s$ 由轻量不确定性估计块（UEB）从 $[I_s^{LQ},\text{sg}(\hat{I}_s^{HQ}),I_s^{HQ}]$ 预测，总损失 $\mathcal{L}_{H2B}=\frac{1}{TS}\sum_{t}\sum_{s}\left(\frac{1}{2\sigma_{t,s}^2}\mathcal{L}_{rec}^{(t,s)}+\log\sigma_{t,s}\right)$。

方法步骤详情

流程分四步。第一步特征提取：LQ 图像经 3×3 卷积得 $I^S$，逐级通过 HGATB 编解码器，每级配合对应尺度的 H2M。第二步先验蒸馏（仅训练期）：配对图像经 pixel-unshuffle 下采样加 3×3 卷积得 $F^{LQ},F^{HQ}\in\mathbb{R}^{H'W'\times C'}$；可学习原型 $P\in\mathbb{R}^{L\times C'}$ 作 query 执行 $V^{HQ}=\text{CrossAttention}(P,F^{LQ},F^{HQ})$，让 $P$ 学会聚合 HQ 特征中最具代表性的解剖结构；按当前任务索引选出共享部分 $V^{Sh}$ 与专属部分 $V^{Sp}$，以 $M_{Sh/Sp}\leftarrow\alpha M_{Sh/Sp}+(1-\alpha)V^{Sh/Sp}$ 写入初始化为零的记忆库。第三步先验检索（训练与测试通用）：以 LQ 特征为 query 执行 $V^{H}=\text{CrossAttention}(F^{LQ},P,M)$，得到与输入最相关的干净先验，投影后注入各 HGATB 的 HGA。第四步训练优化：重建损失为 L1，总损失由 H2B 按 $\sigma_{t,s}$ 加权（UEB 输入含 stop-gradient 的预测图以解耦梯度）；Muon 优化器训练 $6\times10^5$ 次迭代，学习率 $3\times10^{-4}$ 余弦退火至 $1\times10^{-7}$，patch 128×128、batch 14；UniH3-3D 用 patch $64^3$、batch 6、学习率 $5\times10^{-5}$。

技术新颖性

技术新颖性体现在与三类既有方案的对照中。与任务区分型 All-in-One 方法相比：AirNet 用对比学习学任务表征、PromptIR/AdaIR 学判别性提示、AMIR 用 MoE 路由，它们只在\u2018异质性\u2019一个维度做文章；本文首次把同质性提升为可蒸馏、可检索、可注入的一等公民，并证明两条线索互补——Table 7 显示基线 36.52 dB，仅加 H2M 达 36.66 dB，仅加 H2B 达 36.64 dB，联合 36.77 dB。与先验注入方式相比：SFT 与交叉注意力把 LQ 特征当主体、先验当补充；本文 HGA 用 $\pm I$ 偏置反转主从关系，Table 8 显示其以 28.96M 参数取得 36.77 dB，优于 SFT（37.46M，36.74 dB）和交叉注意力（29.60M，36.67 dB）。与损失加权策略相比：经典不确定性加权每任务只有标量 $\sigma_t$，无法处理扫描仪/中心差异引起的任务内偏移；H2B 的层级化 $\sigma_{t,s}$ 在 All-in-One MedIR 中首次同时化解两级冲突。此外 Table 6 证明 H2M+H2B 是即插即用模块：移植到 Uformer、Restormer、PromptIR、AdaIR 四个骨干分别带来 +0.17、+0.18、+0.16、+0.19 dB，说明框架与骨干解耦，具有通用组件价值。

**Fig. 2**: Hierarchical Homogeneity Memory.

实验结果

核心结论：（1）2D All-in-One（Table 2）：UniH3 在 MedIR-2D-500K 七个任务全部第一，平均 36.77 dB/0.9045，超第二名 AdaIR（36.56 dB）0.21 dB；参数 28.96M、FLOPs 26.33G，低于 PromptIR（39.49G）。（2）3D All-in-One（Table 4）：UniH3-3D 平均 44.50 dB/0.9674，超 Restore-RWKV-3D（43.95 dB）0.55 dB，MRI 任务 40.03 dB 领先 0.60 dB。（3）单任务设定：2D 平均 36.90 dB 超 MambaIR 0.15 dB（Table 3）；3D 平均 45.29 dB 超 Spach Transformer（43.81 dB）1.48 dB，CT 任务领先 1.77 dB（Table 5）。（4）消融：HGA 优于 SFT（36.74）与交叉注意力（36.67）（Table 8）；H2M 两类同质性槽各贡献 +0.09/+0.20 dB（Table 9）；H2B 两级平衡各贡献 +0.04/+0.07 dB，联合 36.77 dB（Table 10）。（5）可迁移性：Fig. 8 显示 All-in-One UniH3 36.77 dB 已匹敌单任务 MambaIR 36.75 dB，逐任务微调后达 36.97 dB、反超单任务 UniH3 的 36.90 dB，证明其可作预训练骨干。（6）机制证据：Fig. 6 中两个 PET 脊柱 token 的 top-10 槽位重叠 8/10，PET 与 CT 脊柱跨模态重叠 2/10，PET 脊柱与病灶重叠 0/10，直接证明 H2M 按解剖语义存储与检索先验。

**Table 1**: Overview of the MedIR-2D-500K and MedIR-3D-3K datasets.

**Table 2**: All-in-one MedIR comparison results on the MedIR-2D-500K dataset.

**Table 3**: Single-task MedIR comparison results on the MedIR-2D-500K dataset.

**Table 4**: 3D all-in-one MedIR results on the MedIR-3D-3K dataset.

**Table 6**: Performance of H2P and H2B on different backbones on the MedIR-2D-500K dataset.

**Fig. 4**: Visual comparison of methods for all-in-one medical image restoration on the MedIR-2D-500K dataset.

**Fig. 5**: Visual comparison for component analysis.

**Fig. 6**: Retrieval attention map in H2M. The Top-10 scores are marked by red rectangles.

**Fig. 7**: Estimated uncertainty distribution.

查看结构化数据

任务	指标	本文	基线	提升
2D All-in-One 医学图像修复（MedIR-2D-500K，7 任务平均）	PSNR / SSIM	36.77 dB / 0.9045	AdaIR：36.56 dB / 0.9026	+0.21 dB PSNR，七个任务全部第一
3D All-in-One 医学图像修复（MedIR-3D-3K，3 任务平均）	PSNR / SSIM	44.50 dB / 0.9674（UniH3-3D）	Restore-RWKV-3D：43.95 dB / 0.9644	+0.55 dB PSNR
2D 单任务医学图像修复（7 任务平均）	PSNR / SSIM	36.90 dB / 0.9058	MambaIR：36.75 dB / 0.9046	+0.15 dB PSNR
3D 单任务医学图像修复（3 任务平均）	PSNR / SSIM	45.29 dB / 0.9698（UniH3-3D）	Spach Transformer：43.81 dB / 0.9634	+1.48 dB PSNR

局限与改进

作者在讨论中承认的局限：研究只覆盖每个模态的主要修复任务（PET/CT 去噪、MRI 超分等），未涉及同一模态内的其他退化类型（如 CT 伪影、MRI 运动伪影），距\u2018真正通用\u2019的 MedIR 尚有距离。我补充的观察：第一，H2M 的槽位结构与任务数强耦合（$T+1$ 个槽、$T=7$），新增模态或任务需要扩容记忆库并重新训练，扩展成本不低；第二，蒸馏先验来自训练分布的 HQ 图像，跨中心、跨扫描仪部署时记忆库可能过拟合于源域解剖统计，论文未测试跨域泛化；第三，H2B 的 UEB 需要输入 HQ 真值（$I_s^{HQ}$），本质是训练期损失加权技巧，对推理质量只有间接影响，测试时模型对自己的不确定性没有感知；第四，All-in-One（36.77 dB）与单任务（36.90 dB）仍有 0.13 dB 差距，任务间干扰未被完全消除；第五，3D 实验中 CT 测试集仅 30 对体数据，结论的统计强度有限；此外论文未与扩散模型类生成式修复方法对比，而后者在感知质量上可能另有优势。

独立分析的弱点

独立分析四点弱点。其一，记忆库静态性：$M$ 的槽位在训练前固定（7 个任务专属+1 个共享），面对开放世界的连续新任务只能重建整个记忆系统；改进方向是引入参数化槽位生成器或外挂检索库，使新任务无需改动架构即可增量接入。其二，域偏移风险：蒸馏的先验本质是训练集 HQ 图像的统计摘要，若部署医院的扫描协议、设备型号、人群构成不同，检索到的先验可能误导修复；改进方向是加入跨中心训练数据、先验置信度门控，或在检索时联合设备/协议元数据。其三，样本级不确定性依赖真值：$\Delta\sigma_s$ 由包含 $I_s^{HQ}$ 的输入预测，推理时无法使用，意味着模型在测试期对自己的不确定性没有感知；改进方向是训练一个仅依赖 LQ 与预测图的不确定性估计器，兼作推理期质量预警。其四，效率评估不完整：FLOPs 在 128×128 输入下报告，而 3D 版 $64^3$ patch、batch 6 的显存开销与推理时延未给出，临床部署关心的吞吐指标缺失；改进方向是补充分辨率扫描实验，并结合半精度、蒸馏或滑动窗口推理优化。此外，All-in-One 与单任务之间 0.13 dB 的残差差距也提示任务干扰建模仍有提升空间。

未来方向

作者提出的方向：扩展任务谱系，纳入更多模态与退化类型，迈向更通用的 MedIR 模型，服务临床诊断与更多下游任务。基于其成果可延伸的研究：第一，预训练-微调范式——Fig. 8 显示 All-in-One 模型逐任务微调后达 36.97 dB、反超单任务训练的 36.90 dB，可系统研究其作为\u2018医学修复基础模型\u2019的缩放规律，并与医学自监督预训练结合；第二，跨机构泛化——把 H2M 记忆库设计为可联邦更新的共享解剖知识库，在多中心数据上验证检索先验的域鲁棒性；第三，先验扩展——用扩散模型或生成式先验替代确定性记忆，覆盖记忆库未见过的解剖变异与罕见病灶；第四，任务内异质性的显式利用——当前 $\Delta\sigma_s$ 只作用于损失权重，可让检索先验本身也感知扫描仪/协议元数据，实现条件化引导；第五，向 3D 全身影像、动态（4D）影像以及低剂量 CT、加速 MRI 等更多重建链路推进，并在分割、检测、组学提取等下游任务上量化修复质量带来的实际临床收益。

复现评估

复现评估：代码已在 GitHub 开源（github.com/Yaziwel/UniH3）。数据方面，MedIR-2D-500K 与 MedIR-3D-3K 由多个公开数据集与私有收集混合构成，其中 PET 与 CT 部分含 Private 数据，因此完整训练集无法完全获得，但公开部分足以复现大部分任务与方法对比。实现细节披露充分：网络规模（$N_1=2,N_2=N_3=3,N_4=4$，$C=48$）、H2M 配置（$T=7$、$L=128$、$\alpha=0.99$）、训练配置（Muon 优化器、$6\times10^5$ 次迭代、学习率 $3\times10^{-4}$ 余弦退火至 $1\times10^{-7}$、patch 128×128、batch 14、L1 损失）以及 3D 版配置（$C=16$、lr $5\times10^{-5}$、patch $64^3$、batch 6）均有明确数值。算力方面，全量训练约需单卡高端 GPU（A100 量级）连续数天，3D 版显存压力更大；在公开数据子集上用较少迭代复现趋势是可行的。总体难度中等：架构可复现，但 50 万对全量基准与私有数据部分的绝对数字难以严格对齐，复现者应重点关注相对提升而非绝对 PSNR。

论文图表

Fig. 3: Homogeneity-Guided Attention.

HGA 的构造示意：在转置自注意力的值路径上，把常规的 $AV$ 替换为 $A[(1-\lambda_1)V+\lambda_1 V^H]+\lambda_2(V^H-V)$，仅用加法和减法操作就实现了\u2018偏向干净同质性先验\u2019的注意力偏置，图中标注了通道级可学习权重 $\lambda_1,\lambda_2$。

HGA 是先验注入机制的核心，此图直观呈现了它与标准自注意力（$\lambda_1=\lambda_2=0$ 退化）以及 Eq. 6/7 两个中间变体的递进关系。

Fig. 8: All-in-One vs. Single-Task.

柱状图比较四种设定在 MedIR-2D-500K 上的平均 PSNR：All-in-One UniH3 36.77 dB、单任务 MambaIR 36.75 dB、单任务训练的 UniH3 36.90 dB、All-in-One UniH3 逐任务微调（迁移）后 36.97 dB。

支撑论文的核心论点：单一 All-in-One 模型性能已可匹敌专门模型，且微调后还能反超，说明它可作为可迁移的预训练骨干，这是 All-in-One 路线实用性的关键证据。

Table 5: 3D single-task MedIR results on the MedIR-3D-3K dataset.

3D 单任务设定对比：UniH3-3D 平均 45.29 dB/0.9698，超第二名 Spach Transformer（43.81 dB）1.48 dB，其中 CT 任务 45.37 dB 领先 1.77 dB，是全文最大幅度的单项提升。

与 Table 4 对照可见 3D 上 All-in-One 与单任务差距（0.79 dB）比 2D 更大，同时证明 UniH3-3D 架构在 3D 场景的优势尤为显著。

Table 9: Ablation studies on H2M.

H2M 槽位消融：将任务专属/共享槽替换为普通可学习参数（脱离 HQ 蒸馏）作基线 36.52 dB；仅保留任务间同质性 36.61 dB，仅保留任务内同质性 36.72 dB，两者兼备 36.77 dB。

验证层级同质性设计的必要性：两类先验各自有效且组合最优，其中任务内同质性贡献更大（+0.20 dB）。