【ICLR 2024】A Flexible Generative Model for Heterogeneous Tabular EHR with Missing Modality

type

status

date

slug

summary

原文摘要

现实的合成电子健康记录（EHRs） 可以用于加速研究方法的开发，同时缓解与数据共享相关的隐私问题。然而，生成对抗网络（GANs） 的训练仍然具有挑战性，通常会导致模式崩溃等问题。尽管扩散模型（diffusion models） 在生成高质量的表格型 EHR 数据时取得了一定进展，但当面对异质性表格 EHR 数据中的缺失模态时，它们的性能会下降。例如，一些 EHR 数据仅包含静态测量值，而另一些仅包含时间序列测量值，或者两者兼有。

为了解决这一问题，我们提出 FLEXGEN-EHR——一种专为异质性表格 EHR 设计的多功能扩散模型，能够在集成学习框架下处理缺失模态。我们定义了一个最优传输（optimal transport）模块，以对齐并突出 EHR 数据的公共特征空间，从而更好地捕捉数据的异质性。实验结果表明，我们的模型在生成数据的真实性（fidelity）方面比当前最先进的合成 EHR 生成方法高出最多 3.10%，在实用性（utility）方面提升最多 7.16%。此外，我们还证明了该方法可以成功应用于隐私敏感的场景，即原始的患者级数据无法共享的情况下。

速览表

T 目标	EHR模态缺失影响了EHR数据合成精确度与实用性-》先去修复缺失的模态，再拿来生成
I 输入	部分模态缺失的数据
P 处理	跨模态统一空间映射 + 扩散模型生成
O 输出	修复好缺失部分的完善数据

P 问题	表型EHR数据的一些特点导致GAN模式崩溃。作者采用Diffusion的做法来生成EHR数据
C 条件	EHR（医学表格类型数据）、数据异质（离散、连续、文本）、长尾分布、多峰分布
D 难点	怎么学到各项指标的深层联系？对于人体医学数据，各项指标一定是有深层联系的（假如有一个人体模拟器，把生成出来的人体参数放进去，要能“拟人”的正常运行，而不是直接闪退。
L 水平	Accepted by ICLR 2024

1. 解决了什么问题？

医疗EHR数据具有很高的隐私性。从前普通的隐私化方法容易遭受攻击，使得数据再被重新标识——而用直接用生成式模型生成数据，那就完全不存在这个问题了——反正都是“假”的

出于隐私保护，医院数据将去掉部分列后公开给科研机构用作研究。比如一个100岁以上的人的就诊记录，就算匿名化，你也能找出来这个人到底是谁

用扩散模型生成这个缺失的模态

这篇文章做的是固定缺失某个模态

过往的工作

GAN系列

MedGAN (Choi et al., 2017), CorGAN(Torfi & Fox, 2020), and EHR-WGAN (Zhang et al., 2019)

mode collapse problem (Li et al., 2021c) and the alignment between heterogeneous data types

C-RNN-GAN 和 TimeGAN 采用 RNN 作为生成器和判别器，训练成本高，难以扩展到大规模 EHR 数据。
EHR-M-GAN 采用 Autoencoder + BLSTM 组合，但两者的表示空间不兼容，导致数据生成质量下降。

Diffusion系列

MedDiff 仅解决了连续数据，没有管离散数据
TabDDPM 能够解决EHR异质问题，但是他离散和连续分别用了两个encoder分别处理，这将导致模型对离散和连续之间联系的学习能力受损。

=》本工作 FlexGen-EHR 熟练地辨别并表示静态和时间特征之间的基本关系。

2. 核心思想&创新点

先分别用两个Encoder处理两种类型的数据，嵌入到两个空间过后再映射到同一个空间中去。在这个映射过程中实现模态的补全——Latent Space Alignment

在隐空间里补全后，再交给LDM后分割为两部分，两个Decoder分别恢复出对应的模态。

最优传输（Optimal Transport, OT）用于寻找两个点集（离散分布）之间的最小成本映射。在 EHR 生成任务中，OT 可用于对齐静态特征和时间序列特征的嵌入表示，使得它们共享一个统一的潜在空间。

具体来说，我们假设有两个经验分布（empirical distributions），分别对应：

静态特征的嵌入

时间序列特征的嵌入

它们的经验分布可以表示为：

其中：

p 和 q 是 非负向量（长度分别为 I 和 J），并且它们的元素总和为 1（即它们是概率分布）。

μ 和 ν 分别表示两个不同的概率分布（静态嵌入和时间序列嵌入）。

目标是找到一个概率耦合矩阵（probabilistic coupling），使得从 μ 到 ν 的传输成本最小。

