[论文解读] Hierarchical BERT for Medical Document Understanding

Ning Zhang, Maciej Jankowski|arXiv (Cornell University)|Mar 11, 2022

Topic Modeling被引用 10

一句话总结

本文提出医学文档 BERT（MDBERT），一种分层 BERT 架构，通过自下而上地在标记、句子和文档层级处理文本，实现了长医学文档的有效、高效且通用的表征学习。MDBERT 在 ICD 编码分配任务上达到最先进性能（在 MIMIC-III 数据集上相对提升高达 20%），并将自注意力复杂度显著降低至原始 BERT 的 1/100 以下，同时支持可迁移的句子级与文档级嵌入，适用于多种下游 NLP 任务。

ABSTRACT

Medical document understanding has gained much attention recently. One representative task is the International Classification of Disease (ICD) diagnosis code assignment. Existing work adopts either RNN or CNN as the backbone network because the vanilla BERT cannot handle well long documents (>2000 to kens). One issue shared across all these approaches is that they are over specific to the ICD code assignment task, losing generality to give the whole document-level and sentence-level embedding. As a result, it is not straight-forward to direct them to other downstream NLU tasks. Motivated by these observations, we propose Medical Document BERT (MDBERT) for long medical document understanding tasks. MDBERT is not only effective in learning representations at different levels of semantics but efficient in encoding long documents by leveraging a bottom-up hierarchical architecture. Compared to vanilla BERT solutions: 1, MDBERT boosts the performance up to relatively 20% on the MIMIC-III dataset, making it comparable to current SOTA solutions; 2, it cuts the computational complexity on self-attention modules to less than 1/100. Other than the ICD code assignment, we conduct a variety of other NLU tasks on a large commercial dataset named as TrialTrove, to showcase MDBERT's strength in delivering different levels of semantics.

研究动机与目标

• 为解决标准 BERT 在处理长临床记录（>2000 个标记）时因 512 个标记序列长度限制而带来的局限性。
• 克服现有 ICD 编码分配模型的任务特异性问题，这些模型缺乏通用的文档和句子嵌入。
• 开发一个统一的端到端模型，同时学习多层次语义表征（标记、句子、文档），以提升在各种 NLP 任务中的可迁移性。
• 通过降低自注意力机制的计算复杂度，实现在长医学文档上的高效编码。
• 展示 MDBERT 学习到的嵌入在多种临床 NLP 任务中的泛化能力，包括 ICD 编码、患者分组分类和适应性标准化。

提出的方法

• MDBERT 采用自下而上的分层架构，将来自不同文档的句子重新分批，以实现在文档层级的联合学习。
• Transformer 模块在每个语义层级（标记、句子、文档）使用多头自注意力与相对位置编码处理表征。
• 通过分层平均池化聚合语义表征，支持在任意层级附加特定任务头的端到端训练。
• 将代码描述增强直接整合到训练过程中，无需额外的代码嵌入网络。
• 句子级嵌入来自首个平均池化层，而文档级嵌入则来自最终的池化表征。
• 在多个下游任务（如 ICD 编码、性别分类、患者分组）上应用多任务学习，以提升泛化能力和鲁棒性。

实验结果

研究问题

• RQ1 基于 BERT 的模型是否能有效处理超过标准 BERT 512 标记限制的长医学文档，且不造成性能下降？
• RQ2 统一模型是否能学习到可跨多种临床 NLP 任务迁移的通用文档级与句子级嵌入？
• RQ3 分层 BERT 架构是否能在保持或提升长文档理解任务性能的同时，降低计算复杂度？
• RQ4 MDBERT 的统一嵌入空间在支持下游任务（如 ICD 编码分配、患者分组分类和适应性标准化）方面的有效性如何？
• RQ5 与任务特定微调相比，MDBERT 的多任务训练策略是否能在低资源或复杂临床 NLP 任务中提升性能？

主要发现

• MDBERT 在 MIMIC-III 数据集上的 ICD 编码分配任务中实现了高达 20% 的相对性能提升，达到最先进水平。
• 该模型将自注意力模块的计算复杂度降低至原始 BERT 的 1/100 以下，实现了对长文档的高效处理。
• MDBERT 的文档级嵌入在多个下游任务中表现出良好的泛化能力，在 TrialTrove 数据集上对年龄组分类达到 98.4% AUC-ROC，对患者分组分类达到 99.8%。
• MDBERT 的句子级嵌入在适应性标准化任务中优于现有基线模型，在 SNOMED 匹配任务中取得 96.8% AUC 分数，超过 SBERT（95.3%）和 BioBERT（94.5%）。
• 多任务学习在 ICD10 编码预测任务中提升了性能，证明了在相关临床 NLP 任务中共享表征学习的优势。
• MDBERT 表现出过拟合迹象，性能在约 10 个周期后趋于平稳，尤其在小数据集上，提示未来工作需引入更强正则化或句子层级约束。