Stable Cox Regression用于分布偏移下的生存分析

生存分析旨在估计协变量对事件发生时间的影响。现有方法通常假设相似的训练和测试分布，然而，现实世界中变化的数据源会削弱其可靠性。这迫切需要生存分析方法利用不同队列中的稳定特征进行预测，而不是依赖于虚假相关性。为此，作者提出了基于稳定学习的 Cox 模型，用稳定学习理论来识别稳定变量。通过对模拟和现实世界的组学和临床数据的广泛评估，Stable Cox 不仅在不同的独立测试集中表现出强大的泛化能力，而且还根据已识别的生物标志物对患者的亚型进行分层。

来自：Stable Cox regression for survival analysis under distribution shifts, Nature Machine Intelligence, 2024

背景概述

生存分析是统计学的一个分支，用于评估协变量对某个事件发生时间的影响。在流行的生存分析方法中，Cox比例风险（PH，proportional hazards）模型在历史上最为突出，因为它可以灵活地处理删失数据，适应各种协变量，无需指定底层分布。虽然现有的生存分析方法在假设训练和测试数据具有相似的分布的情况下显示出有希望的结果，但当这一假设不成立时，就会出现挑战。训练和测试数据会从不同的中心收集，在医疗保健场景中，分布变化是不可避免的。例如，两类地区人群间通常存在异质性且治疗方案各异，因此预后标志物在两类地区之间的分布不可避免地会发生变化。

更具体地说，在肿瘤学中，预后标志物在患者的管理决策中起着关键作用，而识别这些标志物是临床研究的主要目标之一。然而，现有研究发现同一种生物标志物在不同的研究中显示出不同的预后价值。例如，在针对肝细胞癌 (HCC) 中国患者的研究中，通过免疫组织化学确定的肿瘤组织中上皮细胞粘附分子 (EPCAM) 表达的预后价值显示出不同的结果。一项研究确定 EPCAM 表达是良好预后的预测因子，而另一项研究发现高水平的 EPCAM 表达与不良预后有关。

作者对两个 HCC 转录组队列进行了单变量 Cox 回归分析。如图 1a（左）所示，与文献一致，两个队列中具有相同预后值的基因重叠有限。此外，一些基因甚至显示出完全相反的预后预测值。从数据角度来看，这些生物标志物与预后之间的不一致关系可能是由协变量的分布变化或基因与预后之间的真实功能关系引起的。

人们通常假设基因与特定癌症类型患者预后之间的真实功能关系是稳定的，并且不会在队列之间发生变化。相反，由于某些基因在不同人群中的表达水平不同，协变量分布很容易发生变化。在图 1a（右）中进一步可视化了这两个队列的协变量分布。显然，两个队列的协变量分布存在显著差异，表明存在分布偏移。分布偏移对生存分析构成了严峻挑战。分布偏移的主要挑战在于识别与不同队列的结果保持一致关系的稳定变量。当前方法可能会盲目地从训练集中存在的虚假相关性中学习错误模式。这种相关性是不稳定的，在将训练模型应用于新队列时会带来很大风险。

• 图1a：左图：维恩图显示了肝细胞癌基因组图谱 (TCGA-LIHC) 队列 (n = 351) 和 Roessler 等人 (n = 209) 的 HCC 转录组数据中与预后相关的基因的交集（两个队列之间有 11,512 个重叠基因）。满足以下标准的基因被视为预后相关基因：(1) HR（单变量 Cox 分析）大于 1 且对数秩 P 值（中位数分层）小于 0.05 为不利；(2) HR 小于 1 且对数秩 P 值小于 0.05 为有利。右图：这两个队列（cohort）对应的t-SNE图。
• 图1b：乳腺癌转录组数据包含 20,388 个基因表达对数谱作为 t-SNE 的输入。样本量如下：队列 1，n = 763；队列 2，n = 521；队列 3，n = 288；队列 4，n = 238。肺癌临床数据具有 51 个临床特征，这些特征经过最小-最大标准化。根据肿瘤的位置将样本分为两个亚群：中心性，n = 141；外周，n = 255。每个点代表一名患者。并根据其所属的队列（左）或肺部肿瘤的位置（右）着色。这些可视化突出显示了不同队列和亚群之间协变量分布的差异。

