TocBERT: Medical Document Structure Extraction Using Bidirectional Transformers

Abstract

Text segmentation holds paramount importance in the field of Natural Language Processing (NLP). It plays an important role in several NLP downstream tasks like information retrieval and document summarization. In this work, we propose a new solution, namely TocBERT, for segmenting texts using bidirectional transformers. TocBERT represents a supervised solution trained on the detection of titles and sub-titles from their semantic representations. This task was formulated as a named entity recognition (NER) problem. The solution has been applied on a medical text segmentation use-case where the Bio-ClinicalBERT model is fine-tuned to segment discharge summaries of the MIMIC-III dataset. The performance of TocBERT has been evaluated on a human-labeled ground truth corpus of $250$ notes. It achieved an F1-score of $84.6\%$ when evaluated on a linear text segmentation problem and $72.8\%$ on a hierarchical text segmentation problem. It outperformed a carefully designed rule-based solution, particularly in distinguishing titles from subtitles.

I Introduction

Text segmentation is the task of partitioning a document into topically coherent segments [1, 2, 3, 4]. It represents an important pre-processing step in several Natural Language Processing (NLP) downstream tasks inducing information retrieval and summarization[1, 2, 3, 4, 5, 6]. Two types of text segmentation can be distinguished: 1) linear segmentation where a document is divided into sequential contiguous segments, and 2) hierarchical segmentation where higher level segments are further segmented into smaller sub-segments [3].

Existing solutions to text segmentation are applied either to unformatted, i.e. plain-text, documents like [3, 4] or to formatted ones like [7, 8, 9]. In the first family, recent text segmentation methods are based on the semantic representation of words and/or of sentences. For instance, word embeddings and sentence embeddings are used to detect topic changes in a sequence of sentences in [3, 4]. In the second family, document layout, orthographic features, geometric features and stylistic properties, e.g. font properties and line spaces, are exploited to extract the document structure. Some of the aforementioned features can be extracted using the Document Image Analysis (DIA) [10, 9] while others can be extracted using regex matches of predefined patterns [8]. In both families, titles and subtitles play an important role in text segmentation. For example, in semantic-aspects-based methods, supervised solutions like [3] and [4] are trained on a large corpus named WIKI-727K [4] which is automatically labeled such that topic-change cutoff points are completely dependent on titles. Similarly, visual-aspects-based methods like [9] are mainly based on title detection.

Figure 1: An extract from a discharge summary from the MIMIC-III database annotated with titles (yellow) and subtitles (blue)

In this article, the focus is set on the hierarchical segmentation of unformatted medical documents by detecting titles and subtitles using their semantic vector representations. Particularly, the considered use-case is the segmentation of the discharge summaries of the MIMIC-III database [11] which represents an unformatted free-text corpus. Figure 1 shows an extract of a discharge summary where we highlighted titles in yellow and subtitles in blue. The motivation underlying the aforementioned segmentation task is two-fold:

• To clean the MIMIC-III discharge summaries corpus such that it can be efficiently used to domain-adapt pretrained language models. For example, noisy sections like ”Discharge medications” or ”Admission labs” can be detected and removed since their content doesn’t represent a fluent language and can harm the domain adaption task.

• To enable the extraction of specialized sub-corpora and facilitate the information retrieval tasks. For instance, sub-corpora about radiology reports or cardiac diseases reports can be extracted and used to train relevant information extraction systems.

While the aforementioned advantages have been discussed on the considered use-case, the proposed method is general and can be extended to solve many similar tasks as will be shown in Section IV.

The rest of the paper is organized as follows: we start by exploring related works in Section II. Training and test corpora are described in Section III. The proposed method is presented in Section IV. Experimental results are discussed in Section V while Section VI concludes the paper and discuses future work.

II Related work

Figure 2: Two-level neural networks for text segmentation: the first-level network builds sentence embeddings from token embeddings while the second-level network labels the sequence of sentence representations as: topic-change $(1)$ or no-topic-change $(0)$

