ICML 2024 | 情境化的策略恢复：用自适应模仿学习来建模和解释医疗决策

DRUGAI

今天为大家介绍的是来自卡内基·梅隆大学的Eric P. Xing团队的一篇论文。可解释策略学习旨在从观察到的行为中估计可理解的决策策略；然而，现有模型在准确性和可解释性之间存在权衡，这限制了基于数据的人类决策过程的解释。从根本上说，现有方法之所以受到这一权衡的困扰，是因为它们将底层决策过程表示为一个通用策略，而实际上人类决策是动态的，可以在不同情境下发生显著变化。因此，作者开发了情境化策略恢复（CPR），将复杂决策过程的建模问题重新定义为一个多任务学习问题，每个情境代表一个独特的任务，可以通过多个简单的情境特定策略逐步构建复杂的决策策略。CPR将每个情境特定策略建模为一个线性映射，并随着新的观测数据的加入生成新的策略模型。作者提供了两种CPR框架的实现方式：一种侧重于精确的局部可解释性，另一种保留了完整的全局可解释性。作者通过模拟数据和实际数据进行了评估，在预测重症监护病房中的抗生素处方和预测阿尔茨海默症患者的MRI处方方面，达到了最先进的性能。通过这一改进，CPR弥合了可解释方法和黑箱方法之间的准确性差距，允许对情境特定决策模型进行高分辨率的探索和分析。

可解释策略学习旨在从一组展示行为的数据集中恢复底层的决策过程，并将这一过程表示为一个可解释的模型，该模型可以量化、审计并直观地理解。这种方法在医疗信息学社区中引起了广泛关注，被认为是一种有望通过检测偏见、解释次优结果以及量化区域和机构差异来提高护理标准的方法。经典的策略推断机器学习算法基于逆强化学习或模仿学习，并使用如递归神经网络之类的黑箱架构。这些方法已应用于多个医疗领域，最显著的是肿瘤学预后。然而，由于缺乏可解释性以及无法识别灾难性失败模式和泛化问题，黑箱方法在医学界受到了怀疑。

为了实现对可解释策略的需求，近年来出现了一系列透明的策略参数化方法用于模仿学习。这些方法包括递归决策树、视觉决策边界、高级编程语法或结果偏好。这些方法通常对临床医生来说更具可解释性，但其可解释性源于限制性的建模架构，从而牺牲了性能或施加了模糊的约束，使得实际应用充满挑战。主要的挑战在于，人类决策受多种因素影响，包括患者背景、病史、实验室测试等，真正的人类决策过程是复杂的。因此，将决策策略压缩为单一的观察-行动映射需要使用大型非参数模型（如神经网络），这些模型排除了直接可解释性，或者生成一个无法捕捉人类决策复杂性的模型。简而言之，模型必须既准确又可解释，才能有效支持临床决策。

模型方法

表 1

作者试图在每个时间步学习一个可解释的观测行为参数化，以理解代理在部分可观测的离线环境中如何采取行动，这扩展了经典模仿学习的目标（仅仅寻求复制展示的行为）。CPR结合了以前工作的优势（见表1），在每个时间步保持观测到行动的映射可解释，同时能够适应完整的过去观测，使得学习到的策略更接近展示的行为。虽然POETREE能够在时间上携带一个隐藏状态，但这个隐藏状态在每个时间步起到加性偏置的作用，而不是调整底层的模型参数θ，后者在整个时间内是静态的。

情境化建模

给定一个由目标y∈Y，观测x∈X和情境c∈C组成的数据集，相应的随机变量分别表示为Y，X和C，作者希望学习一个模型P(Y∣x, c, θ)基于x和c来预测y。由此，概率模型定义如下：

它由Al-Shedivat等人最初描述，并被Lengerich等人推广，使作者能够通过任何黑箱模型来建模P(θ|c)，同时保持P(Y|x, θ) 在由θ参数化的简单模型类中。由于数据集的复杂性、异质性和规模的增加，样本特定推断在许多应用领域引起了兴趣。情境化建模已在多个不同框架中用于估计情境特定参数θ，使用一个元模型P(θ|C)，将情境信息C与θ的变化联系起来。