例如，如图 1b 所示，不同队列或亚群（协变量）分布不同，其中“批次效应”是造成队列异质性的主要原因。可以假设不同队列或亚群之间的变化主要是由部分协变量（即不稳定的协变量）引起的，这些不稳定的协变量可能与其他稳定协变量存在虚假相关性。如果我们在特定队列上训练模型，当前生存分析方法的相关性驱动性质大大增加了它捕捉到特定于该队列的虚假相关性的可能性。因此，将已确定的高风险因素作为生物标志物应用于未知人群具有严重后果（例如错误治疗分配）。鉴于高风险应用中不可接受的风险，要求模型必须识别能够适应分布变化的稳定特征。

增强 Cox PH 识别与结果变量最相关特征的能力的最常见做法是结合稀疏性范式，比如岭回归----线性回归添加惩罚的形式：
$L=\sum (y-\hat{y})+{\sum }_i{\beta }_i$等。这些正则化技巧本质上缺乏识别稳定变量和不稳定变量的能力。稳定学习是机器学习方法的一个分支，它将因果关系引入学习方法，旨在弥合因果推理中精确建模与机器学习黑箱之间的差距。受益于因果推理，稳定学习旨在在对协变量和结果之间的关系进行建模，并识别稳定的因果因子变量，而不是相关性。

StableCox 回归模型，旨在识别用于预测的稳定变量，从而确保基于这些选定变量在分布偏移下具有强大的泛化性。StableCox旨在消除协变量之间的虚假相关性，并专注于使用稳定变量进行预测。该模型流程如下：

• 独立性驱动的样本重加权和 Cox 回归。在独立性驱动的样本重加权阶段，使用一个模块来学习使协变量独立的患者样本权重。在随后的加权 Cox 回归阶段，使用这些学习到的权重对受试者进行重加权，目标是加权对数似然损失。这种损失在优化过程中有效地隔离了每个变量的影响。从理论上讲，Stable Cox 证明，在一些温和的假设下，即使模型设定错误，Stable Cox 模型也完全依赖稳定变量进行预测。这意味着，只要学习到的样本权重在所有协变量之间保持严格的相互独立，不稳定变量的系数将为零。

作者在模拟数据和两种关键的现实应用上验证了所提出方法的有效性。结果证明了所提出方法对未见过的测试队列或亚群的泛化能力。值得注意的是，从该方法得出的系数在下游任务中表现出了显著的稳定性和可解释性。学习到的系数可用于发现潜在的生物标志物并对具有显著不同生存风险的亚型进行分层。

Stable Cox通用框架

令$X=(X_1,X_2,...,X_p)\in {\mathbb{R}}^p$是$p$维特征，描述每个受试者（subject）的特征，这些特征是与生存时间相关的协变量，用于预测死亡或者疾病复发的时间，$T\in [0,\infty )$是一个值，表示死亡或者疾病复发的时间（失效时间）， δ ∈ { 0 , 1 } δ\in\left\{0,1\right\} δ∈{0,1}表示截尾标志，$\delta =1$表示$T$是完全观测的，$\delta =0$表示$T$本质没有观测完，但是被人为在此时刻截断了。假设基于随机变量$X,T,\delta$得到$n$个独立同分布的数据$\left\{x^{(i)},t^{(i)},{\delta }^{(i)}\right\}$组成训练分布$p^{tr}$，其中 $x^{(i)}$ 和 $t^{(i)}$ 分别表示subject $i$ 的特征（协变量）和失效时间。令 $p^{te}$ 表示未知的测试分布。