One remarkable family of solutions to topical text segmentation was proposed in [4, 3]. As depicted in Figure 2, these methods are based on two-level neural networks that divide the segmentation problem into two sub-problems: 1) starting from token embeddings, build sentence embeddings, and 2) based on the sentence representations, annotate the sentences’ sequence with binary labels i.e. $0$ for no-topic-change and $1$ for topic-change. Specifically, these methods start by dividing the input document into sentences (sentence tokenization) using the NLTK PUNKT tokenizer [12]. Then, they use pretrained token embeddings to represent input tokens $\textbf{w}_{i}$. Concretely, word2vec embeddings [13, 14] are used in [4] while FastText embeddings [15] are used in [3]. Token embeddings of each sentence are fed to a sentence embedding model in order to create sentence representations i.e. each sentence is represented by a dense vector $\textbf{s}_{j}$. This sentence embedding model is a 2-layer bidirectional LSTM [16] in the method proposed in [4] while it represents a transformer encoder [17] in the method proposed in [3]. The sequence of sentence embeddings is then fed to a sentence labeling model which assigns a binary label $l_{i}\in\{0,1\}$ to each sentence where $l_{i}=1$ encoding the existence of a topic change. This sentence labeling model is a 2-layer bidirectional LSTM in [4] while it is a transformer encoder in [3]. Both of the aforementioned solutions are trained in a supervised manner where the labeled corpus WIKI-727K [4] is used for training. Note that the transformer-based solution adds another feed-forward head besides the sentence labeling model. This head is responsible for explicit coherence modeling [3]. A similar work was proposed in [18] where authors tested a solution based on a transformer to encode sentences (first level) and a bidirectional LSTM to detect topic changes (second level). For deeper details about how transformers and LSTMs work, interested readers are referred to this tutorial [19].

While the aforementioned solutions are supervised, many solutions based on unsupervised-learning have been proposed. For instance, the method proposed in [20] uses RoBERTa embeddings [21] to calculate semantic similarity scores between sentences. These scores are then used to perform segmentation. Particularly, the variation in similarity scores over time is leveraged to recognize changes in topics. The problem of representing sentences from token embeddings was addressed in two ways: 1) to use sentence-BERT [22], a version of BERT [23] fine-tuned on sentence representation tasks, and 2) to aggregate token embeddings from the 2nd to the last layer of RoBERTa using max pooling.

Before the deep-learning era, several interesting statistical solutions were proposed. Interested readers are referred for example to 1) TextTiling [24] which captures semantic similarity between sentences using word frequency vectors, 2) TopicTiling [25] an algorithm based on TextTiling [24] and Latent Dirichlet Allocation (LDA) [26], 3) LDA-based topic modeling [27], and 4) the statistical models for text segmentation proposed in [2].

The solutions discussed so far represent the semantic-aspects-based text segmentation family. Let us explore now some visual-aspects-based text segmentation methods i.e. segmentation applied on formatted documents. In fact, in the financial domain, several interesting solutions to the document-structure-extraction problem have been proposed thanks to the FinTOC shared tasks [8, 7, 9]. These latter aim at extracting the structure of complex layout formatted documents stored in pdf format. The objective is to detect titles and to hierarchically organize them into a Table of Contents (ToC). For instance, authors of [9] proposed a title detection algorithm based on Document Image Analysis (DIA). Particularly, they started by Faster R-CNN [28] model which is pertained on the PubLayNet dataset [29] and they fine tuned it on the FinTOC-2022 training set. In order to recognize the hierarchical levels of the detected titles, they extracted orthographic and layout features from the detected titles and they proposed a supervised random forest model to predict the title level.

III Data analysis and annotation

