TOIS 2025 | 提示微调的多中心药品推荐模型_tois 大模型论文-优快云博客

作者：香港城市大学、西安交通大学刘启东

今天跟大家分享一篇来自西安交通大学、香港城市大学、腾讯优图实验室天衍研究中心、吉林大学和中国科学技术大学的论文，该论文已被TOIS接收。这篇文章提出了一种新颖的提示微调的多中心药品推荐模型，名为TEMPT。该方法首次研究了药品推荐中的多中心问题。

论文地址：https://dl.acm.org/doi/pdf/10.1145/3706631
代码链接：https://github.com/Applied-Machine-Learning-Lab/TEMPT

1. 简介

医生开药通常是治疗病人的一个重要手段，但是面对病人各种各样的疾病，这一过程对医生的专业水平要求非常高并且耗时费力。为了应对这一任务专业性和复杂性需求，当前已经有许多的相关研究寻求开发出自动化的药品推荐系统来辅助医生进行决策。然而，当前的研究集中于为有丰富医疗数据的医院单独开发一个药品推荐系统，而忽视了在真实情况下有许多的医院并没有积累足够多的医疗数据。如图1(a)所示，我们分析了一个多中心的医疗数据库eICU，发现大部分的医院仅有不足1500个样本数据，这对于充分训练一个药品推荐模型是具有很大挑战的。

虽然，使用所有医院的数据训练一个统一的药品推荐模型是一个可能的方案。但是，不同医院因为医疗条件、专业性等差别，开药的分布通常是不同的。如图1(b)所示，不同医院的平均每次开药数量具有很大的差别。此外，为了更近一步验证这一分布差异，我们计算了每两个医院之间开药分布的Jesen-Shannon divergence(JSD)距离，并可视化了所有医院间JSD距离的热力图，如图2所示。从图中可以发现大部分医院对的JSD距离是大于0.4的，甚至很多都大于0.8。这说明了不同医院的开药情况具有很大的不同。

近期，一些多域推荐的工作(MDR)可以缓解上述不同域间分布不同的问题。一类是基于多任务学习的方法，如MMOE和PLE。但是，这些方法经常对域的数量较小的情况更加适配，难以适配多中心医院数量较大的情况，因为通常专家数量是与域的数量相同的。另一类方法是多域CTR预估的工作，比如STAR。但是，这类方法是基于MLP进行构建的，难以适配更先进的药品推荐模型架构。为了解决这一问题，本文提出了预训练-提示微调的多中心药品推荐新范式。

2. 问题定义

多中心的电子医疗记录(EHR)同时包含多个医院的数据，表示为。每个治疗样本又是由其诊断集合、治疗集合和用药集合构成的，记为。由此，我们可以将多中心药品推荐任务形式化为：在全部医院的EHR数据上训练一个模型，给定任何某个医院中一个病人的诊断和治疗情况和，该模型可以为为病人推荐药品组合。

3. 方法

3.1 方法概述

所提出的TEMPT药品推荐模型主要包含两个训练阶段：预训练和提示微调。为了能够推荐合适的药品组合，我们使用嵌入矩阵来获得诊断和治疗的表征。然后，使用两层transformer来编码诊断和治疗集合，名为medical encoder。为了学习多中心的通用医疗知识，我们先在全部医院的数据上对模型进行预训练。然后，为了适配到不同医院的开药分布，我们使用每个医院的数据对TEMPT进行了提示微调。

3.2 输入表征和medical encoder

我们为诊断和治疗分配了两个可学习的embedding矩阵，用于将ID集合转换为embedding集合，作为模型的输入。输入集合可以表示为和。

Medical Encoder主要是由两层transformer构成的，用于编码诊断向量集合和治疗向量集合。值得注意的是，为了更好地学习公有的医疗知识，我们所设计的Medical Encoder使用相同的参数对诊断集合和治疗集合进行编码。类似于Bert的做法，我们在诊断和治疗集合的最前面插入"[CLS]"token，最后使用该token位置对应的输出表征作为这个集合的整体表征，本文记为和。

3.3 预训练阶段

在预训练阶段，我们使用全部医院的数据来使模型学习一般性的医疗知识。但是，由于数据稀疏问题，这种学习通常是困难的。并且，诊断和治疗之间的关系之前很少被考虑到。因此，本文设计了两个自监督任务来进行预训练，分别是是Mask Prediction和Contrastive任务，如图3所示。

