Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

该论文介绍了 CausalLongPFN，这是一种在合成时间因果模型上预训练的冻结先验拟合网络，它无需特定领域的训练即可预测反事实纵向结果，在癌症、HIV 和 MIMIC-III 等基准测试中展示了具有竞争力的性能，同时为传统的特定队列建模提供了一种更具成本效益的替代方案。

要点

想象一下，你是一名医生，正试图预测一名患者在未来几周内对特定治疗方案的反应。问题在于，患者是非常复杂的：他们的健康状况每天都在变化，过去的治疗会影响他们当前的状态，而且其他因素（如饮食或压力）也会干扰结果。通常，为了做出这些预测，你必须为遇到的每一个新的患者群体构建一个全新的、高度专业化的计算机模型。这就像是每当你搬到一个新社区时，都要雇佣一位新的建筑师来设计一座房子。这既缓慢又昂贵，而且需要大量的数据。

这篇论文介绍了一个名为 CAUSALLONGPFN 的新工具。你可以把它想象成一个 “通用医学直觉引擎”，它在见到真正的患者之前，就已经学习好了游戏规则。

以下是它的工作原理，通过简单的概念进行拆解：

1. “训练营”（合成预训练）

与其等待真实的患者到来，创作者们建立了一个庞大的虚拟“训练营”。在这个营地里，他们模拟了数百万个虚假患者，涵盖了数百万种不同的体型、疾病和治疗反应。他们为这些虚假患者编写了复杂的行为逻辑：

• 有些人恢复缓慢；有些人则迅速恶化。
• 有些治疗效果立竿见影；有些则具有延迟效应。
• 有些患者会根据其隐藏的遗传基因做出不同的反应。

这个 AI 模型在整个训练营期间，一直在学习如何预测这些虚假患者的结果。它不仅仅是在死记硬背答案，而是在学习治疗、时间和生物学是如何相互作用的底层逻辑。

2. “冻结的专家”（无需重新训练）

这里有一个神奇的技巧：一旦模型完成了训练营的任务，他们就冻结了它。他们锁定了它的“大脑”。它无法学习任何新知识，也无法改变其内部设置。

当一批新的真实患者到来时（例如医院里的癌症患者），该模型并不会从零开始。它不需要重新训练。相反，它表现得像一个阅读病历的超级实习生：

• 支持轨迹（Support Trajectories）： 你向模型展示当前医院中一些患者的示例（即“支持数据”）。这些示例就像是给实习生看一些案例文件，以便他们了解这家医院的具体情况。
• 查询（The Query）： 你问道：“如果我们在接下来的 5 天内给这位特定患者使用治疗方案 A，会发生什么？”
• 答案： 模型会立即结合你在案例文件中展示的上下文信息，利用它在训练营中学到的知识，预测出结果。它完成这一切时，完全不需要进行任何“梯度下降”（这是描述通常重新训练过程的一个技术术语）。

3. “时空旅行模拟器”

该模型旨在处理**纵向（longitudinal）**数据，这意味着它理解时间。它不仅仅是猜测明天的结果，它还会逐步模拟未来。

• 它预测第 1 天。
• 它将该预测作为第 2 天的起点。
• 它重复这个过程，以观察在第 5 天、第 6 天或第 7 天会发生什么。

这就像是一个飞行模拟器，它不仅能展示起飞过程，还能根据飞行员的决策模拟整个飞行路径，即使飞行途中天气发生了变化。

4. 为什么这很重要（结果）

作者将这个“冻结的专家”与旧的方法（即为每个医院构建一个新模型）进行了对比测试。

• 测试： 他们使用了癌症、艾滋病毒（HIV）、血液稀释剂（华法林）以及真实的重症监护室（ICU）记录数据。
• 结果： 冻结模型表现得与那些为特定疾病专门训练的模型一样出色，甚至更好。
• 重大胜利： 它在真实的 ICU 数据上表现尤为出色，因为在这些数据中，人们无法进行“如果……会怎样”的实验（因为在模拟中对真实患者尝试不同治疗方案是不符合伦理的）。该模型仅根据观察到的数据，就能预测接下来可能发生的情况。

核心结论

该论文声称，你并不总是需要为每一个新的医疗数据集构建一个定制的新模型。相反，你可以通过在广泛的“假设场景”（合成数据）上训练一个庞大的模型，然后将其作为一个冻结的、即插即用的工具来使用。

