Text Segmentation by Cross Segment Attention (2020)

Abstract: Document and discourse segmentation are two fundamental NLP tasks pertaining to breaking up text into constituents, which are commonly used to help downstream tasks such as information retrieval or text summarization. In this work, we propose three transformer-based architectures and provide comprehensive comparisons with previously proposed approaches on three standard datasets. We establish a new state-of-the-art, reducing in particular the error rates by a large margin in all cases. We further analyze model sizes and find that we can build models with many fewer parameters while keeping good performance, thus facilitating real-world applications.

## introduction text segmentation is a traditional nlp task that breaks up text into constituents, according to predefined requirements. it can be applied to documents, in which case the objective is to create logically coherent sub-document units. these units, or segments, can be any structure of interest, such as paragraphs or sections. this task is often referred to as document segmentation or sometimes simply text segmentation. in figure 1 we show one example of document segmentation from wikipedia, on which the task is typically evaluated @xcite . documents are often multi-modal, in that they cover multiple aspects and topics; breaking a document into uni-modal segments can help improve and/or speed up down stream applications. for example, document segmentation has been shown to improve information retrieval by indexing subdocument units instead of full documents @xcite . other applications such as summarization and information extraction can also benefit from text segmentation @xcite . the aim of document segmentation is breaking the raw text into a sequence of logically coherent sections (e.g., "early life and marriage" and "legacy" in our example). a related task called discourse segmentation breaks up pieces of text into sub-sentence elements called elementary discourse units @xcite . edus are the minimal units in discourse analysis according to the rhetorical structure theory @xcite . in figure 2 we show examples of edu segmentations of sentences. for example, the sentence "annuities are rarely a good idea at the age 35 because of withdrawal restrictions" decomposes into the following two edus: "annuities are rarely a good idea at the age 35" and "because of withdrawal restrictions", the first one being a statement and the second one being a justification in the discourse analysis. in addition to being a key step in discourse analysis @xcite , discourse segmentation has been shown to improve a number of downstream tasks, such as text summarization, by helping to identify fine-grained sub-sentence units that may have different levels of importance when creating a summary @xcite . multiple neural approaches have been recently proposed for document and discourse segmentation. @xcite proposed the use of sentence 1: annuities are rarely a good idea at the age 35 because of withdrawal restrictions sentence 2: wanted: an investment that's as simple and secure as a certificate of deposit but offers a return worth getting excited about. figure 2 : example discourse segmentations from the rst-dt dataset @xcite . in the segmentations, the edus are separated by the character. however, these approaches are evaluated on different datasets and as such have not been compared against one another. furthermore they mostly rely on rnns instead of the more recent transformers @xcite and in most cases do not make use of contextual embeddings which have been shown to help in many classical nlp tasks @xcite . in this work we aim at addressing these limitations and make the following contributions: in particular, multiple neural approaches have been proposed for the task. in one, a sequence labeling algorithm is proposed where each sentence is encoded using a bi-lstm over tokens, and then a bi-lstm over sentence encodings is used to label each sentence as ending a segment or not @xcite . authors consider a large dataset based on wikipedia, and report improvements over unsupervised text segmentation methods. in another work, a sequence-to-sequence model is proposed @xcite , where the input is encoded using a bigru and segment endings are generated using a pointer network @xcite . the authors report significant improvements over sequence labeling approaches, however on a dataset composed of 700 artificial documents created by concatenating segments from random articles from the brown corpus @xcite . lastly, @xcite consider an attention-based cnn-bi-lstm model and evaluate it on three small-scale datasets. discourse segmentation contrary to document segmentation, discourse segmentation has historically been framed as a supervised learning task. however, a challenge of applying supervised approaches for this type of segmentation is the fact that the available dataset for the task is limited @xcite . for this reason, approaches for discourse segmentation usually rely on external annotations and resources to help the models generalize. early approaches to discourse segmen-tation were based on features from linguistic annotations such as pos tags and parsing trees @xcite @xcite . the performance of these systems was highly dependent on the quality of the annotations. recent approaches started to rely on end-to-end neural network models that do not need linguistic annotations to obtain high-quality results, relying instead on pretrained models to obtain word or sentence representations. an example of such work is by @xcite , which proposes a sequenceto-sequence model getting a sequence of glove @xcite word embeddings as input and generating the edu breaks. another approach utilizes elmo pretrained embeddings in the crf-bi-lstm architecture and achieves state-of-the-art results on the task @xcite . ## architectures we propose three model architectures for segmentation. one uses only local context around each candidate break, while the other two leverage the full context from the input (by candidate break, we mean any potential segment boundary). all our models rely on the same preprocessing technique and simply feed the raw input into a word-piece (sub-word) tokenizer @xcite . we use the word-piece tokenizer implementation that was open-sourced as part of the bert release @xcite , more precisely its english, uncased variant, which has a vocabulary size of 30,522 word-pieces. ## results in to train our models, we used the adamw optimizer @xcite with a 10% dropout rate as well as a linear warmup procedure. learning rates are set between 1e-5 and 5e-6, chosen to maximize the f1-score on the validation sets from each dataset. for the more expensive models, and especially on the wiki-727k dataset, we trained our models using google cloud tpus. we can see from the table that our models outperform the baselines across all datasets, reducing the relative error margins from the best baseline by 20%, 16% and 79% respectively on the wiki-727k, rst-dt and choi datasets. the improvements are statistically significant for all datasets. the errors are impressively low on the choi dataset, but it is important to point out that it is a small-scale synthetic dataset, and as such limited. since each document is a concatenation of extracts from random news articles, it is an artificially easy task for which a previous neural baseline achieved an already low error margin. moreover, on this dataset, the crosssegment bert model obtains very good results compared to the hierarchical models which do not attend across the candidate break. this aligns with the expectation that locally attending across a segment break is sufficient here, as we expect large semantic shifts due to the artificial nature of the dataset. hierarchical models, with a sentence encoder followed by a document encoder, perform well on the rst-dt dataset. as a reminder, this discourse segmentation task is about segmenting individual sentences so there is no notion of document context. in order to study whether the hierarchical structure is really necessary for discourse segmentation, we also trained a model without the bi-lstm (that is, making predictions directly using bert): this decreased the f1-score by 0.4%. it is also worth noting that several known lstm downsides were particularly apparent on the wiki-727k dataset: the model was harder to train and significantly slower during both training and inference. regarding the hierarchical bert model, different initialization methods were used for the two document segmentation datasets. on the choi dataset, a hibert initialization (a model fully pretrained end-to-end for hierarchical bert, similarly to @xcite was necessary to get good results, due the small dataset size. on the contrary, we obtained slightly better results initializing both levels of the hierarchy with bert base on the wiki-727k dataset, even though the model took longer to converge. other initializations, e.g., random for both levels of the hierarchy or bert base at the lower level and random at the upper level, gave worse results. perhaps the most surprising result from ## analyses in this section we perform additional analyses and ablation studies to better understand our segmentation models. experiments revolve around the cross-segment bert model. we choose this model because it has several advantages over its alternatives: ## conclusion in this paper, we introduce three new model architectures for text segmentation tasks: a crosssegment bert model that uses only local context around candidate breaks, as well as two hierarchical models, bert+bi-lstm and hierarchical bert. we evaluated these three models on document and discourse segmentation using three standard datasets, and compared them with other recent neural approaches. our experiments showed that all of our models improve the current state-of-theart. in particular, we found that a cross-segment bert model is extremely competitive with hierarchical models which have been the focus of recent research efforts @xcite . this is surprising as it suggests that local context is sufficient in many cases. due to its simplicity, we suggest at least trying it as a baseline when tackling other segmentation problems and datasets. naturally these results do not imply that hierarchical models should be disregarded. we showed they are strong contenders and we are convinced there are applications where local context is not sufficient. we tried several encoders at the upperlevel of the hierarchy. our experiments suggest that deep transformer encoders are useful for encoding long and complex inputs, e.g., documents for document segmentation applications, while bi-lstms proved useful for discourse segmentation. moreover, rnns in general may also be useful for very long documents as they are able to deal with very long input sequences. finally, we performed ablation studies to better understand the role of context and model size. consequently, we showed that distillation is an effective technique to build much more compact models to use in practical settings. in future work, we plan to further investigate how different techniques apply to the problem of text segmentation, including data augmentation @xcite and methods for regularization and mitigating labeling noise @xcite .