Efficient Document Embeddings via Self-Contrastive Bregman Divergence Learning
Year: 2023

Abstract:
Learning quality document embeddings is a fundamental problem in natural language processing (NLP), information retrieval (IR), recommendation systems, and search engines. Despite recent advances in the development of transformer-based models that produce sentence embeddings with self-contrastive learning, the encoding of long documents (Ks of words) is still challenging with respect to both efficiency and quality considerations. Therefore, we train Longfomer-based document encoders using a state-of-the-art unsupervised contrastive learning method (SimCSE). Further on, we complement the baseline method siamese neural networkwith additional convex neural networks based on functional Bregman divergence aiming to enhance the quality of the output document representations. We show that overall the combination of a self-contrastive siamese network and our proposed neural Bregman network outperforms the baselines in two linear classification settings on three long document topic classification tasks from the legal and biomedical domains.

heading: 1 Introduction
text: The development of quality document encoders is of paramount importance for several NLP applications, such as long document classification tasks with biomedical (Johnson et al., 2016), or legal (Chalkidis et al., 2022b) documents, as well as information retrieval tasks (Chalkidis et al., 2021a; Rabelo et al., 2022; Nentidis et al., 2022). Despite the recent advances in the development of transformer-based sentence encoders (Reimers and Gurevych, 2019; Gao et al., 2021; Liu et al., 2021; Klein and Nabi, 2022a) via unsupervised contrastive learning, little do we know about the potential of neural document-level encoders targeting the encoding of long documents (Ks of words).
∗ The authors contributed equally to this work.
The computational complexity of standard Transformer-based models (Vaswani et al., 2017; Devlin et al., 2019) (PLMs) given the quadratic self-attention operations poses challenges in encoding long documents. To address this computational problem, researchers have introduced efficient sparse attention networks, such as Longformer (Beltagy et al., 2020), BigBird (Zaheer et al., 2020), and Hierarchical Transformers (Chalkidis et al., 2022a). Nonetheless, fine-tuning such models in downstream tasks is computationally expensive; hence we need to develop efficient document encoders that produce quality document representations that can be used for downstream tasks out-ofthe-box, i.e., without fully (end-to-end) fine-tuning the pre-trained encoder, if not at all.
Besides computational complexity, building good representation models for encoding long documents can be challenging due to document length. Long documents contain more information than shorter documents, making it more difficult to capture all the relevant information in a fixed-size representation. In addition, long documents may have sections with different topics, which increases the complexity of encoding that usually leads to collapsing representations (Jing et al., 2022). More-
12181
over, long documents can be semantically incoherent, meaning that content may not be logically related or may contain irrelevant information. For these reasons, it is challenging to create a quality representation that captures the most important information in the document.
To the best of our knowledge, we are the first to explore the application of self-contrastive learning for long documents (Table 1). The contributions of our work are threefold:
(i) We train Longfomer-based document encoders using a state-of-the-art self-contrastive learning method, SimCSE by Gao et al. (2021).
(ii) We further enhance the quality of the latent representations using convex neural networks based on functional Bregman divergence. The network is optimized based on self-contrastive loss with divergence loss functions (Rezaei et al., 2021).
(iii) We perform extensive experiments to highlight the empirical benefits of learning representation using unsupervised contrastive and our proposed enhanced self-contrastive divergence loss. We compare our method with baselines on three long document topic classification tasks from the legal and biomedical domain.

heading: 2 Related Work
text: Document Encoders The need for quality document representations has always been an active topic of NLP research. Initial work on statistical NLP focused on representing documents as Bag of Words (BoW), in which direction TF-IDF representations were the standard for a long time. In the early days of deep learning in NLP, models developed to represent words with latent representations, such as Word2Vec (Mikolov et al., 2013), and GloVe (Pennington et al., 2014). Within this research domain, the use of word embedding centroids as document embeddings, and the development of the Doc2Vec (Le and Mikolov, 2014) model were proposed. Given the advanced compute needs to encode documents with neural networks, follow-up work mainly developed around sentence/paragraph-level representations, such as Skip Thoughts of Kiros et al. (2015), which relies on an RNN encoder. In the era of pre-trained Transformer-based language models, Reimers and Gurevych (2019) proposed the Sentence Transformers framework in order to develop quality dense sentence representations. Many
works followed a similar direction relying on a selfsupervised contrastive learning setup, where most ideas are adopted mainly from Computer Vision literature (Chen et al., 2020; Bardes et al., 2022).
Self-Supervised Contrastive Learning in NLP Several self-contrastive methods have been proposed so far for NLP applications. To name a few: MirrorRoBERTa (Liu et al., 2021), SCD (Klein and Nabi, 2022b), miCSE (Klein and Nabi, 2022a), DeCluTR (Giorgi et al., 2021), and SimCSE (Gao et al., 2021) – described in Section 3.2–, all create augmented versions (views) of the original sentences using varying dropout and comparing their similarity. The application of such methods is limited to short sentences and relevant downstream tasks, e.g., sentence similarity, while these methods do not use any additional component to maximize diversity in latent feature representations.