3.3.1 Mask Prediction Task

诊断集合和治疗集合中每个元素之间的共现关系通常反映了重要的医疗知识，因此捕获这种共现关系是非常重要的。针对此，我们设计了mask预测任务。具体地，我们随机mask掉诊断集合和治疗集合中的一部分元素，即使用mask表征进行替代。由此得到两个集合表征分别为和。

然后我们使用这个表征经过一个MLP来预测被mask掉的诊断和治疗。我们这里以诊断预测为例，经过MLP的预测如下公式：

然后，我们根据该预测可以构建mask预测的损失函数如下。针对治疗的损失函数类似，所以最终该任务的总体损失函数为两者相加。

3.3.2 Contrastive Task

通常，治疗也是针对诊断疾病的，因此治疗集合和诊断集合之间的关联关系也包含了非常重要的医疗知识。故，我们构建了一个对比学习任务，来捕获这两个集合之间的关系。具体地，我们同样利用上一个任务中得到的诊断和治疗表征，经过MLP的转换，获取用于对比学习的表征。

然后我们设计了基于batch的对比学习方式。即一个batch内，对应位置的诊断和治疗表征为一对正样本，batch内其他对应位置均为负样本。然后，我们就可以得到对比学习任务的损失函数如下：

3.4 提示微调阶段

在微调阶段，我们以药品推荐为目标进行训练。具体地，使用诊断和治疗集合表征的拼接，经过一个MLP得到最终的推荐结果。

实际上，直接使用预训练好的模型在每个医院的数据上进行微调也是一个可行方案。但是，全量参数微调会导致灾难性遗忘问题，导致预训练阶段学习到的通用医疗知识丢失。因此，我们提出了针对每个医院的提示微调方法。具体地，我们为每个医院设计一个可微调的提示向量。这个提示向量分别插入到诊断和治疗向量集合的最前面，由此得到了新的的输入向量集合：

加入了提示向量的输入集合得到的新的诊断和治疗表征记为和。由该表征可以进行最后的药品推荐。

值得注意的是，提示微调的时候只微调最后的MLP层和提示向量的参数，而medical encoder的参数都是固定不动的，从而更好地保留了学习到的通用医疗知识。

4. 实验

4.1 实验设置

我们在eICU数据集上进行了广泛的实验，因为该数据集是当前仅有的公开的可用于多中心药品推荐的数据集。eICU处理后的的统计信息如下：

4.2 对比方法

我们一共对比了四组方法：（1）传统药品推荐模型：Leap, GAMENet, G-Bert, COGNet；（2）基于大语言模型的药品推荐方法：GPT-4, TALLRec；（3）多域推荐模型：MMOE, PLE, STAR；（4）TEMPT的变体方法：Single-Train, Full-Train和Finetune。

4.3 整体性能比较

从总体效果来看，传统药品推荐模型的效果都远差于MDR和TEMPT。原因在于现在的模型都没有考虑到多中心的场景，大部分医院的数据量不足以支持充分训练药品推荐模型。基于大模型的方法同样表现较差，因为大模型更擅长零样本和少样本的设定，并且它们并不具有充分的医疗领域的知识。虽然MDR的方法表现较好，但是依旧差于TEMPT，因为MMOE和PLE难以适配医院数量特别多的情况，而STAR难以应用transformer这类能力更强的模型架构。TEMPT(finetune)的效果验证了预训练-微调的方式非常适合于多中心药品推荐的任务，但是其面临灾难性遗忘的问题，导致不如我们所提出的提示微调的方法。

4.4 消融实验

TEMPT-DE表示我们使用不同参数的medical encoder编码诊断和治疗。TEMPT-SM表示在预训练阶段对两个预训练任务使用相同的MLP。这两个变体效果均有下降，说明了我们设计的合理性。w/o MP和w/o CL表示分别去掉两个预训练任务，结果说明了每个预训练任务对于学习通用医疗知识都是有贡献的。

4.5 效率实验

FT, ST, TF和TE分别表示Full-Train，Single-Train，全量参数微调和提示微调。我们发现全量训练的耗时最长，但是因为其用一个模型服务全部医院，所以存储消耗是最小的。ST和TF因为都要对全部参数进行微调，所以训练时间依然大于提示微调。而存储方面，因为提示微调只需要存储不同医院的提示向量和MLP参数，因此其消耗虽然大于FT，但显著少于ST和TF。