CONTEXT CLUES: EVALUATING LONG CONTEXT MODELS FOR CLINICAL PREDICTION TASKS ON EHRS

該文章于2025年發表在ICLR（CCF A），早在2024年12月發布在arxiv。

文章地址：[Context Clues: Evaluating Long Context Models for Clinical Prediction Tasks on EHR Data](https://proceedings.iclr.cc/paper_files/paper/2025/hash/91a5bb5e5939cb075f5f2464d7b8bbf0-Abstract-Conference.html)

arXiv：[2412.16178] Context Clues: Evaluating Long Context Models for Clinical Prediction Tasks on EHRs

代碼倉庫：som-shahlab/hf_ehr: Training HuggingFace models on EHR data

openreview打分： 6 10 6 8 5

一、概述

在醫療領域大模型的評估和嘗試。首先，探索了更長的上下文模型可以提高預測性能（基于Mamba模型）。其次，在先前在模型上未充分探索的三個EHR屬性進行進一步分析（1.“復制轉發”診斷的流行，在EHR序列內有大量重復。2.EHR事件之間的不規則時間間隔，兩個事件之間可能相差1天、一年等。3.疾病的復雜性隨時間自然增加）。此外，用結構化EHR構建模型的詞匯表，訓練EHR序列大模型。

二、研究的核心問題及背景

研究的核心問題：對醫療領域大模型在三方面的評估和嘗試（長上下文、EHR屬性的影響、EHR建模）

背景：

1.用于EHR的基礎模型

基礎模型（FM）是通過無監督學習在大量未標記數據上訓練的大規模深度學習模型，在本文中，只考慮結構化EHR數據（忽略筆記、圖像等）。現如今的大多數基于EHR數據的序列建模框架使用基于Transformer的架構，如BERT或GPT，在數百萬EHR數據上用因果語言建模或掩碼語言建模進行預訓練。

2.長上下文的基礎模型

更長的上下文已經顯示出通過使模型能夠推理更多的信息來積極影響基礎模型的性能，token級別的困惑度隨著上下文長度增加而減少，說明更長的序列提高模型的理解能力。從理論上講，更多了解患者病史也應該有助于更好的臨床決策，成本也對應提高。為此，已經有提出針對較長上下文而優化的次二次架構，利用線性或對數線性的方法替換的Transformer中的注意力機制，但這些架構很大程度上仍未在EHR數據上測試。

三、現有方法面臨的挑戰

1.考慮到醫院的計算資源有限，現如今基于EHR數據的模型通常上下文長度不高，例如512。但是單個病人的EHR長度可能達到10K的token，極大的限制了這些基礎模型。

2.這些基礎模型也沒考慮最先進的次二次架構，減少成本。

3.與自然語言相比，EHR數據表現出特定類型的標記重復和噪聲，使較長上下文的預期好處復雜化。對此沒有更深入的分析

四、針對挑戰，解決思路

1.針對挑戰1、2，提出利用次二次架構的Mamba模型，并且進一步將token數擴展到16384，來探索上下文長度對EHR數據建模的影響。

2.針對挑戰3，本文探索了三個未充分探索的EHR屬性的定量分析，包括①.“復制轉發”診斷的流行，在EHR序列內有大量重復。②.EHR事件之間的不規則時間間隔，兩個事件之間可能相差1天、一年等。③.疾病的復雜性隨時間自然增加

五、模型框架與具體實現

問題定義

數據集由n個患者的EHR數據構成，每個患者的EHR數據包含一系列的事件三元組（事件，醫療實體，數值/分類值），作者稱為patient's timeline。

Tokenization（分詞）

將每一組時間三元組映射到一個token序列上，本文的詞匯表用的是先前工作EHRSHOT（一個few-shot基準，用于評測模型在EHR各種任務上的表現，2023）提供的，這樣可以和之前SOTA模型公平比較。具體來說，患者每一個臨床事件都有一個對應的編碼c來表示，在數據集中至少出現一次的所有唯一代碼c都被分配了一個唯一的token。與分類值相關的所有代碼c都被分配了一個唯一token。與數值相關的代碼c，在數據集中獲得的值范圍內，每十分位數分配一個唯一token。詞匯表共記錄了39811個token。

架構

本文評估了四種模型，GPT、Llama、Mamba和Hyena。

對于不同模型的上下文長度不同，在預訓練時做了截斷，預訓練使用的是交叉熵損失，以下一個token預測為目標。

評估

使用的是EHRSHOT基準，由15種臨床預測任務組成，例如預測ICU轉移、30天內再入院，預測實驗室檢查結果某項指標（血小板減少、貧血）是否異常，預測未來一年內是否會得高血壓等新診斷。主要評估指標是AUROC，此外還有Brier（預測與實際結果之間的均方誤差）。本文在零樣本、少樣本和全部 數據設置下評估模型，為了與原始EHRSHOT一致，只是在基礎模型之上訓練一個邏輯回歸頭。

EHR特殊屬性

對于之前劃分的三種EHR特性，首先，應用于EHR-OMOP驗證集，衡量真實世界EHR大規模語料在多大程度上表現出這些特性。其次，將其中兩種屬性重復性和不規則形應用于EHRSHOT數據集，以根據每個病人表現出這些特征的程度進行分層，這種分層能夠使我們評估模型性能如何隨著這些特征水平的不同而變化。

六、實驗

數據集：EHR-OMOP，從一個學術醫療中心按照OMOP格式化后的數據集。

EHRSHOT，一個few-shot學習基準數據集

1.在14個臨床預測任務上評估每個模型

2.分析了三個EHR特有屬性對模型性能的影響

3.“復制轉發”會影響模型性能

EHR數據表現出比自然語言更高的重復程度

將患者EHR重復率由低到高分為四個階段

4.不規則的token間時間間隔更難建模

EHR-OMOP數據集中的關于不規則性的時間數據統計

三種方法；每個患者的事件間隔時間的平均值，標準差和四分位數范圍

表2將這一分析擴展到在這項工作中訓練的EHR FM。雖然模型性能仍然隨著不規則性的增加而下降，但在所有四分位數中，Mamba和Llama的較長上下文版本始終優于其較短的對應版本。

5.更長的背景更好地模擬疾病進展效應

患者時間軸中較晚的token更難預測（更高的困惑度）

與較短的上下文相比，較長的上下文版本的Mamba和Llama在所有標記位置上始終實現較低的困惑度，并且在較晚的標記處差距擴大。這表明更完整的患者時間軸視圖有助于處理由于老化而增加的標記復雜度。相比之下，Hyena的較長上下文模型表現較差，復制了我們最初的EHRSHOT結果。對于GPT，結果好壞參半：較長上下文（2k和4k）在后面的token上實現了較低的困惑，但表現出顯著的尖峰。這似乎是由GPT使用絕對位置嵌入引起的-用旋轉位置嵌入（ROPE）取代它們。