在涉及多个协变量的生存分析问题中，通常只有一小部分协变量会显著影响事件时间，而其余协变量可能代表噪声或与结果显示出虚假相关性，这些相关性在unseen的测试分布中是不稳定的。例如，在组学数据中，一些基因和 tumor 展示出因果关系，比如某个基因的高表达导致某些类型的癌症具有侵袭性，例如 ERBB2（Erb-B2 受体酪氨酸激酶 2，也称为 HER2）阳性乳腺癌往往更具侵袭性。然而，某些基因的表达（例如，乳糖酶持久性的基因）可能与人居住的地点高度相关，因此该基因与其预后的关系在不同城市之间是不稳定的。这样，基因就会因地点而产生虚假相关性。为了形式化这种情况，作者对协变量进行了结构性假设，将它们分为稳定变量 $S$ 和不稳定变量 $V$，其中失效时间 $T$ 仅取决于稳定变量 $S$。稳定变量是结果的真实预测因子，而不稳定变量则通过与稳定变量的相关性与结果相关联。

在迁移的场景中，通常假设$P(T|X)$不变，而$P(X)$可能会在训练集和测试集之间发生变化。例如，在生存分析中，一些基因在不同队列中持续表现出与不良预后相关的稳定趋势。如图1c所示，由于偏差，稳定协变量$S$和不稳定协变量$V$之间存在虚假相关性，从而导致$P(X)$变化，因此，意外的相关性会误导模型学习 $V$ 和 $T$ 之间的虚假相关性。这种相关性 $P(T|V)$ 在测试分布中是不稳定的，将导致泛化性下降。为了摆脱不稳定的相关性并捕捉 $S$ 和 $T$ 之间的稳定关系，Stable Cox提出学习一组样本权重来消除观测数据中协变量之间的相关性，然后在加权分布中优化 Cox 模型。

• 图1c：所提出的框架与传统生存分析方法的比较说明。Stable Cox旨在消除稳定变量 $S$ 和不稳定变量 $V$ 之间的虚假相关性，从而消除 $V$ 与生存结果 $T$（即 $P(T|V)$）之间的虚假相关性，并专注于跨环境（Envs）的稳定关系 $P(T|S)$。因此，Stable Cox 在多个独立测试队列和各种下游任务中比 Cox PH 有了显着的改进。

Stable Cox 回归模型由两个阶段组成。第一阶段，利用样本重加权模块来学习样本权重，以便 $X$ 在加权分布中具有统计独立性（图 2a）。在实施过程中，使用典型的独立性驱动算法，即样本重加权去相关算子 (SRDO)。先前的研究-SRDO提出通过估计训练分布 $P^{tr}$ 和特定加权分布 $\tilde{P}$ 的密度比来学习加权函数 $w(X)$。假设$\tilde{P}$通过重采样得到$\tilde{P}(X_1,...,X_p)={\prod }_{j=1}^p{P}^{tr}(X_j)$。因此，加权函数为：$$w(X)=\frac{\tilde{P}(X)}{P^{tr}(X)}=\frac{{\prod }_{j=1}^p{P}^{tr}(X_j)}{P^{tr}(X_1,X_2,...,X_p)}$$等式中的密度比可以通过类概率估计问题有效解决。因此，如果准确估计密度比，SRDO 可以保证协变量 $X$ 之间的统计独立性。

• 图2a：第 1 阶段，独立性驱动的样本重加权。样本重加权模块会消除协变量之间的虚假相关性。
• 图2b：第 2 阶段，加权 Cox 回归。该模块利用学习到的样本权重 $W$ 在对数似然损失中重新加权每个subject的事件。加权对数似然损失可以驱动模型利用稳定变量 $S$ 而不是不稳定变量 $V$ 进行预测。