The experimental work has been performed on the discharge summaries corpus, a sub-corpus from the MIMIC-III [11] database. MIMIC-III contains more than two million reports including discharge summaries, EEG reports, radiology reports, and many others. Figure 3-a shows the distribution of MIMIC-III reports’ families. The discharge summaries corpus contains around $60~{}000$ reports. Figure 3-b shows the distribution of discharge summaries’ lengths given in number of tokens. The average discharge summary length is $1435$ tokens (roughly $3$ pages).

We divided the discharge summaries corpus to a test corpus of $250$ randomly selected reports and a training corpus containing the rest of reports. The test corpus has been manually annotated using the web application ACUITEE [30] as illustrated in Figure 1. Note that two hierarchical level of titles are considered: titles (highlighted in yellow) and sub-titles (highlighted in blue).

In order to annotate the training corpus, we conducted a deep analysis of the discharge summaries structure and we proposed a rule-based title detection system, namely TocRegex. This latter has been implemented using regular expressions. For example, the first considered pattern was any character sequence that:

• begins with a newline ”\n” followed by a valid title content and terminates with a colon followed by newline ”:\n”; or

• begins with a double newline ”\n\n” followed by a valid title content and terminates with a colon ”:”; or

• begins with a beginning-of-document character followed by a valid title content and terminates with a colon ”:”

The definition of ”valid title content” as well as the full list of the $12$ considered regular expressions are discussed in the Appendix A.
Figure 4 shows the frequencies of the top 35 matches detected using the above-described pattern. For instance, the title ”history of present illness” has been detected $10e5$ times in the corpus. We note that several false positives have also matched the considered pattern like e.g. ”tablet(s)*refills” which appears frequently in medication lists. In order to filter-out such kind of false positives, a manual curation has been applied to the detected titles. As a result, the total remaining numbers of unique titles is $56~{}357$. They have appeared $2~{}930~{}725$ times in the corpus i.e. the average number of titles is around $49$ per discharge summary.

Figure 3: a) distribution of sub-corpora sizes of the MIMIC-III database. 2) distribution of reports’ lengths of the discharge summaries corpus.

Figure 4: Top $35$ ”candidate” titles extracted using the first pattern of the TocRegex solution including true positives (blue) and false positives (light brown)

IV The proposed solution: TocBERT

In this section, we explore the proposed solution TocBERT (Table of Content BERT). The hierarchical title detection task is formulated as a sequence labeling problem i.e. token classification problem. Particularly, it is considered as a named entity recognition (NER) problem with three entity types: ”I-title” for titles, ”I-Stitle” for subtitles and ”O” for other tokens. Note that we adopt the ”Inside–outside–beginning” (IOB) standard labeling format. Figure 5 shows two examples of labeled sequences where tokens are tagged with title, subtitle or outside labels. For instance, the second sequence has one title ”Physical exam” and four subtitles ”HEENT”, ”Neck”, ”Lungs” and ”Extremities”.

Figure 5: Labeling token sequences in TocBERT: Example 1) two titles ”Past Medical History” and ”Family History” are labeled with ”I-title” while other tokens are labeled with ”O”; Example 2) one title ”Physical exam” is labeled with ”I-title” and four subtitles ”HEENT”, ”Neck”, ”Lungs” and ”Extremities” are labeled with ”I-Stitle”, while other (outside) tokens are labeled with ”O”.

TocBERT is based on fine-tuning the pretrained model Bio-ClinicalBERT [31] on the above-mentioned NER task. Note that bidirectional transformers, like the BERT family, are more convenient for tackling sequence labeling tasks, like NER, compared to generative i.e. auto-regressive transformers like GPT [32]. Bio-ClinicalBERT represents a variant of BERT adapted to biomedical and clinical domains. We fine-tuned this pretrained model using the discharge summaries training corpus that has been semi-automatically labeled using TocRegex followed by a manual curation as described in the previous section III.
Bio-ClinicalBERT, like BERT, is pretrained using a fixed size vocabulary: $28~{}996$ tokens. This vocabulary was created using the word-piece algorithm [33]. Some of the resultant tokens represent sub-words. For this reason, the first step of TocBERT is to project labels of the full-word tokens (extracted by simple pre-tokenization procedure) to sub-word tokens (extracted using Bio-ClinicalBERT tokenizer). The second step consists in preparing convenient training windows from the training corpus. This is important because the maximum window size of BERT is $512$ (sub-word) tokens while the average length of a discharge summary is $1435$ (full-word) tokens. To this end, discharge summaries have been segmented into windows of $384$ words where we used the approximate formula: $1$ token $=0.75$ words. The total size of the resultant training set is $144~{}000$ labeled windows. Finally, the last step consists of training the TocBERT model.