这就像拥有一位大师级厨师，他已经在虚拟厨房里练习过各种类型的烹饪。当你带去一组新的食材（一组新的患者）时，他不需要重新学习如何烹饪；他只需观察你的食材，就能立即运用他那广博的、预先习得的直觉，知道如何准备这道菜。

重要提示： 作者特别强调，该工具是用于预测和研究，而非在临床中做出最终的生死决策。它有助于医生理解潜在的结果，但它仍然依赖于任何其他因果模型所遵循的标准医学假设。它是一个强大的研究工具，而不是医生判断力的替代品。

---

资源链接：
如果您希望查看实现细节或使用预训练模型，可以访问以下链接：

• 代码 (GitHub): https://github.com/Amirhossein-Zare/causal-long-pfn
• 预训练模型权重 (Hugging Face): https://huggingface.co/Amirhossein-Zare/causal-long-pfn

技术摘要

技术摘要：用于反事实结果预测的因果纵向先验拟合网络 (Causal Longitudinal Prior-Fitted Networks)

问题定义

本文旨在解决在纵向观测数据中，预测未来治疗序列下的潜在结果（potential outcomes）这一挑战。核心任务是估计历史条件潜在结果 $E[Y(\bar{a})_{t+\tau} | H_t]$，其中 $H_t$ 代表截至时间 $t$ 的观测到的历史记录，$\bar{a}$ 是计划的未来治疗序列。

该问题受以下三个主要因素的影响而变得复杂：

1. 时变混杂（Time-varying confounding）： 每个步骤的治疗分配取决于协变量，而这些协变量本身又是先前治疗的结果。
2. 异质性患者动态（Heterogeneous patient dynamics）： 个体轨迹表现出复杂的、非线性的状态演化和潜在的异质性。
3. 数据局限性： 观测队列通常规模较小，难以为每一个新领域或模拟器从头训练可靠的深度序列模型。

现有的纵向因果估计器（如 RMSN、CRN、G-Net、Causal Transformer）通常需要针对每个新队列进行单独的有监督训练过程，包括超参数选择和倾向评分建模。当需要重复进行特定领域的训练时，这种流程成本高昂且不切实际。

方法论：CAUSALLONGPFN

作者引入了 CAUSALLONGPFN，一种先验拟合上下文预测器（prior-fitted in-context predictor），旨在通过广泛的纵向结构因果模型（TSCMs）先验来摊销（amortize）纵向因果预测的任务。

1. 基于广泛先验的合成预训练

该模型完全在从广泛的 TSCM 先验中采样的合成样本上进行预训练。该先验旨在覆盖广泛的纵向因果动态类，而非仅仅复制单个模拟器。合成先验的关键特征包括：

• 因果时间图（Causal Temporal Graphs）： 具有稀疏的共时性和滞后依赖关系，以及无环的瞬时图。
• 非线性机制： 状态坐标遵循稀疏的非线性自回归更新，使用多样化的基本非线性函数（恒等、tanh、正弦、ReLU 等）和各种噪声分布。
• 动力学模态（Dynamical Motifs）： 叠加了结构化模态，如动作记忆（action-memory）、饱和（saturating）、稳态（homeostatic）、反馈控制（feedback-control）和光滑输出（smoothed-readout）通道，以捕捉诸如延迟效应和调节反馈等定性机制。
• 受混杂影响的行为策略： 治疗样本取自受潜在个体异质性 ($Z_i$) 影响的状态相关随机策略，从而产生时间变化的治疗-混杂反馈。
• 结果模型： 标量结果通过具有直接和累积治疗效应的自回归读出生成。

2. 架构

CAUSALLONGPFN 采用双编码器架构：

• 因果历史编码器（Causal History Encoder）： 一个轨迹级因果 Transformer（使用掩码自注意力机制），将纵向序列映射为历史表示，确保时间 $r$ 的表示仅依赖于截至该时刻可获得的信息。
• PFN 上下文编码器（PFN Context Encoder）： 一个基于 Transformer 的上下文编码器，执行上下文内适配（in-context adaptation）。它通过自注意力机制共同处理支持轨迹（视为无序集合）和查询标记（query token）。不对支持轨迹的顺序进行位置编码。
• 高斯混合预测头（Gaussian-Mixture Prediction Head）： 最终的查询表示参数化一个包含 5 个分量的正态混合分布，用于归一化的下一时刻结果，从而同时提供点预测和不确定性估计。