在第二阶段，如图 2b 所示，Stable Cox通过学习到的权重重新加权 Cox PH 模型的对数似然损失，如下所示：$$\begin{gathered} L_w(\beta )=\sum _{i=1}^n{\delta }^{(i)}log{L}_w^{(i)}(\beta )=\sum _{i=1}^n{\delta }^{(i)}w(x^{(i)})log\frac{\lambda (t^{(i)};x^{(i)})}{{\sum }_{j:t^{(j)}\ge t^{(i)}}w(x^{(j)})\lambda (t^{(j)};x^{(j)})} \\ =\sum _{i=1}^n{\delta }^{(i)}w(x^{(i)})log\frac{{\lambda }_0(t^{(i)})exp({\beta }^T{x}^{(i)})}{{\sum }_{j:t^{(j)}\ge t^{(i)}}w(x^{(j)}){\lambda }_0(t^{(j)})exp({\beta }^T{x}^{(j)})} \\ =\sum _{i=1}^n{\delta }^{(i)}w(x^{(i)})log\frac{exp({\beta }^T{x}^{(i)})}{{\sum }_{j:t^{(j)}\ge t^{(i)}}w(x^{(j)})exp({\beta }^T{x}^{(j)})} \end{gathered}$$其中$\beta$是需要学习的协变量的系数。${\lambda }_0(t^{(i)})$表示时间$t^{(i)}$时的基线风险函数值。$L_{w=1}^{(i)}(\beta )$表示在时间$t^{(i)}$对于subject $i$而言观察到的事件发生的未加权似然。似然考虑了在时间 $t^{(i)}$ 之前事件尚未发生的任何subject $j$（包括subject $i$ 本身）的总和。在似然最大化过程中，模型对subject $i$ 在任何subject $j$ 之前发生事件的预测概率被优化。$L_w^{(i)}(\beta )$是每个subject似然的加权版本，其中每个subject发生 $w(X)$ 次。在温和的假设下，可以证明不稳定协变量$\hat{\beta }(V)$的估计系数将以高概率接近于零。

方法

引理

标记
令$X$表示特征，对于向量$x\in X$，定义$x^{\otimes 0}$表示标量1，$x^{\otimes 1}$表示$x$本身，$x^{\otimes 2}$表示$x{x}^T$。在概率期望的背景下，${\mathbb{E}}_Q[\cdot ]$和${\mathbb{E}}_Q[\cdot |\cdot ]$分别表示分布$Q$的期望和条件期望。为了简化符号，当提到训练分布 $P^{tr}$ 下的期望时，省略下标并使用$\mathbb{E}[\cdot ]$和$\mathbb{E}[\cdot |\cdot ]$表示。

加权函数
考虑集合$W$，定义为满足以下条件的加权函数集合：

对于每个$w\in W$，加权分布${\tilde{P}}_w$为：

为了简化标记，${\mathbb{E}}_w[\cdot ]$表示${\tilde{P}}_w$的期望。进一步，对于任意函数$f(X)$，有$\mathbb{E}[w(X)f(X)]={\mathbb{E}}_w[f(X)]$。

此外，独立性驱动的样本重加权算法专注于学习子集$W_{\perp }\subseteq W$，而不是完全的集合$W$，$W_{\perp }$上的加权函数确保特征$X$在相应的加权分布${\tilde{P}}_w$中相互独立。

Cox PH
在生存分析中，关键变量$T$和$\delta$通过以下过程生成：

其中，$T^{failure}$代表死亡或复发时间（事件已发生），$T^{censored}$表示截尾时间。因此，如果观测变量$T$在截尾时间之前发生，则可以得到死亡或复发时间。

在Cox PH模型（或者称为Cox regression）中，失效时间$T^{failure}$的风险函数$\lambda (u)$，测量在时间$u$发生的事件的瞬时速率，条件是直到时间$u$才发生事件。数学上表示为$\lambda (u;X)=li{m}_{h\to 0^{+}}{P}^{tr}(T^{failure}\le u+h|T^{failure}\ge u,X)/h$，对于以协变量$X$为特征的个体，标准Cox模型假定该个体的风险函数具有以下形式（方程 7）：