V Experimental results

V-A Experimental configurations

TocBERT was trained on one NVIDIA A100 GPU equipped with 80 GB of RAM. It was trained for $20$ epochs with a batch size of $16$ training samples. The training time was around $17$ hours. The inference hardware was an NVIDIA RTX A3000 laptop GPU equipped with 6 GB of dedicated RAM.

V-B Results

Table I shows the experimental results of the proposed text segmentation solutions, TocBERT and TocRegex, in the hierarchical segmentation configurations, i.e. detecting both titles and sub-titles. While TocBERT is trained on a corpus labeled using TocRegex, the former considerably outperforms the latter in all the considered performance criteria, i.e. precision, recall and F1-score.

TABLE I: Experimental results: hierarchical text segmentation

	Precision	Recall	F1-score
TocBERT	0.714	0.754	0.728
TocRegex	0.667	0.563	0.606

Table II shows the results of the aforementioned solutions in the linear text segmentation configurations. TocBERT and TocRegex show comparable overall performance measured by the F1-score metric. TocBERT shows higher sensitivity (recall) while it is less specific. Note that TocBERT is completely based on the semantic representation of titles and their context. It doesn’t depend on the existence of visual aspects like newlines and colons. This explains why it detects more titles (higher recall). On the other hand, TocRegex is more specific since it is entirely based on patterns that almost always exist only in titles.

Comparing the results of hierarchical and linear segmentation, the value of TocBERT is clearly its strong capacity to exploit the context and the semantic aspects to discriminate between titles and sub-titles.

TABLE II: Experimental results: linear text segmentation

	Precision	Recall	F1-score
TocBERT	0.829	0.878	0.846
TocRegex	0.932	0.784	0.845

TABLE III: Regular expressions used to detect titles

In terms of inference execution time, TocRegex takes $87$ ms while TocBERT takes $195$ ms, in average, to segment a discharge summary. While TocRegex is faster, both solutions satisfy near-real-time requirements.

Figure 6: Grou** of some key concepts towards constructing topics’ ontology

Finally, we observed that the detected groups of titles and subtitles can be exploited to automatically construct an ontology of topics. Figure 6 shows an example of initial manual grou** of key concepts that help constructing such ontology. The semantic vector representations of the detected titles can play an important role in organizing the topics’ ontology. The interest of building this ontology is to facilitate tasks like text segmentation and information retrieval. This idea will be investigated in future works.

VI Conclusion

In this paper, we proposed a new solution, TocBERT, for the hierarchical segmentation of medical reports. The segmentation task was formulated as a named entity recognition problem. TocBERT was initialized by a pretrained model, Bio-ClinicalBERT, and fine tuned on a the MIMIC-III discharge summaries corpus. This latter was semi-automatically labeled with titles and sub-titles. TocBERT showed very good results considerably outperforming a carefully-designed rule-based system. Particularly, it showed a good performance in discriminating between titles and subtitles by leveraging their semantic representations and employing their context.

The semantic representations of the extracted titles can be exploited to automatically construct an ontology of topics. Such an ontology can further facilitate tasks like text segmentation and information retrieval. The investigation of this idea is left for future work.

Appendix A TocRegex

Table III shows the regular expressions list used in the proposed rule-based solution, TocRegex. In the first line, the four supported title-content patterns are defined while the remaining $12$ lines list the full-title patterns.