Structure-Aware Language Model Pretraining Improves Dense Retrieval on Structured Data (2023)

Abstract: This paper presents Structure Aware Dense Retrieval (SANTA) model, which encodes user queries and structured data in one universal embedding space for retrieving structured data. SANTA proposes two pretraining methods to make language models structure-aware and learn effective representations for structured data: 1) Structured Data Alignment, which utilizes the natural alignment relations between structured data and unstructured data for structure-aware pretraining. It contrastively trains language models to represent multi-modal text data and teaches models to distinguish matched structured data for unstructured texts. 2) Masked Entity Prediction, which designs an entity-oriented mask strategy and asks language models to fill in the masked entities. Our experiments show that SANTA achieves state-of-the-art on code search and product search and conducts convincing results in the zero-shot setting. SANTA learns tailored representations for multi-modal text data by aligning structured and unstructured data pairs and capturing structural semantics by masking and predicting entities in the structured data. All codes are available at https://github.com/OpenMatch/OpenMatch .

## introduction dense retrieval has shown strong effectiveness in lots of nlp applications, such as open domain question answering @xcite , conversational search @xcite , and fact verification @xcite . it employs pretrained language models (plms) to encode unstructured data as high-dimensional embeddings, conduct text matching in an embedding space and return candidates to satisfy user needs @xcite . besides unstructured data, structured data, such as codes, html documents and product descriptions, is ubiquitous in articles, books, and web pages, and plays the same important roles in understanding text data. learning the semantics behind text structures to represent structured data is crucial to building a more self-contained retrieval system. the structured data modeling stimulates researchers to build several benchmarks to evaluate model performance, such as code search and product search @xcite . the structured data retrieval tasks require models to retrieve structured data according to user queries. dense retrieval @xcite shows a promising way to build a retrieval system on structured data by encoding user queries and structured data in an embedding space and conducting text matching using the embedding similarity. nevertheless, without structure-aware pretraining, most plms lack the necessary knowledge to understand structured data and conduct effective representations for retrieval @xcite @xcite . lots of structure-aware pretraining methods are proposed to continuously train plms to be structure-aware and better represent structured data @xcite . they design task-specific masking strategies and pretrain plms with mask language modeling. nevertheless, only using mask language modeling may not sufficiently train plms to conduct effective representations for structured data @xcite . some natural alignment signals between structured and unstructured data, such as code-description documentation and product description-bullet points, provide an opportunity to pretrain the structured data representations. using these alignment signals, plms can be contrastively trained @xcite to match the representations of aligned structured and unstructured data and understand the semantics of structured data with the help of natural language. in this paper, we propose structure aware dense retrieval (santa), a dense retrieval method on structured data. as shown in figure 1 , santa encodes queries and structured data in an embedding space for retrieval. santa designs two pretraining tasks to continuously train plms and make plms sensitive to structured data. the structured data alignment task contrastively trains plms to align matched structured-unstructured data pairs in the embedding space, which helps to represent structured data by bridging the modality gap between structured and unstructured data. the masked entity prediction task masks entities and trains plms to fill in the masked parts, which helps to capture semantics from structured data. our experiments show that santa achieves state-of-the-art in retrieving structured data, such as codes and products. by aligning structured and unstructured data, santa maps both structured and unstructured data in one universal embedding space and learns more tailored embeddings for multi-modal text data matching. the masked entity prediction task further guides santa to capture more crucial information for retrieval and better distinguish structured and unstructured data. depending on these pretraining methods, santa can even achieve comparable retrieval results with existing code retrieval models without finetuning, showing that our structure-aware pretraining can benefit structured data understanding, multi-modal text data representation modeling and text data match-ing between user queries and structured data. ## related work dense retrieval @xcite @xcite encodes queries and documents using pretrained language model (plm) @xcite @xcite and maps them in an embedding space for retrieval. however, during retrieving candidates, the documents can be passages in natural language @xcite , images @xcite , structured data documents @xcite or multi-modal documents @xcite , which challenges existing dense retrieval models to handle different kinds of modalities of knowledge sources to build a self-contained retrieval system. existing work @xcite also builds dense retrievers for retrieving structured data and mainly focuses on learning representations for code data. leaning more effective representations with plms is crucial for dense retrieval @xcite , thus several continuous training models are proposed. they usually employ mask language modeling to train plms on structured data and help to memorize the semantic knowledge using model parameters @xcite @xcite . codebert uses replaced token detection @xcite and masked language modeling @xcite to learn the lexical semantics of structured data @xcite . dobf @xcite further considers the characteristics of code-related tasks and replaces class, function and variable names with special tokens. codet5 @xcite not only employs the span mask strategy @xcite but also masks the identifiers in codes to teach t5 @xcite to generate these identifiers, which helps better distinguish and comprehend the identifier information in code-related tasks. nevertheless, the mask language modeling @xcite may not sufficiently train plms to represent texts and show less effectiveness in text matching tasks @xcite @xcite @xcite . the recent development of sentence representation learning methods has achieved convincing results @xcite . the work first constructs sentence pairs using backtranslation @xcite , some easy deformation operations @xcite , original sequence cropping @xcite or adding dropout noise @xcite . then they contrastively train plms to learn sentence representations that can be used to distinguish the matched sentence pairs with similar semantics. ## methodology in this section, we introduce our structure aware dense retrieval (santa) model. first, we introduce the preliminary of dense retrieval (sec. 3.1). and then we describe our structure-aware pretraining method (sec. 3.2). ## experimental methodology in this section, we describe the datasets, evaluation metrics, baselines, and implementation details in our experiments. dataset. the datasets in our experiments consist of two parts, which are used for continuous training and finetuning, respectively. continuous training. during continuous training, two datasets, codesearchnet @xcite and esci (large) @xcite , are employed to continuously train plms to conduct structure-aware text representations for codes and shopping products. in our experiments, we regard code documentation descriptions and product bullet points as unstructured data for aligning structured data, codes and product descriptions, during training. more details of pretraining data processing are shown in appendix a.2. finetuning. for downstream retrieval tasks on structured data, we use adv @xcite , and esci (small) @xcite to finetune models for code search and product search, respectively. all data statistics are shown in evaluation metrics. we use mrr@100 and ndcg@100 to evaluate model performance, which is the same as the previous work @xcite @xcite . baselines. we compare santa with several dense retrieval models on code search and product search tasks. we first employ three pretrained language models to build dense retrievers for structured data retrieval, including bert @xcite , roberta @xcite and t5 @xcite , which are widely used in existing dense retrieval models @xcite @xcite . all these models are trained with in-batch negatives @xcite . for the code search task, we also compare santa with three typical and task-specific models, codebert @xcite , codet5 @xcite and coderetriever @xcite . codebert inherits the bert architecture and is trained on code corpus using both mask language modeling and replaced token detection. codet5 employs the encoder-decoder architecture for modeling different code-related tasks and teaches the model to focus more on code identifiers. coderetriever is the state-of-the-art, which continuously trains @xcite with unimodal and bimodal contrastive training losses. implementation details. this part describes the experiment details of santa. we initialize santa with t5-base and codet5base for product search and code search. for masked entity prediction, we regard code identifiers and some noun phrases as entities in codes and product descriptions, respectively. more details about identifying entities are shown in appendix a.3. during continuous training, we set the learning rate as 1e-4 and 5e-5 for product search and code search, and the training epoch as 6. during finetuning, we conduct experiments by training santa using inbatch negatives and hard negatives. we set the training epoch to 60 and learning rate to 5e-5 for product search, while the training epoch and learning rate are 6 and 1e-5 for code search. and we follow ance @xcite , start from inbatch finetuned santa (inbatch) model and continuously finetune it with hard negatives to conduct the santa (hard negative) model. the learning rates are set to 1e-5 and 1e-6 for product search and code search. these hard negatives are randomly sampled from the top 100 retrieved negative codes/product descriptions from the santa (inbatch) model. all models are implemented with pytorch, huggingface transformers @xcite and open-match @xcite . we use adam optimizer to optimize santa, set the batch size to 16 and set the warmup proportion to 0.1 in our experiments. ## evaluation results in this section, we focus on exploring the performance of santa on code search and product search tasks, the advantages of santa in representing structured data, and the effectiveness of proposed pretraining methods. ## conclusion this paper proposes santa, which pretrains language models to understand structure semantics of text data and guides language models to map both queries and structured texts in one universal embedding space for retrieval. santa designs both structured text alignment and masked entity prediction tasks to continuously train pretrained language models to learn the semantics behind data structures. our experiments show that santa achieves state-of-the-art on code and product search by learning more tailored representations for structured data, capturing semantics from structured data and bridging the modality gap between structured and unstructured data.