LinkBERT: Enhancing Language Models with Document Link Information

Language models (LMs) like BERT and the GPT series have become foundational to modern Natural Language Processing (NLP) systems, powering applications from search engines to personal assistants. Their success stems from a pretraining phase using self-supervised learning on vast amounts of text data, allowing them to learn rich world knowledge and adapt quickly to new tasks without extensive fine-tuning. Common pretraining objectives include masked language modeling (MLM), where the model predicts masked words, and causal language modeling, where it predicts the next word in a sequence.

The Challenge of Isolated Documents

A significant limitation in many existing LM pretraining strategies is their tendency to process documents in isolation. This approach overlooks the inherent interconnectedness of information, particularly in data sources like the web and scientific literature, which are rich with hyperlinks and citation links. These links are not merely navigational aids; they represent crucial dependencies that can convey knowledge spanning multiple documents. For instance, an article about the Tidal Basin in Washington D.C. might mention the National Cherry Blossom Festival, and a linked article could detail that the festival celebrates Japanese cherry trees. Together, these linked documents provide multi-hop knowledge, such as “Tidal Basin has Japanese cherry trees,” which is not explicitly available within either document alone.

Models trained without considering these inter-document relationships may fail to capture this distributed knowledge, which is vital for tasks requiring multi-hop reasoning or comprehensive understanding, such as question answering and knowledge discovery.

Introducing LinkBERT: A Graph-Aware Approach

To address this limitation, researchers have developed LinkBERT, a novel pretraining method designed to incorporate document link information for training more knowledgeable language models. LinkBERT’s approach involves three key steps:

1. Document Graph Construction: The first step is to build a document graph by identifying and linking related documents within a corpus. This is typically achieved by treating each document as a node and establishing directed edges between documents based on hyperlinks or citation links, which are generally indicative of high relevance.
2. Link-Aware Input Creation: To enable the LM to learn dependencies across linked documents, LinkBERT creates special input instances. Each document is segmented into smaller pieces. Then, pairs of segments are concatenated to form training sequences. These pairs can be:
   • Contiguous segments: Segments from the same document, mimicking previous LM approaches.
   • Random segments: Segments sampled from two different, unrelated documents.
   • Linked segments: Segments sampled from documents that are linked in the document graph.

   This variety of input types helps the model learn to recognize relationships between different text segments.

3. Link-Aware LM Pretraining: The final step involves pretraining the LM using these link-aware input instances with two joint self-supervised tasks:
   • Masked Language Modeling (MLM): By masking tokens within these concatenated segments, LinkBERT encourages the model to learn multi-hop knowledge. For example, if segments from articles about the Tidal Basin, the Cherry Blossom Festival, and Japanese cherry trees are placed together, the LM can learn the relationships between these concepts.
   • Document Relation Prediction (DRP): This task trains the model to classify the relationship between two segments (contiguous, random, or linked). This objective helps the LM learn document relevance and dependencies, as well as bridging concepts that connect related documents.

   These pretraining tasks can be conceptualized as graph-based self-supervised learning: node feature prediction (MLM) and link prediction (DRP).

Performance and Strengths of LinkBERT

LinkBERT has been evaluated on various downstream NLP tasks in both general and biomedical domains. In the general domain, it uses Wikipedia with its hyperlinks, while in the biomedical domain, it utilizes PubMed with citation links, resulting in models named LinkBERT and BioLinkBERT, respectively.

The evaluations demonstrate that LinkBERT consistently outperforms baseline language models pretrained without document links across a range of tasks and domains. The improvements are particularly notable in the biomedical domain, where citation links are crucial for understanding scientific literature. BioLinkBERT, for instance, has achieved state-of-the-art results on benchmarks like BLURB, MedQA, and MMLU.

LinkBERT exhibits several key strengths:

• Effective Multi-hop Reasoning: LinkBERT significantly improves performance on tasks requiring multi-hop reasoning. For example, in question answering scenarios, it can correctly infer answers by connecting information across multiple documents, unlike standard BERT which might be limited to information within a single document.
• Improved Document Relation Understanding: The DRP task during pretraining equips LinkBERT with a better ability to understand relationships between documents. This makes it more robust to noisy or irrelevant documents in open-domain question answering, maintaining accuracy where BERT might falter.
• Data Efficiency and Few-Shot Learning: LinkBERT demonstrates superior performance in few-shot and data-efficient learning scenarios. When fine-tuned with only a small fraction of training data, it shows substantial gains over BERT, indicating that it has internalized more world knowledge during pretraining.

Using LinkBERT

LinkBERT is designed to be a drop-in replacement for existing BERT models. Pretrained LinkBERT and BioLinkBERT models are available on HuggingFace, allowing easy integration into NLP projects. Researchers can leverage these models by loading them with standard libraries and then fine-tuning them for specific downstream tasks, such as question answering or text classification.

Pros and Cons

Pros:

• Significantly enhances language model capabilities by incorporating inter-document relationships.
• Improves performance on knowledge-intensive tasks, particularly those requiring multi-hop reasoning.
• Demonstrates greater robustness to irrelevant information and better data efficiency in low-resource settings.
• Achieves state-of-the-art results in specialized domains like biomedical NLP.
• Easy to integrate as a direct replacement for existing BERT-based models.

Cons:

• Requires a corpus with well-defined document links (hyperlinks, citations) for effective pretraining.
• The process of constructing a document graph can be computationally intensive for very large corpora.
• Effectiveness may vary depending on the density and quality of links within the training data.

Key Takeaways

• LinkBERT is a novel pretraining method that leverages document links to imbue language models with richer world knowledge and improved reasoning abilities.
• By treating documents as nodes in a graph and incorporating link-aware pretraining tasks (MLM and DRP), LinkBERT learns to connect information across multiple documents.
• This approach leads to significant performance gains in tasks requiring multi-hop reasoning, document relation understanding, and few-shot learning.
• LinkBERT models are readily available and can be used as a direct substitute for standard BERT models, offering a straightforward way to enhance NLP applications.
• The research opens avenues for future work, including extending this approach to other model architectures and modalities.