✅ 基于潜在空间嵌入（Latent Space Embedding Models） 的方法，即在不同特征上训练的嵌入模型，往往能够展现相似的几何模式和行为（analogous geometric patterns and behaviors）。

因此，我们提出一个假设：

异构 EHR 数据的潜在嵌入空间可以通过线性变换（linear transformations）相互转换，即即使某个模态缺失，我们仍然可以通过学习一个线性变换，使得一个模态的嵌入能映射到另一个模态的嵌入，从而填补缺失的信息。

3. 核心方法

🔹 1. 背景：缺失模态对齐问题

传统方法假设所有模态（静态特征 & 时间序列特征）都存在，但在实际中，EHR 数据常存在模态缺失（如患者可能缺失某些时间序列数据）。为了处理这种情况，我们定义了一个最优传输（Optimal Transport, OT）问题，用于对齐静态特征嵌入和时间序列特征嵌入，即：

xi=(xiS,NA,yi)或xi=(NA,xiT,yi)

其中 NA 表示数据缺失。

🔹 2. 传统方法的局限性：线性变换的监督依赖

一种常见的方法是学习一个线性变换 AA 来对齐静态和时间序列特征嵌入：

min⁡A∈O(l)∥ZS−AZT∥F2A∈O(l)min∥ZS−AZT∥F2

其中：

ZT（时间序列嵌入矩阵） 和 ZS（静态特征嵌入矩阵） 需要样本数量相等。

O(l) 约束 A 为正交矩阵，确保变换的稳定性。

📌 局限性：这种方法需要已知的样本对应关系，因此无法处理缺失模态导致的不匹配问题。

🔹 3. 解决方案：Gromov-Wasserstein 传输（GW）对齐

为了解决缺失模态导致的样本不匹配问题，我们采用 Gromov-Wasserstein（GW）流形对齐 方法。

✅ 核心思想：

通过 GW 传输，利用 嵌入空间的内部结构 和 标签信息 进行对齐，而不依赖于最近邻方法推测缺失数据。

GW 通过 相似度矩阵（cosine similarity） 计算不同模态嵌入之间的匹配关系，并构造 4 阶损失张量 进行优化。

✅ 损失函数（衡量不同模态间的相似性误差）：

L(CikT,CjlS)=12(yiykCikT−yjylCjlS)2L(CikT,CjlS)=21(yiykCikT−yjylCjlS)2

✅ GW 传输优化目标（添加熵正则化加速求解）：

GW(CT,CS,p,q)=min⁡Γ∈Π(p,q)∑i,j,k,lLijklΓijΓkl−ϵH(Γ)GW(CT,CS,p,q)=Γ∈Π(p,q)mini,j,k,l∑LijklΓijΓkl−ϵH(Γ)

其中：

Γij表示匹配 ziT（时间序列嵌入）到 zjS（静态嵌入）的概率。

H(Γ) 是熵正则化项，提高计算效率。

🔹 4. 结果：基于 GW 传输的对齐与数据填补

通过求解 GW 传输问题，我们得到最优传输耦合矩阵 ΓΓ，它提供了一种 软匹配（soft matching） 机制，即：

提供任意两个嵌入是否匹配的概率，从而允许对缺失模态数据进行填补。

解决了 传统线性变换方法无法处理的缺失模态问题。

💡 填补缺失特征的方法：

如果静态特征缺失，可通过逆变换估计：

ziS=A:iziTΓi:−1ziS=A:iziTΓi:−1

如果时间序列特征缺失，则将 ziS 作为源、ziT 作为目标，重新求解问题。

📌 结论

方法	作用	局限性
线性变换方法	监督学习静态 & 时间序列特征之间的映射	需要已知的样本对应关系，无法处理缺失模态
Gromov-Wasserstein 传输	通过测度空间对齐填补缺失数据	计算复杂度较高，但可用熵正则化加速

✅ GW 传输方法的优势：

不需要监督样本对应关系，可处理缺失模态数据。

利用嵌入空间的几何结构，更合理地进行特征对齐。

可解释性强，提供不同模态嵌入的匹配概率。

4. 总结

个人觉得这篇文章和Diffusion Model没什么关联，加这个Diffusion Model的意义或许只是涨点，DM的生成效果比Flow、GAN等架构的生成模型的效果好。

实际上确实模态的修复并不是依靠的LDM，而是Gromov-Wasserstien Distance。消融实验没有给出使用不同DM的效果对比。

模态缺失的划分有问题，怎么可能全缺失离散或者全缺失连续？但是仔细一想，也只有这样划分吧