[论文解读] Three-level Hierarchical Transformer Networks for Long-sequence and Multiple Clinical Documents Classification

Yuqi Si, Kirk Roberts|arXiv (Cornell University)|Apr 17, 2021

Machine Learning in Healthcare被引用 5

一句话总结

本文提出了一种三级层次化的Transformer网络，通过从词到句子、句子到临床记录、记录到患者层面结果的分层处理，对长时序临床序列进行建模。通过扩展BERT的上下文长度并引入时间感知的自适应分割机制，该模型在MIMIC-III数据集上的院内死亡率和表型预测任务中实现了最先进性能，AUC最高提升21%，PRC提升51%，F1提升46%。

ABSTRACT

We present a Three-level Hierarchical Transformer Network (3-level-HTN) for modeling long-term dependencies across clinical notes for the purpose of patient-level prediction. The network is equipped with three levels of Transformer-based encoders to learn progressively from words to sentences, sentences to notes, and finally notes to patients. The first level from word to sentence directly applies a pre-trained BERT model as a fully trainable component. While the second and third levels both implement a stack of transformer-based encoders, before the final patient representation is fed into a classification layer for clinical predictions. Compared to conventional BERT models, our model increases the maximum input length from 512 tokens to much longer sequences that are appropriate for modeling large numbers of clinical notes. We empirically examine different hyper-parameters to identify an optimal trade-off given computational resource limits. Our experiment results on the MIMIC-III dataset for different prediction tasks demonstrate that the proposed Hierarchical Transformer Network outperforms previous state-of-the-art models, including but not limited to BigBird.

研究动机与目标

• 解决BERT输入长度限制在512个token时对纵向临床记录建模的局限性。
• 捕捉每位患者多份临床文档之间的层次化与时间依赖关系。
• 提升对院内死亡率和表型分类等临床结局的患者层面预测准确性。
• 通过超参数调优与分布式训练，在计算资源约束下优化模型性能。
• 证明使用完整纵向记录序列而非仅依赖孤立记录（如出院小结）的必要性。

提出的方法

• 在词到句子层级使用预训练BERT模型，提取上下文表征。
• 引入两个额外的基于Transformer的编码器堆栈：一个用于句子到记录的聚合，另一个用于记录到患者分类。
• 引入时间感知的自适应分割机制，根据时间接近度对临床记录进行分组，以保留纵向结构。
• 使用mean_max池化方法对句子和记录层级的表征进行聚合，增强层次化特征学习。
• 在所有三个层级应用位置编码，以保持序列顺序的感知能力。
• 采用分布式训练以突破GPU显存限制，实现对超过BERT标准限制的更长输入序列的处理。

实验结果

研究问题

• RQ1 层次化Transformer架构能否有效建模跨越多个患者记录的长时序临床序列？
• RQ2 通过自适应分割引入时间结构后，与将记录视为独立样本相比，临床预测性能是否显著提升？
• RQ3 在临床序列建模中，模型规模、输入长度与计算成本之间的最优权衡是什么？
• RQ4 与仅使用出院小结相比，使用所有可用临床记录是否能显著提升预测性能？
• RQ5 哪些架构组件（如池化策略、编码器层数量或位置编码）对模型性能影响最大？

主要发现

• 所提出的层次化Transformer网络在院内死亡率预测任务中，AUC相比先前SOTA模型（如BigBird）最高提升21%。
• 在表型预测任务中，PRC提升51%，F1分数提升46%，表明在多个指标上均有显著增益。
• 使用所有可用临床记录（而非仅出院小结）可显著提升性能，当仅使用单个记录时AUC下降至约0.85。
• 表现最佳的模型在句子和文档层级均使用2个编码器层，表明更大的模型并不总能带来更好结果，因优化瓶颈存在。
• 自适应分割至关重要——若移除该机制，AUC降至0.8558，PRC降至0.7887，凸显其在保留时间关系方面的作用。
• mean_max池化优于其他池化策略，且位置编码在保持各层级序列顺序感知方面至关重要。