heading: 3 Methods

heading: 3.1 Base Model - Longformer
text: We experiment with Longformer (Beltagy et al., 2020), a well-known and relatively simple sparseattention Transformer. Longformer uses two sets of attention, namely sliding window attention and global attention. Instead of using the full attention mechanism, the sliding-window attention gives local context higher importance. Given a fixed window size w, each token attends to 12 w tokens on the respective side. The required memory for this is O(n × w). Sliding-window attention is combined with global attention from/to the [CLS] token.
Domain-Adapted Longformer: As a baseline, we use LongformerDA models which are Longformer models warm-started from domain-specific PLMs. To do so, we clone the original positional embeddings 8× to encode sequences up to 4096 tokens. The rest of the parameters (word embeddings, transformers layers) can be directly transferred, with the exception of Longformer’s global attention K, Q, V matrices, which we warm-start from the standard (local) attention matrices, following Beltagy et al. (2020). All parameters are updated during training.
For legal applications (Section 4.1), we warmstart our models from Legal-BERT (Chalkidis et al., 2020), a BERT model pre-trained on diverse English legal corpora, while for the biomedical one, we use BioBERT (Lee et al., 2020), a BERT model pre-trained on biomedical corpora.

heading: 3.2 Self-supervised Contrastive Learning
text: To use our LongformerDA for self-supervised contrastive learning, we need to use a Siamese network architecture (left part of Figure 1). Assume we have mini-batch D = {(xi)}Ni=1 of N documents. As positive pairs (xi, xi+), the method uses augmented (noised) versions of the input feature xi. As negative pairs (xi, xi−), all remaining N-1 documents in a mini-batch are used. The augmentations take place in the encoder block fθ of the model. θ is the parameterization of the encoder. We use the SimCSE (Gao et al., 2021) framework, in which case the encoder fθ is a pre-trained language model, LongformerDA in our case, and augmentation comes in the form of varying token dropout (masking) rate (τ). The loss objective used in the unsupervised version of SimCSE is the multiple negatives ranking loss (ℓmnr):
ℓmnr = −1n n∑
i=1
exp ( f (si, s̃i)) ∑
j exp ( f ( si, s j )) (1)
where s̃i is the positive augmented input sequence in the mini-batch, and s̃ j are the negatives. Multiple negatives ranking loss takes a pair of representations (si, s̃i) and compares these with negative samples in a mini-batch. In our experiments, we train such models, dubbed LongformerDA+SimCSE.

heading: 3.2.1 Bregman Divergence Loss
text: We complement this method with an additional ensemble of subnetworks optimized by functional Bregman divergence aiming to improve the output document latent representations further. Specifically, the embedding of self-contrastive networks
further passes to k-independent subnetworks to promote diversity in feature representations. The si and s j vectors from the contrastive framework are mapped to k-independent ensemble of neural networks that are optimized using functional Bregman divergence.
Gϕ(sa, sb) = ϕ(sa) − ϕ(sb)−∫ [sa(x) − sb(x)]δϕ(sb)(x)dµ(x)
(2)
sa and sb are vectors output by the self-contrastive network, and ϕ is a strictly convex function and can be described via a linear functional, consisting of weights wk and biases ϵk. The function ϕ(sa) is approximate by:
ϕ(sa) = sup (w,ϵw)∈Q
∫ sa(x)w(x)dx + ϵw (3)
We take the empirical distribution of the projection representation to compute ŝa and ŝb. Specifically we define: ŝi=argmaxk[ ∫ sa(x)wk(x)dx + ϵk] for i = (a,b). Using the above specification and ϕ(sa), we get the following functional divergence term:
G(sa, sb) = ( ∫ sa(x)wŝa(x)dx + ϵŝa)−
( ∫ sa(x)wŝb(x)dx + ϵŝb) (4)
Each sub-network produces a separate output (right part of Figure 1). The divergence is then computed using the output at point ŝa and ŝb using the projections as input. We convert the divergence to similarity using a Gaussian kernel as done by Rezaei et al. (2021).1
ψ = exp ( −G/2σ2 ) (5)
The mini-batch has size N. For empirical distributions saα(zi), sb(z j) where i and j are the respective index for the two branches and z the projector representation, we have:
ℓBregman(sa(zi),sb(z j)) = − log(exp
( ψi, j )
∑N t=1 exp ( ψi,k ) )
(6)
The final objective function is computed on the combination of as follows:
LTotal = ℓmnr + λ · ℓBregman (7) 1Rezaei et al. (2021) explore various monotone transformations. The Gaussian kernel performed best compared to other transformation methods.
Where λ is a scalar hyperparameter to weigh the relative importance of the Bregman divergence and contrastive loss. In our experiments, we train such models, dubbed LongformerDA+SimCSE+Bregman.