情境化策略恢复

CPR基于最近在可解释、离线策略学习方面的发展。设数据集D={(xi1 ,ai1),…,(xiTi ,aiTi )}i=1N由N个治疗轨迹组成，其中每个患者i在Ti个离散时间步内观察到症状x∈X和医生行为a∈A。这些数据是由医生代理的未知策略生成的P(at∣x1,a1 ,…,xt−1, at−1, xt)，其中时间t的行为概率是代理当前状态的函数，该状态由当前和过去的患者症状以及过去的行为定义。

为了恢复一个既可处理又可解释的策略，CPR做出了一个实用假设，即在许多现实世界的环境中，最近的信息在决策时具有最高的重要性，而历史信息则为新信息的效果提供了背景。在医学领域，治疗历史对于将患者当前的疾病表现置于疾病进展和过去治疗尝试的背景中是有用的，但决策总是基于当前的表现。为了表示这种信息层次结构，CPR利用情境和历史特征生成情境特定的策略模型。

其中f是一个可解释的模型类，例如逻辑回归，由情境特定的θ参数化，θ是通过历史情境编码器g生成的。当前观测xt对行为概率at的影响通过简单的情境特定模型fθ直接解释。此外，g可以采用任何功能形式，而不排除f的可解释性。情境特定策略模型fθit为每个患者i在每个时间点t生成，研究先前的行为、先前的症状、患者协变量和治疗时间如何影响策略。CPR灵活地允许情境编码器g和观察到行为函数f自由选择，但它们必须是可微的，以便在适当的损失ℓ下进行联合优化。

在实验中，g由一个普通的RNN或LSTM参数化，f是一个逻辑函数，行为A:={0,1}是二元的，损失ℓ是二元交叉熵损失。CPR应用一个lasso正则化器到θ以学习鲁棒的策略参数。

全局可解释性

CPR结合了黑箱组件和可解释模型，能够提供高度准确的策略模型，这些模型在特定情境中是直接可解释的。尽管情境特定的策略是一种新颖、准确且高度适应的解释形式，但在某些情况下，全局可解释性仍然是必要的。为了理解每一个历史特征对每一个行为的确切影响，作者开发了CPR的第二个版本，称为CPR global。在CPR global中，作者利用CPR框架对全局可解释策略进行分段更新，该策略可以分解为所有观察到的特征和行为的线性组合。使用逻辑策略形式，

其中偏置项μt通过与β的线性组合进行每次新的观测更新。超参数α将这一更新与之前对偏置项的更新进行加权。这个形式以每次新的观测扩展为一个全局可解释的线性组合。因此，每次观测在情境特定的全局策略中都有一个确切的线性影响。

实验结果

表 2

作者将CPR应用于两个医疗数据集中的典型模仿学习任务：ADNI和MIMIC-III。这些数据集是部分可观察决策环境的典型例子，它们只能从展示行为中学习并且所学习的策略有可能改善临床操作。CPR显著优于其他可解释的基准模型，甚至在这两个电子健康记录（EHR）数据集上表现与完全黑箱模型相当（见表2）。低Brier得分表明CPR在实现最先进的AUROC和AUPRC的同时，校准良好。

情境化策略揭示了医疗决策过程中的异质性

图 1

