Contrastive Document Representation Learning with Graph Attention Networks (2021)

Abstract: Recent progress in pretrained Transformer-based language models has shown great success in learning contextual representation of text. However, due to the quadratic self-attention complexity, most of the pretrained Transformers models can only handle relatively short text. It is still a challenge when it comes to modeling very long documents. In this work, we propose to use a graph attention network on top of the available pretrained Transformers model to learn document embeddings. This graph attention network allows us to leverage the high-level semantic structure of the document. In addition, based on our graph document model, we design a simple contrastive learning strategy to pretrain our models on a large amount of unlabeled corpus. Empirically, we demonstrate the effectiveness of our approaches in document classification and document retrieval tasks.

## introduction document representations that capture the semantics are crucial to various document-level natural language processing (nlp) tasks, including sentiment analysis @xcite ), text classification @xcite and information retrieval @xcite . in recent years, an increasing volume of work has focused on learning a task-agnostic universal representation for long documents. while improved performance in downstream tasks have been achieved, there are two challenges towards learning a high quality document representation: (1) absence of document structure. most works treat the document as a sequence of tokens without considering high-level structure. transformers-based pretrained language models are ubiquitously state-of-the-art across many nlp tasks. transformer models such as bert @xcite and its variants have shown great success in learning contextual representation of text. representation from large language models can partially mitigate the data scarcity issue due to pretraining on a large amounts of unlabeled data. however, those models mostly consider token-level information and their pretraining tasks are not directly targeting long document representations. another issue of directly applying transformer-based models is the limit of the input text length. due to the quadratic complexity of self-attention, most of the pretrained transformers models can only handle a relatively short text. a wide spectrum of efficient, fast transformer models (collectively called "xformers") have been proposed to tackle this problem; e.g., longformer @xcite and bigbird @xcite use sparse attention to improve the computational and memory efficiency for long sequence text. nevertheless, these models still focus on token-level interactions without considering high-level semantic structure of the document. recently, there is a resurgence of interest in contrastive learning (cl) due to its success in selfsupervised representation learning in computer vision @xcite . contrastive learning offers a simple method to learn disentangled representation that encodes invariance to small and local changes in the input data without using any labeled data. in nlp domain, contrastive learning has been employed to learn sentence representation @xcite under either self-supervised or supervised settings. in this work, we propose a graph attention network (gat) based model that explicitly utilizes the high-level semantic structure of the documents to learn document embeddings. we model the document as not just a sequence of text, but a collection of passages or sentences. specifically, the proposed model introduces a graph on top of the document passages (fig. 1 ) to utilize multi-granularity information. first, passages are encoded using roberta @xcite to collect word-level knowledge. then passages are connected to leverage the higher-level structured information. at last, a graph attention network @xcite is applied to obtain the multi-granularity document representation. to better learn the document embedding, we propose a document-level contrastive learning strategy to pretrain our models. in our contrastive learning framework, we split the document into random sub-documents and train the model to maximize the agreement over the representations of the sub-documents that come from the same document. this simple strategy allows us to pretrain our models on a large unlabelled corpus without any additional priors. as we will see, this simple pretraining task indeed helps the model on the downstream tasks. the contributions of this paper can be summarized as follows. ## methodology in this section, we describe our main model and contrastive pretraining strategy. an example of the graph-roberta architecture for the document representation. ## conclusions in this work, we propose a simple graph attention network model to learn document embeddings. our model not only can leverage the recent advancement of pretrained transformer models as building blocks, but also explicitly utilize the high-level structure of the documents. in addition, we propose a simple document-level contrastive learning strategy that does not require handcraft transformations. with this strategy, we conduct large scale contrastive pretraining on a large corpus. empirically we demonstrate our methods achieve great performance on both document classification and document retrieval tasks.