UniX:统一自回归与扩散模型的胸部X光影像理解与生成 UniX: Unifying Autoregression and Diffusion for Chest X-Ray Understanding and Generation
双分支架构解耦理解与生成,跨模态自注意力实现协同胸部X光分析
前置知识
自回归模型(Autoregressive Model)
自回归模型是一类基于序列建模的生成模型,其核心思想是将联合概率分布分解为条件概率的链式乘积 $p(x) = \prod_{i} p(x_i | x_{<i})$。在本文中,自回归模型被用于理解分支,通过逐token预测的方式生成医学报告。具体实现上,模型将视觉token和文本token拼接成统一序列 $S = [V, T_{in}, T_{out}]$,然后通过交叉熵损失 $\mathcal{L}_{CE} = -\sum_{i=m}^{n-1} \log p(S_{i+1} | S_{\leq i}; \omega_u)$ 训练模型预测下一个token。这种设计天然适合医学报告生成任务,因为报告本身具有顺序结构,模型需要基于视觉和文本上下文进行因果推理。
自回归分支是UniX理解能力的核心,掌握其原理才能理解模型如何将影像信息转化为结构化的医学报告,以及与扩散分支的协同机制
扩散模型(Diffusion Model)
扩散模型是一类基于概率的生成模型,通过逐步向数据添加高斯噪声(前向过程)然后学习去噪(反向过程)来生成高质量样本。本文采用潜空间扩散框架,在VAE编码的潜空间而非像素空间进行扩散,显著提高效率。模型通过最小化均方误差 $\mathcal{L}_{MSE} = \mathbb{E}_{t,p_t(x)} \|v_t(x; \omega_g) - u_t(x)\|^2$ 学习估计目标速度场 $u_t(x)$,其中 $v_t$ 是模型预测的速度场,$\omega_g$ 是生成分支参数。潜空间扩散避免了直接在高维像素空间操作的计算开销,同时保留了生成高质量医学影像的能力。
扩散分支负责高保真图像生成,理解扩散模型的工作原理对于认识UniX如何生成逼真的胸部X光图像至关重要
跨模态自注意力机制(Cross-Modal Self-Attention)
跨模态自注意力是本文提出的核心创新机制,用于连接理解分支和生成分支。不同于传统的交叉注意力(用静态上下文条件化一种模态),该机制对统一的多模态token序列执行联合自注意力。具体实现上,定义统一序列 $S = [T_{in}, N]$,其中 $T_{in}$ 是理解分支产生的文本token,$N$ 是生成分支的噪声条件潜变量嵌入。对于每个token $S_i$,通过模态选择器 $\delta_u(i)$ 和 $\delta_g(i) = 1 - \delta_u(i)$ 进行模态特定投影:$\{Q_i, K_i, V_i\} = \delta_u(i) W_{q,k,v}^{u} S_i + \delta_g(i) W_{q,k,v}^{g} S_i$。最终注意力计算为 $\text{Attn}(S) = \text{softmax}(QK^\top / \sqrt{d}) V$,所有跨模态交互通过联合注意力分数隐式学习。
这是UniX实现理解与生成协同的关键机制,理解它才能明白语义信息如何动态指导图像生成过程
VAE(变分自编码器)
VAE是一种生成模型,由编码器和解码器组成,将数据映射到低维潜空间再重建。在本文中,生成分支使用16倍下采样、16通道的VAE进行图像编码和解码。VAE将原始图像压缩到16维潜空间,然后通过两个单层MLP实现16维潜空间与2048维语言空间之间的双向投影。这种设计使得扩散过程在紧凑的潜空间中进行,大幅降低计算复杂度,同时保持足够的表达能力来重建医学影像的精细结构。
VAE是生成分支的关键组件,理解其压缩和重建机制对于认识UniX如何高效生成高分辨率医学图像很重要
医学基础模型(Medical Foundation Model)
医学基础模型是指在大规模医学数据上预训练、可适配多种下游任务的通用模型。本文关注两类:单任务模型(如专注于理解的LLaVA-Med、专注于生成的RadEdit)和统一模型(如LLM-CXR、HealthGPT)。现有统一模型通常采用参数共享和多任务头策略,但理解任务需要语义抽象而生成任务需要像素级重建,这种目标冲突导致特征干扰,限制整体性能。UniX提出通过架构解耦解决这一根本矛盾。
理解医学基础模型的发展现状和挑战,才能认识到UniX研究的动机和贡献价值
研究动机
当前医学基础模型在统一理解与生成方面面临根本性挑战。理解任务(如疾病诊断、报告生成)需要语义抽象,将视觉信息压缩为高层语义表示;而生成任务(如图像合成)需要像素级重建,保留和恢复精细的结构细节。这两个本质上相互冲突的目标在现有统一模型中通过参数共享学习,导致特征干扰和性能下降。以LLM-CXR为代表的现有方法采用参数共享和联合多任务头进行集成学习,不幸的是,这种方法会引入任务竞争,导致理解或生成性能(或两者)退化。HealthGPT通过H-LoRA模块缓解了这一问题,提供了结构化的折中方案,但并非根本解决方案。此外,大多数当前统一医学基础模型仍然依赖离散生成范式,其输出受词汇粒度约束,无法恢复医学图像中的精细结构细节。在具体性能上,LLM-CXR(12B参数)在理解任务上的Micro F1仅为36.0,在生成任务上的FD-RadDino为71.243,远未达到单任务模型的水平。
本文的目标是本文的目标是提出一种下一代统一医学基础模型UniX,能够在胸部X光领域同时实现高质量的理解和生成,且不产生性能妥协。具体而言,UniX旨在:(1)通过架构解耦消除理解与生成之间的目标冲突;(2)利用扩散模型的连续生成能力替代离散tokenization,捕获医学图像的细粒度结构;(3)引入跨模态协同机制,使理解特征能够动态指导生成过程;(4)在显著减少参数量的情况下(相比LLM-CXR减少75%),达到甚至超越单任务模型的性能水平。最终目标是为医学基础模型建立一个可扩展的协同理解与生成范式。
与已有工作不同的是,本文的独特切入角度体现在三个方面。首先,在架构设计上,UniX是首个将自回归和扩散范式在医学影像领域进行结构性解耦的模型,不同于以往通过参数共享或LoRA适配器实现的表面统一。其次,在生成范式上,UniX采用连续的潜空间扩散替代离散的视觉词袋生成,避免了离散方法固有的粒度损失,能够捕获连续、细粒度的病理模式。第三,在协同机制上,不同于简单的级联或堆叠,UniX通过跨模态自注意力机制实现理解与生成的深度协同,允许语义表示直接调制生成轨迹,同时允许生成状态反馈到语义空间。这种设计借鉴了BAGEL关于无瓶颈多模态交互价值的洞见,在医学影像领域实现了真正的架构协同而非形式统一。
核心方法
UniX的整体方法采用双分支解耦架构,通过跨模态自注意力机制实现协同。直觉上,理解任务和生成任务虽然目标不同,但存在潜在的协同关系——语义推理可以指导视觉合成,视觉合成可以增强语义理解。技术路线如下:首先,构建一个自回归理解分支,基于视觉语言模型处理语义抽象和报告推理;其次,构建一个扩散生成分支,直接继承理解分支的LLM骨干网络,专注于合成高保真胸部X光图像;最后,引入跨模态自注意力模块连接两个分支,实现动态特征交换和语义条件化。为充分利用这一架构的潜力,设计了严格的数据清洗管线和三阶段训练策略,通过渐进式冻结分支确保高效知识迁移和稳定训练。
UniX的核心创新在于通过架构解耦解决理解与生成之间的根本冲突,同时保持两者的协同作用。与已有方法的本质区别体现在:(1)现有统一模型(如LLM-CXR)采用参数共享的联合学习,导致理解需要压缩信息而生成需要保留细节的冲突;UniX通过完全分离的双分支架构消除这一冲突,理解分支专注语义,生成分支专注像素。(2)现有方法大多依赖离散生成(视觉词袋),将连续像素数据压缩到固定码本,不可避免地丢失高频细节和微妙纹理变化;UniX采用潜空间扩散,捕获连续的细粒度病理模式。(3)不同于HealthGPT的H-LoRA模块提供的结构化折中,UniX通过跨模态自注意力实现真正的架构协同——语义表示直接调制生成轨迹,生成状态反馈到语义空间,这种双向信息流是之前方法所不具备的。
方法步骤详情
UniX的方法包含三个关键步骤:数据处理、模型构建和三阶段训练。数据处理阶段,使用DeepSeek大语言模型对MIMIC-CXR数据集的原始报告进行清洗,去除下划线、技术元数据和对话填充词等非诊断信息,构建更语义密集的训练目标。模型构建阶段,理解分支采用SigLIP-large-patch16-384作为视觉编码器产生1024维嵌入,通过两层MLP映射到2048维LLM空间,语言骨干包含24层Transformer并采用QK归一化和QKV偏置;生成分支使用16倍下采样、16通道VAE进行编码解码,两个单层MLP提供16维潜空间与2048维语言空间的双向投影。三阶段训练策略:Stage 1冻结生成分支,仅微调视觉编码器、视觉连接器和语言模型骨干,使用384x384分辨率的图像-报告对进行监督微调,学习率$10^{-4}$,3840步;Stage 2冻结理解分支,预训练生成分支,使用256x256分辨率的文本-低分辨率图像对,应用Representation Alignment将语言模型第8层隐藏状态与RadDino图像特征对齐,学习率$2\times10^{-4}$,75K步;Stage 3保持相同冻结策略,使用512x512分辨率的文本-高分辨率图像对微调生成分支,扩展位置编码并移除特征级监督,学习率$10^{-4}$,5K步。
技术新颖性
UniX的技术新颖性体现在多个层面。首先,它是医学影像领域首个将自回归和扩散范式进行结构性解耦并协同工作的模型,之前的统一模型要么采用参数共享(如LLM-CXR),要么采用LoRA适配(如HealthGPT),都是在共享架构上做任务分离,而UniX从架构层面彻底分离了两个任务。其次,跨模态自注意力机制是原创性贡献,不同于传统交叉注意力用静态上下文条件化一种模态,该机制对统一多模态token序列执行联合自注意力,模态选择器 $\delta_u(i)$ 和 $\delta_g(i)$ 为理解token和生成token提供不同的参数空间,同时保持共享的注意力操作,这种设计使得跨模态交互通过联合注意力分数隐式学习,而非显式条件化。第三,三阶段渐进式训练策略通过合理冻结分支实现了高效知识迁移:Stage 1建立理解能力作为语义特征提供者,Stage 2在低分辨率上建立基础生成能力并应用Representation Alignment加速收敛,Stage 3在高分辨率上精调以获得更好的报告-图像对齐和更清晰的病灶描绘。最后,结合严格的数据清洗管线,整个系统在仅1.5B参数下达到了12B参数模型的性能,展示了架构设计对参数效率的重要作用。
实验结果
UniX在两个代表性基准测试上取得了显著成果。在理解任务上,基于MIMIC-CXR数据集使用CheXbert F1评分评估,UniX(1.5B参数)在"不确定"作为阴性时的Micro F1-14达到53.6,Micro F1-5达到56.6,Macro F1-14为33.2,Macro F1-5为47.3;在"不确定"作为阳性时,Micro F1-14为52.6,Micro F1-5为57.9,Macro F1-14为35.5,Macro F1-5为49.8。相比LLM-CXR(12B参数)的36.0 Micro F1-14(阳性标准),UniX实现了46.1%的提升。相比HealthGPT(3.8B参数)的24.2 Micro F1-14,UniX提升了121.5%。在生成任务上,UniX在512x512分辨率下的FD-RadDino为54.022(越低越好),相比LLM-CXR的71.243(256x256分辨率)提升了24.2%,且与专门的生成模型Sana(54.225)性能相当甚至略优。在14种病理类型的具体生成评估中,UniX在13个类别中取得最佳性能,展示了强大的细粒度视觉合成能力。消融实验表明,冻结理解分支能快速获得生成增益且不损害理解能力,而解冻理解分支且不提供理解数据会严重退化理解性能(Micro-F1从53.2降至13.9)。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 胸部X光报告生成(理解) | CheXbert Micro F1-14(不确定为阳性) | UniX: 52.6 | LLM-CXR: 36.0, HealthGPT: 25.5 | 相比LLM-CXR提升46.1%,相比HealthGPT提升106.3% |
| 胸部X光报告生成(理解) | CheXbert Micro F1-14(不确定为阴性) | UniX: 53.6 | HealthGPT: 24.2, LLaVA-Med: 27.2 | 相比HealthGPT提升121.5%,相比LLaVA-Med提升97.1% |
| 胸部X光图像生成 | FD-RadDino(越低越好) | UniX (512px): 54.022 | LLM-CXR (256px): 71.243, Sana (512px): 54.225 | 相比LLM-CXR提升24.2%,与Sana持平 |
| 胸部X光图像生成 | Alignment Score(图文一致性) | UniX (512px): 0.635 | LLM-CXR (256px): 0.319 | 相比LLM-CXR提升99.1% |
| 参数效率 | 模型参数量 | UniX: 1.5B | LLM-CXR: 12B | 参数量仅为LLM-CXR的1/4,性能更优 |
局限与改进
本文存在以下局限性:首先,实验仅在MIMIC-CXR胸部X光数据集上进行验证,模型的泛化能力尚未在其他医学影像模态(如CT、MRI、病理切片)或其他解剖部位上得到验证。其次,三阶段训练策略需要较长的训练时间(Stage 2需要75K步),且每个阶段需要不同的超参数配置(如学习率、分辨率、是否使用REPA),增加了调参的复杂性。第三,数据清洗依赖DeepSeek大语言模型,这引入了额外的计算成本和对闭源模型的依赖,且清洗质量可能影响最终性能。第四,虽然UniX在理解任务上接近单任务模型的性能,但在某些指标上(如Macro F1)与更大的单任务模型仍有差距。第五,生成任务仅在256x256和512x512分辨率下评估,更高分辨率的生成能力尚未探索。最后,跨模态自注意力机制的具体贡献缺乏单独的消融实验,难以量化其相对于简单拼接或交叉注意力的改进幅度。
独立分析的弱点
从独立分析的角度,UniX存在以下可改进的弱点:第一,架构的模块化程度有限——理解分支和生成分支虽然解耦,但共享LLM骨干网络的初始化,这意味着两个分支并非完全独立,潜在的参数纠缠可能影响最优性能。改进方向是探索更彻底的架构分离,如使用不同的预训练骨干或引入知识蒸馏实现分支间的软对齐。第二,跨模态自注意力机制在统一序列上执行全局注意力,计算复杂度为 $O(n^2)$,当序列长度增加时效率下降。可考虑引入稀疏注意力或线性注意力变体。第三,三阶段训练的Stage 2和Stage 3冻结理解分支,这意味着生成分支无法利用理解能力的持续改进。可以探索交替微调或课程学习策略,允许两个分支在后期阶段相互增强。第四,数据清洗管线是离线的、一次性的,无法适应数据分布的变化。可考虑在线数据质量评估和自适应清洗。第五,评估指标主要基于自动指标(CheXbert F1、FD-RadDino),缺乏临床专家的定性评估和临床实用性验证。
未来方向
基于本文成果,未来研究方向包括:第一,将UniX的双分支解耦架构扩展到其他医学影像模态,如CT、MRI、超声和病理切片,验证跨模态泛化能力。第二,探索更高效的跨模态交互机制,如线性注意力、状态空间模型或混合专家(MoE)架构,在保持协同效果的同时降低计算复杂度。第三,引入临床反馈循环,将医生的报告修改和诊断反馈纳入训练信号,实现持续学习和临床对齐。第四,探索UniX在下游临床任务中的应用,如辅助诊断、治疗方案建议和医学教育,评估其实际临床价值。第五,研究理解分支和生成分支之间的知识迁移和增强机制,例如用生成的高质量图像增强理解训练数据,或用理解特征指导更精细的图像编辑。第六,扩展到多模态医学报告生成,整合实验室检查结果、病史等非影像信息。第七,探索模型压缩和部署优化,使UniX能够在临床环境中实时运行。
复现评估
UniX的复现评估如下:代码和模型已在GitHub(https://github.com/ZrH42/UniX)公开发布,这大大降低了复现门槛。数据方面,使用公开的MIMIC-CXR数据集,遵循官方划分,理解分支使用163,344个训练对和2,365个测试对,生成分支使用237,387个训练对和4,352个测试对。算力要求方面,所有模型在8块NVIDIA L20 GPU上进行全参数微调,Stage 1训练3840步,Stage 2训练75K步,Stage 3训练5K步,总计算量较大但属于研究级实验室可承担的范围。复现难度方面,三阶段训练策略需要仔细的超参数调优(见Table 1),特别是REPA损失权重、学习率调度和数据混合比例。数据清洗依赖DeepSeek模型,可能需要API访问权限。整体而言,具备公开代码、数据和详细超参数配置,复现可行性较高,但需要足够的GPU资源和训练时间。
论文图表