为了了解决策函数在不同情境下如何变化，作者在模型空间中对它们进行比较。决策函数系数向量的UMAP嵌入（图1b）显示了三种截然不同的决策函数集群。最右边的集群包含了每个轨迹的初始模型参数θ0。由于在初次访问时没有情境可以区分代理的行为，这些参数对所有患者都是相同的，情境化模型恢复了总体估计值。随后，这种初始的同质性会产生决策策略的异质性。主要的异质性驱动因素是先前的抗生素处方——在先前接受过抗生素的患者更有可能继续接受抗生素，而在先前没有接受过抗生素的患者更有可能继续不接受抗生素。下方的集群包含大多数（99.8%）在前一状态t−1中接受抗生素的模型，而上方的集群包含大多数（99.3%）在t−1中没有接受抗生素的模型。只有情境化策略能够恢复这一强烈的分裂；忽略情境的全局策略无法识别这一异质性（图1a）。作者训练了两个分别以其各自情境为条件的模型，通过将观测限制在ICU的第二天，进一步消除更多的变异性。全局模型仅代表了总体的一小部分。通过对模型异质性的主要驱动因素（即患者在前一次访问中是否接受了抗生素）进行条件处理，并为每种情况训练单独的模型，得到的模型看起来像是分别在各集群的情境化模型上的平均值。

为了识别决策策略异质性的驱动因素，作者检查了决策函数参数随情境的变化（图1c）。参数中最显著的差异是截距值，对于接受抗生素治疗的患者组，这一值为正，对于没有接受抗生素治疗的患者组，这一值为负。对于已经接受抗生素治疗的患者，感染已经被检测到，医生的决策模型转向减轻抗生素治疗可能的副作用的风险，而不仅仅考虑益处。通过肌酐系数的变化支持了这一优先级的变化。高血清肌酐可能是肾功能受损的指示，这是抗生素可能的副作用；因此，高肌酐水平会降低继续抗生素治疗的可能性。最后，作者看到钾增加与开始抗生素治疗的决策相关。

情境化策略揭示异常患者

图 2

CPR在MIMIC抗生素处方数据集中识别出几名这样的异常患者（图2）。首先，年轻患者通常合并症较少，免疫系统更为强健，如果怀疑感染，医生可以更有信心地认为抗生素不会产生任何不良副作用。作者观察到，情境化策略恢复了这一情况，并且在开具抗生素处方时更准确地代表了20岁以下患者的治疗。其次，升高的肌酐是抗生素的罕见副作用，但却是肾功能不佳和可能的肾衰竭的可能指示。CPR识别到，肌酐升高的患者在初次处方后立即停止抗生素治疗，将其置于一个以缺乏抗生素处方为特征的治疗集群中，而这些患者通常不太可能处于这个集群中。由CPR生成的患者特定策略提供了治疗过程的新视角，能够轻松识别这些罕见且异常的情境特定策略参数和错误，从而揭示以前的工作中未能捕捉到的治疗决策中的细微差别。

ANDI MRI扫描

图 3

图3显示了在t=0时估计的策略在四个患者子群之间的差异。作者发现CPR生成的模型中存在有意义的异质性，其中年龄在CDRSB系数和总体截距中占主导地位，而性别在海马体积系数中占主导地位。此外，静态情境显著提高了CPR的行为匹配性能，AUROC达到0.763。

仿真结果

图 4

作者想知道CPR的显式线性策略表示是否是其性能和可解释性的关键，或者是否可以通过后期解释方法从黑箱策略模型中同样恢复准确和稳健的情境特定线性策略。为了测试这一点，作者模拟了一个异质的、行为依赖的马尔可夫决策过程（MDP），并评估了CPR与黑箱基准模型在恢复模拟参数方面的表现：真实的行为概率和情境特定的线性策略的真实系数（图4）。虽然CPR显式生成这些情境特定的线性系数，但黑箱模型隐式地将这些系数建模为特征梯度（即一阶泰勒展开中的线性系数）。类似于如LIME等流行的后期解释方法，作者利用RNNs的可微性Φ(xt,h)→at来恢复隐式的情境特定线性策略θ。在一个已知的异质且行为依赖的MDP上，CPR的显式策略表示相对于具有非结构化策略表示的黑箱模型不仅提高了对MDP参数的表示能力，还提高了整体性能。

编译 | 于洲

审稿 | 曾全晨

参考资料

Deuschel J, Ellington C N, Lengerich B J, et al. Contextualized Policy Recovery: Modeling and Interpreting Medical Decisions with Adaptive Imitation Learning[J]. arXiv preprint arXiv:2310.07918, 2023.