MedSteer: 基于免训练激活引导的反事实内窥镜图像合成 MedSteer: Counterfactual Endoscopic Synthesis via Training-Free Activation Steering
免训练激活引导生成配对内窥镜反事实图像
前置知识
扩散 Transformer (DiT)
DiT 用 Transformer 替代 U-Net 作为扩散模型的骨干,把图像切成 patch token,在每一层用自注意力、交叉注意力和前馈网络构造残差流。PixArt-$\alpha$ 是 DiT 的文本到图像实现,共 28 层,本文就在它的交叉注意力输出处做手术。
本文的核心干预点正是 DiT 的 cross-attention 输出,理解 token 维度 $h_{l,t} \in \mathbb{R}^{d_l}$ 以及 layer 索引 $l$ 和 timestep 索引 $t$ 的含义,是读懂公式 $h'_{l,t} = h_{l,t} - \alpha \sigma_{l,t} v_{l,t}$ 的前提。
DDIM 反演 (DDIM Inversion)
把真实图像沿 ODE 轨迹反向加噪到纯噪声,再用编辑后的 prompt 重新去噪,试图只改语义不改结构。由于 ODE 反演只是近似,即便用更精确的调度器,重建误差也无法被完全消除,带来结构漂移 (structural drift)。
PnP、h-Edit、RadEdit 等基线都依赖此流程,它们在背景保持 (Bg-LPIPS、Bg-SSIM) 上的劣势正是源于反演误差。MedSteer 跳出 generate→invert→edit 范式,直接用共享噪声种子从零生成,因此天然消除该误差。
交叉注意力 (Cross-Attention, CA) 与反事实激活方向
在文生图扩散里,文本语义只能通过 cross-attention 注入图像 token,因而 CA 输出是文本概念的瓶颈。通过计算多个正负 prompt 对的 CA 特征均值差,再 L2 归一化,即可得到一个单位方向 $v_{l,t}$,它编码有 P 与无 P 的语义差异,即病理方向。
这是 MedSteer 的离线估计阶段,均值差假设 (mean-difference principle) 让方向只与概念有关、与图像无关,可以一次估计、永久复用,且天然正交于解剖/纹理等非目标属性。
余弦相似度门控 (Cosine-Similarity Gate, CSG)
对每个视觉 token 计算它与病理方向 $v_{l,t}$ 的对齐度 $\sigma_{l,t} = \max(\langle h_{l,t}, v_{l,t}\rangle, 0)$。去掉 $1/\|h_{l,t}\|_2$ 归一化是有意为之:对齐越强的 token 越应被强力干预;负对齐的 token 完全不动,保留所有正交结构。
这是 MedSteer 实现空间选择性的关键。直接做全局标量减法会过度抑制非目标 token,CSG 让干预天然聚焦在病灶区域,顺便提供每步扩散的可解释热力图 (Fig. 3)。
研究动机
内窥镜下训练可靠的病灶检测器高度依赖病灶-正常配对样本,但临床上同一病人的配对几乎拿不到,只能借助生成模型合成。一条朴素路径是 re-prompting:对同一噪声种子,把 prompt 里的息肉改成正常盲肠重新去噪,然而整条扩散轨迹被重新摇号,解剖、视角、纹理全部跟着变,无法构造真正的最小编辑对。另一条主流路径是 inversion-based editing (PnP、h-Edit、RadEdit):先用 DDIM 反演把源图映射回噪声,再换 prompt 去噪。DDIM 反演本质是 ODE 的近似,即便用更精确的调度器,累积误差也使背景像素的 LPIPS、SSIM、PSNR 持续下降;RadEdit 还要逐图标注 mask,临床部署成本高。更棘手的是染色息肉 vs 未染色息肉这种属性纠缠:text 根本没法把染料靛蓝胭脂红和息肉隆起两个共现概念拆开,现有方法在 dye disentanglement 上的 DDR 普遍高于 0.80,等同于不工作。
本文的目标是本文提出 MedSteer,目标是在冻结的 DiT 上以零训练、零源图、零 mask 的方式,生成结构完全一致、只改单一临床概念的反事实配对图像;并在三个层面验证价值:(1) 多种临床概念对 (息肉↔正常盲肠、溃疡性结肠炎↔正常盲肠、食管炎↔正常 Z 线、染色息肉→息肉) 的高翻转率与高背景保真度;(2) 染色解缠性能远超基线;(3) 用合成反事实增广后,真实测试集 (HyperKvasir) 上的息肉检测器 F1 和 AUC 显著提升。
与已有工作不同的是，现有 steering 类工作 (SDID、SAeUron、概念擦除等) 要么依赖额外训练、要么 fine-tune 权重或 LoRA,要么没有验证临床下游任务;inversion-based editing 又天生被 DDIM 近似拖累。MedSteer 的独特切入点是共享种子 + 路径外引导——两条分支用同一噪声 $z$ 走同一条 ODE 轨迹,只在 cross-attention 处按 CSG 选择性减去病理方向分量,使得非目标结构在数学上完全相同,而代价仅为每概念离线估计一次病理向量,推理时一次点积和一次减法。
核心方法
直觉上,扩散模型在 cross-attention 中编码了该词对图像的作用。如果两个 prompt 只差一个临床概念,它们的 CA 特征均值差就是一个病理方向。推理时,沿这个方向正负移动图像 token,就相当于加上或减去该概念。MedSteer 把这件事做成一个推理时的即插即用模块:离线阶段对每个概念 P 准备 Z 对正负 prompt,在 $l \in \{L_s,\dots,L_e\}$、所有去噪步 $t$ 上抓 CA 特征,做均值后 S&N 得到单位病理向量 $v_{l,t}$;推理阶段两条分支共享噪声种子 $z$ 和正向 prompt,Unsteered 分支原样运行 DiT,Steered 分支在每个目标层/步执行 SSPS:先算每 token 的余弦对齐 $\sigma_{l,t} = \max(\langle h_{l,t}, v_{l,t}\rangle, 0)$,再更新为 $h'_{l,t} = h_{l,t} - \alpha\sigma_{l,t} v_{l,t}$。
核心创新是把反事实生成与扩散反演彻底解耦,改用共享种子的激活引导实现结构保真。具体包含三点:(1) 干预点选在 CA 输出而非 UNet 的 self-attention,因为 CA 是文本注入图像的唯一通道,方向干净;(2) 引入 per-token 余弦门控 $\sigma_{l,t}$ 做空间选择,只对齐的 token 被减去方向分量,正交成分 (解剖、纹理、视角) 完整保留,且门控图天然可解释;(3) 推理时两条分支共享噪声 $z$,结构差异仅来自 SSPS,故差异必为最小编辑。本质区别于 inversion-based 方法:后者在轨迹上绕路 (反演→重去噪),必然引入误差;MedSteer 始终在同一条轨迹上偏置 token,误差在数学上为零 (非目标结构完全相同)。
方法步骤详情
完整流程分三阶段。第一阶段骨干与干预点 (Sec. 3.1):使用 PixArt-$\alpha$ DiT (N=28),冻结全部权重;只在 $l \in \{8,\dots,16\}$、所有 $T$ 步的 cross-attention 输出上插入 SSPS,层窗经验确定。第二阶段病理向量估计 (Sec. 3.2,离线一次性):对每个临床概念 P (如 Polyp),构造 Z 对正负 prompt,只差 P 的有无;对每对 prompt 用不同种子做 Z 次前向,在 layer $l$、timestep $t$ 抓出 CA 特征 $h_{l,t,n}$,对所有种子和空间 token 求均值 $\bar h^{pos}_{l,t},\bar h^{neg}_{l,t} \in \mathbb{R}^{d_l}$,再 S&N 得 $v_{l,t} = (\bar h^{pos}_{l,t} - \bar h^{neg}_{l,t}) / \|\bar h^{pos}_{l,t} - \bar h^{neg}_{l,t}\|_2$。该 $v_{l,t}$ 与图像无关,可永久复用。第三阶段 SSPS (Sec. 3.3,推理时):输入是层 $l$、步 $t$ 的激活 $h_{l,t}$ 与离线算好的 $v_{l,t}$;先算 $\sigma_{l,t} = \max(\langle h_{l,t}, v_{l,t}\rangle, 0)$ 作 CSG;再 update 为 $h'_{l,t} = h_{l,t} - \alpha \sigma_{l,t} v_{l,t}$,其中 $\alpha$ 控制强度。负号表示去掉概念 (例如染染料→息肉),正号表示注入概念。两条分支共用 $z$ 走同一 ODE 轨迹,差异完全可归因于 SSPS。
技术新颖性
技术新颖性体现在四个层面:(1) 框架新颖——把概念控制从生成-反演-编辑的三步范式,重写为单次生成 + 路径内引导的免训练范式,绕过 DDIM 反演误差这一长期瓶颈。(2) 干预点新颖——首次系统地把 steering 落到 DiT 的 cross-attention 输出而非传统 UNet 的 self-attention 块,论文解释 CA 是文本进入图像的唯一瓶颈,故干预最干净。(3) 门控函数新颖——抛弃全局标量减法,采用 per-token 余弦对齐 $\sigma_{l,t}=\max(\langle h_{l,t},v_{l,t}\rangle,0)$,既实现空间选择性,又生成天然可解释热力图 (Fig. 3),把可解释性从后验分析变成推理副产物。(4) 评估新颖——同时报告下游检测器 (ConvNeXt / ViT 在 HyperKvasir 的 F1、AUC) 的 OOD 增益,以及背景-病灶分离的 LPIPS/SSIM/PSNR,这是以往纯感知指标之外的临床导向验证。
实验结果
实验在 Kvasir v3 上训练/评估,HyperKvasir 作 OOD 测试集。反事实生成 (Table 2):三对概念 (A) Polyp↔Normal Cecum、(B) UC↔Normal Cecum、(C) Esophagitis↔Normal Z-line 上,MedSteer 的 Flip 分别达 0.800/0.925/0.950,$\Delta p$ 达 0.721/0.902/0.928,全面优于 PnP (0.627/0.682/0.718) 与 h-Edit (0.714/0.741/0.796);且 A 对上 Bg-LPIPS=0.1449、Bg-SSIM=0.8951、Bg-PSNR=22.95,均为三方法最优,证明共享种子确实保住了背景。染色解缠 (Table 3):MedSteer 的染料检测率 (DDR) 仅 0.250,即 75% 染料被成功剥离,而 h-Edit 与 PnP 分别为 0.900 和 0.800,等效解缠率仅 10% 和 20%;UNet++ 分割下的有效 LPIPS=0.2027、SSIM=0.6584、PSNR=15.57,均显著领先。下游息肉检测 (Table 1):用 1,000 张合成图增广训练后,MedSteer 配 ConvNeXt 得 F1=0.9263、AUC=0.9341,配 ViT 得 F1=0.9591、AUC=0.9755,均高于 Real-only (ViT AUC=0.8942)、Re-prompting (0.9083)、PnP (0.9518) 与 h-Edit (0.9312),证明反事实结构才是性能提升的真正来源,而非数据量。消融 (Table 4):(a) 层窗实验确认 8–16 为语义形成区,0–8 与 12–20 的 Flip 接近 0;宽度同为 12 时 8–20 达 0.700 而 0–12 仅 0.377,显示位置效应。(b) 强度 $\alpha$ 在 2.5 处最优 (0.800),$\alpha=2.0$ 仅 0.513,$\alpha=3.0$ 退化为 0.700,说明过强会破坏非目标结构。(c) 种子数 50 即饱和 (Flip 0.800),100 反而微降到 0.791。可解释性 (Fig. 3):$\sigma_{8,t}$ 热力图在 t=8 时几乎覆盖整张图,t=19 时收缩到稀疏斑块,符合早步定结构、晚步改细节的扩散直觉。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Counterfactual generation: Polyp → Normal Cecum | Flip rate (↑) | 0.800 | h-Edit 0.714, PnP 0.627 | +0.086 vs h-Edit, +0.173 vs PnP |
| Counterfactual generation: Polyp → Normal Cecum | Bg-LPIPS (↓) | 0.1449 | h-Edit 0.1565, PnP 0.1518 | 背景感知距离最低,保真度最佳 |
| Counterfactual generation: Ulcerative Colitis → Normal Cecum | Flip rate (↑) | 0.925 | h-Edit 0.741, PnP 0.682 | +0.184 / +0.243,跨解剖部位泛化 |
| Counterfactual generation: Esophagitis → Normal Z-line | Flip rate (↑) | 0.950 | h-Edit 0.796, PnP 0.718 | +0.154 / +0.232,跨食管/结肠 |
| Dye disentanglement (UNet++ 分割) | Dye Detection Rate DDR (↓) | 0.250 | h-Edit 0.900, PnP 0.800 | 解缠率 75% vs 10% / 20%,绝对提升 55–65 pp |
| Dye disentanglement | Effective SSIM (↑) | 0.6584 | PnP 0.1872, h-Edit 0.0957 | 提升约 3.5× / 6.9× |
| Downstream polyp detection (ViT, OOD HyperKvasir) | AUC (↑) | 0.9755 | Re-prompting 0.9083, Real-only 0.8942, PnP 0.9518, h-Edit 0.9312 | +0.0672 vs Re-prompting,验证反事实结构价值 |
| Downstream polyp detection (ConvNeXt, OOD HyperKvasir) | F1 (↑) | 0.9263 | Real-only 0.8977, Re-prompting 0.9038, PnP 0.9174, h-Edit 0.9097 | 全部基线之上,F1 提升 0.009–0.029 |
局限与改进
作者在结论中点出几个明显局限:(1) 全部实验限于 2D 静态内窥镜图像,尚未扩展到 3D 体数据或视频序列,而真实临床场景是时序的;(2) 评估数据集局限于 Kvasir v3 / HyperKvasir 同一系统的胃肠图像,跨机构 / 跨设备 / 跨癌种 (如食管癌、胰腺 EUS) 的迁移性未验证;(3) 病理向量估计需要约 50 次前向 + Z 对 prompt,虽然可复用但仍依赖基础模型有较强的跨种子稳定性,在更稀疏的数据上未必成立。我自己的额外观察:(a) 推理时虽然不微调,仍要保存所有 $l \in \{8,\dots,16\}$、所有 $t$ 的 $v_{l,t}$,空间开销大;(b) 方向是正负 prompt 的均值差,要求两类样本在 batch 内分布对称,否则 $v_{l,t}$ 会被非概念特征污染,目前对这一假设缺乏敏感性分析;(c) 评估只用了一个 Oracle 分类器 (AUC=0.98) 测 Flip,若换成临床报告或不同 backbone 分类器,Flip 数值可能波动;(d) DDR 指标本质依赖分类器对染色的判定,未在论文里给出染色分类器的具体训练细节,外部复现的可信度受限。
独立分析的弱点
独立分析的弱点集中在三方面。第一,概念解缠假设偏强:MedSteer 假设正负 prompt 仅差一个概念 P,因此 $v_{l,t}$ 才能被解读为纯 P 方向;实际临床 prompt 经常隐含多个共现属性 (例如染色隆起息肉同时含染色和隆起),此时 S&N 差向量会被其他无关属性污染,论文在 dye disentanglement 上虽然部分验证了染色方向对息肉方向的正交性,但缺乏对污染比例的量化分析。第二,层窗与强度的耦合未充分研究:Table 4 把层窗和 $\alpha$ 分开扫,但实际部署时二者高度耦合,在另一对概念上沿用默认 8–16 + $\alpha=2.5$ 是否仍最优,论文没有 cross-concept 验证。第三,推理时开销不可忽略:每个目标 token 都要算一次内积 $\langle h_{l,t}, v_{l,t}\rangle$ 和一次向量减,叠加在 9 层 × 数十步上,使单图推理比 Unsteered 慢数倍,在临床实时性场景下是瓶颈;此外所有 $v_{l,t}$ 常驻显存/内存,部署成本不低。改进方向:对第一个问题,可在 prompt 集里加入多个干扰属性的样本后做残差化 (residualization);对第二个问题,加入 per-concept 网格搜索或基于验证集 Flip 的简单控制器;对第三个问题,可以把 CSG 改为稀疏 top-k 选取或低秩近似,只对被门控命中的 token 执行减法。
未来方向
作者明确提出未来方向是 3D 体数据 (CT/MRI 体积)、视频内窥镜序列以及跨机构部署。基于成果可延伸的方向还包括:(1) 把 MedSteer 拓展到多概念组合引导,例如同时换 息肉形态 + 染色,只需叠加多个 $v_{l,t}$ 并各自经过 CSG;(2) 与 DPO / RLHF 结合,让 $\alpha$、层窗根据放射科医生反馈自适应;(3) 借用 MedSteer 的 CSG 热力图作为可解释病灶定位信号,辅助弱监督病灶分割,直接省去昂贵的像素级标注;(4) 把病理向量估计迁移到其他医学模态 (超声、皮肤镜、病理切片),构成统一的医学 steering 库;(5) 与 RAG 知识库结合,在 prompt 构造阶段注入临床指南,提升方向估计的语义准确性。
复现评估
代码明确声明 will be made publicly available upon acceptance,目前未开源,因此完整复现受限。论文给出了充分细节以支撑再实现:骨干为 PixArt-$\alpha$ (公开) + LoRA rank=64 在 Kvasir v3 (公开) 上微调;训练集 Kvasir v3 8,000 张、8 类,测试集 HyperKvasir 已去重,均为公开下载;每概念 50 张正/负 prompt 生成图,seed 数 Z=50 经验最优;层窗 {8,..,16}、$\alpha=2.5$、N=1,000 合成图增广、ConvNeXt + ViT 二分类器、3-seed 误差棒。DDIM 反演基线 (PnP、h-Edit) 来自各自原始仓库,Oracle 分类器 AUC=0.98,UNet++ / Segformer 分割 Dice 分别为 0.912 / 0.921,这些都方便复现。算力门槛:需 28 层 DiT + LoRA 微调,8 GB 以上单卡即可跑推理,但批量生成 1,000 张反事实 + 训练检测器至少需要 24 GB 显存 (A5000/3090 级),个人工作站一周内可完成。复现难度中等偏低,主要风险在染色 vs 息肉方向的正交性依赖 prompt 措辞细节,以及 Oracle 与 DDR 分类器的训练集不公开,可能导致不同实现间 Flip/DDR 数值出现 ±0.05 波动。论文结构 (方法-实验-消融) 清晰,表 1–4 和 Fig. 1–3 完整呈现所有关键数字与可视化,整体可复现性良好,推荐关注官方代码发布。
论文图表
三个子表:(a) 层窗 {0–8, 8–16, 12–20, 0–12, 8–20} 的 Flip / $\Delta p$,8–16 最优 (0.800/0.721),同宽度 12 时 8–20 优于 0–12;(b) $\alpha \in \{0.5, 1.0, 2.0, 2.5, 3.0\}$,$\alpha=2.5$ 达 0.800,$\alpha=3.0$ 退化为 0.700;(c) 种子数 30/50/80/100,50 即可饱和。
可读出全部超参选择的经验依据,理解'为什么是层 8–16、为什么 $\alpha=2.5$、为什么 50 个 seed 就够'的工程逻辑,也是审稿人/复现者最关心的敏感性分析。