3. 上下文内预测与展开（In-Context Prediction and Rollout）

在测试阶段，模型是冻结的。它接收：

• 支持轨迹（Support trajectories）： 来自新领域/任务的示例。
• 查询历史（Query history）： 截至时间 $t_{obs}$ 的观测记录。
• 提议的未来治疗序列： 计划的干预措施。

模型直接返回预测分布，无需梯度更新、倾向评分建模或对抗平衡。对于多步预测 ($\tau > 1$)，模型执行自回归插件式展开（autoregressive plug-in rollout）：预测单步结果分布，将混合分布的均值插入查询历史中，并在指定的治疗序列下重复此过程。

核心贡献

1. 一种用于纵向因果预测的先验拟合模型： CAUSALLONGPFN 是首个针对计划纵向治疗序列下的历史条件潜在结果预测的 PFN 式模型。它作为一个冻结模型运行，无需测试时适配。
2. 一个关于纵向因果任务的合成先验： 作者设计了一个 TSCM 先验，生成的任务具有多样化的治疗-混杂反馈、潜在异质性、非线性动态、延迟/累积效应以及机制转变。
3. 用于纵向上下文推理的架构： 一种结合了因果 Transformer 历史编码器、PFN 上下文编码器和高斯混合头的创新双编码器。
4. 自回归反事实展开： 将学习到的单步预测器扩展到多步预测，通过确定性插件式展开实现。
5. 零样本评估（Zero-Shot Evaluation）： 通过在分支反事实基准测试和事实真实世界数据上，对单个冻结模型与领域训练的基准模型（MSM、RMSN、G-Net、CRN、Causal Transformer、G-Transformer）进行了全面的评估。

结果

模型在四个基准测试上进行了评估：癌症肿瘤生长、华法林（Warfarin）药代动力学/药效学、HIV 治疗动态以及 MIMIC-III ICU 轨迹。

• 领域均衡性能： CAUSALLONGPFN 实现了最佳的领域均衡单步归一化 RMSE (0.222)，略微超过了 MSM 和 RMSN。在五步预测中，其综合排名第三，仅次于 RMSN 和 G-Net，但优于 MSM、CRN 和基于 Transformer 的基准模型。
• 各领域表现：
  • MIMIC-III（事实数据）： 该模型在单步和五步预测中均排名第一，展示了在无需领域特定训练的情况下，向真实世界临床轨迹的强大迁移能力。
  • 反事实基准（癌症、HIV、华法林）： 模型保持了竞争力，在单步任务中排名第二或第三。然而，在长时程反事实任务（如 5 步癌症预测）中，专门的领域训练循环模型（RMSN、CRN）取得了更低的误差，这表明当存在足够的特定目标领域数据进行针对性拟合时，这类模型具有优势。
• 不确定性校准： 高斯混合头提供了有用的分布信息。校准效果因领域而异，华法林显示出最佳的校准度，而 MIMIC-III 由于更高的异质性显示出更宽的置信区间。

意义与主张

本文声称，广泛的合成因果预训练可以提供一个有用的冻结替代方案，特别是在以下情况：

• 重复进行特定领域的训练成本过高或不切实际。
• 需要快速适应新队列。
• 无法获得反事实监督（如在 MIMIC-III 等真实世界事实预测任务中）。

作者强调，CAUSALLONGPFN 并没有消除解释观测数据所需的标准因果假设（一致性、正值性、序列可交换性）。相反，它通过摊销（amortization）来简化估计问题。结果表明，一个足够广泛的合成先验可以捕捉治疗-响应任务中可重用的结构，使该模型成为一个强大的通用型上下文预测器。然而，论文也谦虚地指出，当存在充足的目标领域数据和验证信号时，特别是对于特定领域的长时程反事实预测，领域特定的训练仍然具有价值。

这项工作将自身定位为因果序列建模和假设生成的研究工具，而非独立的临床决策系统，并提醒不要在因果假设或先验支持不足时过度信任预测结果。

代码与模型可用性

该工作的实现代码已在 GitHub 上公开：https://github.com/Amirhossein-Zare/causal-long-pfn。预训练的模型权重可在 Hugging Face 上获取：https://huggingface.co/Amirhossein-Zare/causal-long-pfn。