heading: 4 Experimental Set-up

heading: 4.1 Datasets and Tasks
text: ECtHR (Chalkidis et al., 2021b) dataset contains 11k cases from the European Court of Human Rights (ECtHR). This is a multi-label topic classification task, where given the facts of an ECtHR case, the model has to predict the alleged violated ECtHR article among ten such articles (labels).
SCOTUS (Chalkidis et al., 2022b) dataset contains 4.7k cases from the Supreme Court of US (SCOTUS). This is a single-label multi-class topic classification task, where given a SCOTUS opinion, the model has to predict the relevant area among 14 issue areas (labels).
MIMIC (Johnson et al., 2016) dataset contains approx. 50k discharge summaries from US hospitals. Each summary is annotated with one or more codes (labels) from the ICD-9 hierarchy, which has 8 levels in total. We use the 1st level of ICD9, including 19 categories, respectively. This is a multi-label topic classification task, where given the discharge summary, the model has to predict the relevant ICD-9 top-level codes (labels).

heading: 4.2 Experimental Settings
text: To get insights into the quality of the learned representations out-of-the-box, we train classifiers using
document embeddings as fixed (frozen) feature representations. We consider two linear classification settings: (i) Linear evaluation plugging a MLP classification head on top of the document embeddings; (ii) Linear evaluation plugging a linear classifier on top of the document embeddings.

heading: 5 Results and Discussion
text: In Table 2, we present the results for all examined Longformer variants across the three examined datasets and two settings using macro-F1 (mF1) and micro-F1 (µ-F1) scores.
Classification performance: In the last line of Table 2, we present the results for the baseline LongformerDA model fine-tuned end-to-end, which is a ‘ceiling’ for the expected performance, comparing to the two examined linear settings, where the document encoders are not updated. We observe that in the SCOTUS dataset training models with an MLP head are really close to the ceiling performance (approx. 1-4p.p. less in µ-F1). The gap is smaller for both models trained with the self-contrastive objective (+SimCSE, +SimCSE+Bregman), especially the one with the additional Bregman divergence loss, where the performance decrease in µ-F1 is only 1 p.p.
In the other two datasets (ECtHR and MIMIC), the performance of the linear models is still approx. 10-15 p.p. behind the ceilings in µ-F1. In ECtHR, we find that self-contrastive learning improves performance in the first settings by 3 p.p. in µ-F1, while the additional divergence Bregman loss does not really improve performance. This is not the
case, in the second linear setting (second group in Table 2), where the baseline outperforms both models. Similarly in MIMIC, we observe that selfcontrastive learning improves performance in the first settings by 3 p.p. in µ-F1, but the performance is comparable given linear classifiers. Overall, our enhanced self-contrastive method leads to the best results compared to its counterparts.
In Table 3, we also present results on SCOTUS in a few-shot setting using the SetFit (Tunstall et al., 2022) framework, where Bregman divergence loss improves performance compared to the baselines.
Given the overall results, we conclude that building subnetwork ensembles on top of the document embeddings can be a useful technique for encoding long documents and can help avoid the problem of collapsing representations, where the model is unable to capture all the relevant information in the input. Our approach has several advantages for long-document processing:
Efficiency considerations: In Table 2, we observe that in both linear settings where fixed document representations are used, the training time is 2- 8× decreased compared to end-to-end fine-tuning, while approx. 0.5% of the parameters are trainable across cases, which directly affects the compute budget. We provide further information on the size of the models in Appendix B.
Avoidance of collapsing representations: When processing long documents, there is a risk that the representation will collapse (Jing et al., 2022), meaning that the model will not be able to capture all the relevant information in the input. By mapping the document embedding from the base encoder into smaller sub-networks, the risk of collapsing representations is reduced, as the divergence loss attempts to reduce redundancy in the feature representation by minimizing the correlation. The results shown in Table 3 in a low-resource setting further highlight the advantage of training a Longformer with contrastive divergence learning.

heading: 6 Conclusions and Future Work
text: We proposed and examined self-supervised contrastive divergence learning for learning representation of long documents. Our proposed method is composed of a self-contrastive learning framework followed by an ensemble of neural networks that are optimized by functional Bregman divergence. Our method showed improvement compared to the baselines on three long document topic classifications in the legal and biomedical domains, while the improvement is more vibrant in a few-shot learning setting. In future work, we would like to further investigate the impact of the Bregman divergence loss in more classification datasets and other NLP tasks, e.g., document retrieval.
Limitations
In this work, we focus on small and medium size models (up to 134M parameters), while recent work in Large Language Models (LLMs) targets models with billions of parameters(Brown et al., 2020; Chowdhery et al., 2022). It is unclear how well the performance improvement from the examined network architecture would translate to other model sizes or baseline architectures, e.g., GPT models.
Further on, it is unclear how these findings may translate to other application domains and datasets, or impact other NLP tasks, such as document retrieval/ranking. We will investigate these directions in future work.