其中，${\lambda }_0(u)$表示基本风险函数，反映当协变量处于baseline水平时事件的时间相关风险。向量 $\beta$ 体现了效应参数。根据 PH 假设，协变量和风险之间存在乘积关系。例如，在系数恒定的场景中，药物治疗可能会在任何时间 $t$ 持续将受试者的风险降低两倍，而不管基线风险的变化如何。

Cox 建议使用对数偏似然函数估计参数 $\beta$，仅关注协变量的影响（方程 8）：

Stable Cox Regression

先前的研究表明，如果方程（7）中的假设准确地代表了真实的风险函数，那么当样本量接近无穷大时，方程（8）中的估计量将收敛到方程（7）中的实际$\beta$。然而，在实际情况中，测试数据可能来自与训练数据不同的中心或人群，从而导致分布偏移。

为了应对挑战，作者提出了Stable Cox Regression。这种方法侧重于利用在不同环境中表现出更高一致性的特征进行预测，而不是依赖于训练数据中发现的虚假相关性。具体来说，Stable Cox 回归将样本重加权模块与加权 Cox 回归模块相结合，以实现跨不同数据集的稳健性能。

样本重加权
该模块的主要目标是最小化协变量的依赖性。首先，取原始协变量矩阵 $M$，并通过随机逐列重采样生成列去相关版本 $\tilde{M}$。此过程会破坏 $X$ 中变量的联合分布，从而产生 $p$ 个独立的边际分布。因此，虽然原始特征分布是 $P^{tr}(X)$，但重采样后的分布变为：$$\tilde{P}(X)=P^{tr}(X_1)P^{tr}(X_2)\cdot \cdot \cdot P^{tr}(X_p)$$接下来，采用密度比估计技术构建了一个联合分布$P^{\prime }(X,Z)$，其中 Z ∈ { 0 , 1 } Z\in\left\{0,1\right\} Z∈{0,1}，$P^{\prime }(Z=0)=P^{\prime }(Z=1)=0.5$，$P^{\prime }(X|Z=1)=\tilde{P}(X)$，$P^{\prime }(X|Z=0)=P^{tr}(X)$，通过贝叶斯理论，密度比被形式化为：

通过训练MLP来区分来自原始矩阵$M$和去相关版本$\tilde{M}$的样本，可以估计出概率$P^{\prime }(Z=1|X)$和$P^{\prime }(Z=0|X)$。为了将样本权重标准化到单位均值，$w(x^{(i)})$ 通过将$w(x^{(i)})$除以平均值$\frac{1}{n}{\sum }_{i=1}^{n}w(x^{(i)})$来缩放。

加权Cox回归
在此模块中，先前获得的样本权重被应用到 Cox PH 损失函数中对每个受试者的事件进行重新加权。具体的，log部分似然$L_w(\beta )$和估计量${\hat{\beta }}_w$被定义为：

• Stable learning小结，我们可以发现，稳定学习主要体现在损失函数层面，这和StableNet是一个思想，用稳定学习训练出适应分布偏移场景的参数。
• 其次，稳定学习的核心就是样本重加权，然后把加权后的样本用于训练。

SRDO估计样本权重

SRDO（Sample Reweighting by Density Ratio Optimization）是一种通过密度比率估计来重新加权样本的方法，主要用于处理分布不匹配的问题。它的核心思想是通过估计源分布 p(x) 和目标分布 q(x) 之间的密度比率，从而为样本分配新的权重，使得重新加权后的样本分布更接近目标分布。

首先定义重采样函数 column_wise_resampling：