heading: Acknowledgments
text: Mina Rezai and Bernd Bisch were supported by the Bavarian Ministry of Economic Affairs, Regional Development and Energy through the Center for Analytics – Data – Applications (ADA-Center) within the framework of BAYERN DIGITAL II (20-3410-2-9-8).M. R. and B. B. were supported by the German Federal Ministry of Education and Research (BMBF) Munich Center for Machine Learning (MCML). This work was also partly funded by the Innovation Fund Denmark (IFD).2

heading: A Hyper-parameter Optimization
text: Continued Pre-training: We define the search space based on previous studies such as Rezaei et al. (2021) and Gao et al. (2021). For the contrastive Bregman divergence, we benchmark the performance for the first-stage hyper-parameters on the downstream task to tune the respective hyper-parameters. We use mean pooling for all settings. The learning rate, the total optimization steps, the use of a batch-norm layer, the σ parameter, the number of sub-networks g, and the batch size are grid-searched. Temperature (.1) and the input length to 4096 are fixed beforehand. The learning rate for these models was 3e-5. We run 50.000 optimization steps for each model.
Training for classification tasks: We used AdamW as an optimizer. Bayesian optimization is used to tune the hype-rparameters learning rate, number of epochs and batch size. We use mean pooling for all settings. Early stopping is set to a patience score of 3.3 These parameters were fixed after some early experiments. We use a learning rate of 1e-4 and run ECTHR and SCOTUS for 20 and
3We also experimented with other patience scores but experiments suggest that 3 epochs results in the best performance.
30 epochs respectively for the MLP head setting. For MIMIC we used 10 epochs for the MLP head and had to truncate the maximum sequence length to 2048 due to computational constraints. For each task we compared multiple different training checkpoints of our encoder. The reported results are the best performing checkpoints.

heading: B Number of parameters
text: Table 5 shows the number of parameters for the different models. Modding the transformer to a Longformer adds 6M parameters for LegalBERT small and 24M parameters for BioBERT medium. By working with LegalBERT-small and BioBERTbase we cover both small and medium sized models.

heading: C Pooling methods
text: We evaluate Mean, Max and [CLS] pooling. Results for end-to-end fine-tuning can be found in the table 6. Our results show that using mean pooling during continued pre-training in combination with max-pooling for classification could further enhance the performance instead of using the same pooling method for both stages.

heading: D Neural network Architecture
text: Our model contains two linear layers with one activation layer and two batch normalization layers.
We also compare the model without batch normalization layers. The comparison is made on the SCOTUS dataset using end-to-end fine-tuning. One can see that removing batch normalization worsens performance.
ACL 2023 Responsible NLP Checklist

heading: A For every submission:

heading: 3 A1. Did you describe the limitations of your work?
text: Unnumbered
7 A2. Did you discuss any potential risks of your work? We don’t see any direct potential risk.

heading: 3 A3. Do the abstract and introduction summarize the paper’s main claims?
text: Section 1
7 A4. Have you used AI writing assistants when working on this paper? Left blank.
B 7 Did you use or create scientific artifacts? Left blank.

heading: 3 B1. Did you cite the creators of artifacts you used?
text: Section 2
7 B2. Did you discuss the license or terms for use and / or distribution of any artifacts? We re-use already available public resources.
B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? Not applicable. Left blank.
B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? Not applicable. Left blank.
B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? Not applicable. Left blank.
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. Left blank.
C 3 Did you run computational experiments? Section 4

heading: 3 C1. Did you report the number of parameters in the models used, the total computational budget
text: (e.g., GPU hours), and computing infrastructure used? Appendix A.2
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.

heading: 3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found
text: hyperparameter values? Section 4
7 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? Left blank.
7 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? Because we don’t.
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? Not applicable. Left blank.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? Not applicable. Left blank.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? Not applicable. Left blank.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? Not applicable. Left blank.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? Not applicable. Left blank.

heading: Future Work
text: In future work, we would like to further investigate the impact of the Bregman divergence loss in more classification datasets and other NLP tasks, e. g. , document retrieval. Limitations
In this work, we focus on small and medium size models (up to 134M parameters), while recent work in Large Language Models (LLMs) targets models with billions of parameters(Brown et al. , 2020; Chowdhery et al. , 2022). It is unclear how well the performance improvement from the examined network architecture would translate to other model sizes or baseline architectures, e. g. , GPT models. Further on, it is unclear how these findings may translate to other application domains and datasets, or impact other NLP tasks, such as document retrieval/ranking. We will investigate these directions in future work