def column_wise_resampling(x, p_s, decorrelation_type="global", random_state = 0, category_groups= None,  **options):
    rng = np.random.RandomState(random_state)
    n, p = x.shape

    # 确定需要重采样的特征, 如果 options 中指定了 sensitive_variables，则使用指定的特征索引
    if 'sensitive_variables' in options:
        sensitive_variables = options['sensitive_variables']
    else: # 否则，默认对所有特征（0 到 p-1）进行重采样
        sensitive_variables = [i for i in range(p)] 
    # 创建一个与输入数据 x 形状相同的零矩阵 x_decorrelation，用于存储重采样后的数据
    x_decorrelation = np.zeros([n, p])
    if decorrelation_type == "global":
        # 对每个特征进行遍历
        for i in sensitive_variables:
            rand_idx = rng.permutation(n) # 生成一个随机排列 rand_idx，用于打乱样本顺序
            x_decorrelation[:, i] = x[rand_idx, i] # 将打乱后的特征值赋值给 x_decorrelation 的第 i 列。
    elif decorrelation_type == "group":
        rand_idx = rng.permutation(n)
        x_decorrelation[:, :p_s] = x[rand_idx, :p_s]
        for i in range(p_s, p):
            rand_idx = rng.permutation(n)
            x_decorrelation[:, i] = x[rand_idx, i]
    elif decorrelation_type == "category":
        print("len", len(category_groups))
        pre_j = 1
        for i, j in enumerate(category_groups):
            if j ==pre_j:
                x_decorrelation[:, i] = x[:, i]
                pre_j = j
                continue
            count = category_groups.count(j)
            if count == 1:
                rand_idx = rng.permutation(n)
                x_decorrelation[:, i] = x[rand_idx, i]
            else:
                print(i, count)
                rand_idx = rng.permutation(n)
                x_decorrelation[:, i:i+count] = x[rand_idx, i:i+count]
            pre_j = j 

    else:
        assert False
    return x_decorrelation

重采样函数中，输入矩阵x，形状(n,p)，n是样本数，p是特征数，可以看到，decorrelation_type=“global"其实就是改变了样本顺序，其实更应该把"重采样"叫"置换”。其次注意到，每个特征都被进行了置换，看起来数据完全乱套了，因为特征1和特征2之间由于样本置换顺序不同，已经没有样本配对关系。但是现在得到的新分布 q(x) 有如下意义：

• 特征之间相关性被破坏
• 置换操作不会改变每个特征的边际分布，只会打破特征之间的联合分布

所以，现在可以在 p(x) 上重采样去近似 q(x)，从而让 p(x) 也具有特征不相关的属性，定义SRDO：

def SRDO(x, p_s, decorrelation_type="global", solver = 'adam', hidden_layer_sizes = (100, 5), category_groups = None, max_iter = 500, random_state = 3):
    n, p = x.shape
    x_decorrelation = column_wise_resampling(x, p_s, decorrelation_type, category_groups = category_groups, random_state = random_state)
    P = pd.DataFrame(x)
    Q = pd.DataFrame(x_decorrelation)
    corr_matrix = np.corrcoef(x, rowvar=False)
    abs_corr_matrix = np.abs(corr_matrix)
    sum_abs_corr = np.sum(abs_corr_matrix) - np.sum(np.diag(abs_corr_matrix))
    P['src'] = 1 # 1 means source distribution
    Q['src'] = 0 # 0 means target distribution
    Z = pd.concat([P, Q], ignore_index=True, axis=0)
    labels = Z['src'].values
    Z = Z.drop('src', axis=1).values
    P, Q = P.values, Q.values
    # train a multi-layer perceptron to classify the source and target distribution
    #clf = LogisticRegression(random_state=0, C=0.1)
    clf = MLPClassifier(solver=solver, hidden_layer_sizes=hidden_layer_sizes, max_iter=max_iter, random_state=random_state)
    clf.fit(Z, labels)
    proba = clf.predict_proba(Z)[:len(P), 1]
    weights = (1./proba) - 1. # calculate sample weights by density ratio
    weights /= np.sum(weights) # normalize the weights to get average 1
    weights = np.reshape(weights, [n,1])
    return weights