BigBird: Expanding the Frontier of Language Models

Over the past few years, transformer-based language models like BERT, GPT, and XLNet have achieved remarkable success across a wide range of natural language processing tasks. These models learn rich linguistic representations by pre-training on massive text corpora, allowing them to perform well with limited labeled data.

However, a key technical limitation of existing transformer models is their quadratic dependency on input sequence length. Due to the dense all-to-all attention mechanism, the memory and computation requirements grow quadratically with the number of tokens. With current hardware, this typically limits transformer models to sequences of just 512 tokens.

For many language understanding tasks, much longer context is needed to capture the full meaning and nuance. Documents like news articles, scientific papers, and books can span thousands of words. Many dialog tasks require tracking context over long conversations with many turns. In the biomedical domain, genomic sequences can stretch to millions of bases.

BigBird is a recent model architecture that aims to address this limitation by extending transformers to much longer sequences. Using a sparse attention mechanism, BigBird can efficiently process sequences up to 8 times longer than standard transformers, while preserving their empirical performance. In this post, we‘ll take an in-depth look at how BigBird works and explore its potential to expand the frontier of language modeling.

Attention Efficiency: The Key to Scaling Transformers

At the heart of the transformer architecture is the self-attention mechanism, which allows each token to attend to every other token in the sequence. While this gives the model maximum flexibility to capture both short- and long-range dependencies, it comes at a steep computational cost.

For a sequence of length n, storing the full attention matrix requires O(n^2) memory. This quadratic dependency means that doubling the sequence length requires 4 times as much memory. In practice, this limits standard transformers to short sequences of just 512 tokens (Figure 1).

Figure 1: Memory usage of full attention scales quadratically with sequence length, limiting standard transformers to 512 tokens on current GPUs. Source: Author.

Several approaches have been proposed to improve the efficiency of transformers, including sparse attention mechanisms like Sparse Transformer (Child et al., 2019), Longformer (Beltagy et al., 2020), and Reformer (Kitaev et al., 2020). The key idea is to selectively attend to a subset of the tokens, reducing the number of attention computations from O(n^2) to O(n).

BigBird builds on these ideas, using a novel block sparse attention mechanism that combines local, global, and random attention patterns (Zaheer et al., 2020). This allows it to efficiently capture both short- and long-range interactions while scaling linearly with sequence length.

Inside BigBird‘s Attention Mechanism

BigBird‘s attention matrix is composed of three main components (Figure 2):

1. Window attention: Each token attends to a fixed-size local window around its position. This captures short-range dependencies.

2. Global attention: A small set of tokens is selected to attend to all other tokens and be attended to by all other tokens. This allows global information to be shared across the sequence.

3. Random attention: Each token attends to a random set of tokens in the sequence. This allows information to flow between distant parts of the sequence.

Figure 2: BigBird‘s block sparse attention matrix combines local window attention (green), global attention (blue), and random attention (red). Source: Author, based on Zaheer et al., 2020.

Mathematically, BigBird‘s attention matrix can be written as:

$$A{BigBird} = A{window} + A{global} + A{random}$$

where $A{window}$, $A{global}$, and $A_{random}$ are binary masks indicating which token pairs should attend to each other in the local window, global, and random attention patterns, respectively.

The key insight is that by carefully choosing the sparsity pattern, BigBird can capture the most important interactions between tokens while keeping the number of attention computations linear in the sequence length. The window size, number of global tokens, and number of random tokens per block are hyperparameters that can be tuned based on the task and available compute.

Compared to a full attention matrix, BigBird reduces the number of attention parameters from $n^2$ to $8n$ for sequences of length 4096 (Zaheer et al., 2020). This allows it to scale to much longer sequences while maintaining reasonable memory usage (Table 1).

Table 1: Memory usage of full attention vs. BigBird for different sequence lengths. BigBird scales linearly, allowing it to handle sequences up to 4096 tokens on current GPUs. Source: Zaheer et al., 2020.

Strong Results on Long-Context Language Tasks

To evaluate the effectiveness of its sparse attention mechanism, BigBird was tested on a range of challenging language tasks requiring long-range context:

• Long document question answering on the Natural Questions and TriviaQA datasets
• Document summarization on the arXiv and PubMed datasets
• Long document classification on the Hyperpartisan News Detection dataset

On all of these tasks, BigBird achieved substantial improvements over strong transformer baselines (Table 2). For example, on Natural Questions with 4096 token contexts, BigBird achieved an exact match score of 57.0, outperforming BERT-base by an absolute 5.6% (Zaheer et al., 2020).

Table 2: Comparison of BigBird and other transformer models on long-context language benchmarks. NQ Long = Natural Questions with long contexts, HND = Hyperpartisan News Detection. Source: Zaheer et al., 2020.

These strong empirical results suggest that BigBird‘s sparse attention mechanism is able to effectively capture the most salient interactions over long sequences, matching or exceeding the performance of full attention transformers. The ability to handle 4096 token contexts allows BigBird to extract more signal from long documents, an important capability for many real-world language understanding applications.

Enabling New Applications

BigBird‘s extended context length opens up exciting opportunities to tackle new language problems that were previously out of reach for transformers. Some potential high-impact application areas include:

• Biomedical text mining: Scientific papers, clinical notes, and electronic health records can span many pages. BigBird could help extract valuable knowledge and insights from these long, unstructured documents.

• Document-level machine translation: Translating long articles, books, and web pages requires maintaining coherence and consistency across many sentences. BigBird‘s ability to track long-range dependencies could improve the fluency of document-level translations.

• Conversational AI: Engaging in open-ended dialog often requires tracking context over long chat histories. BigBird could enable chatbots that can carry on in-depth, coherent conversations.

• Long-form question answering: Answering complex questions may require reasoning over multiple parts of a long document and aggregating information from different sections. BigBird could power QA systems that can handle rich, multi-hop queries.

• Full-text information retrieval: Instead of relying on keyword matching, BigBird could be used to encode the full text of documents and perform semantic search over large corpora. This could unlock more precise and relevant search results.

Beyond these applications, BigBird‘s efficient attention mechanism could be used as a general-purpose building block for models that need to process very long sequences. For example, BigBird could be used as the encoder in sequence-to-sequence models for tasks like abstractive summarization and document-level machine translation.

BigBird also points the way towards future research on efficient transformers that can scale to even longer contexts. Promising directions include dynamic and learned sparsity patterns, content-based sparse attention, and hybrid local-global attention mechanisms. As transformers continue to get bigger and more computationally intensive, these techniques will be key to keeping them tractable and applicable to real-world problems.

Implementing BigBird

BigBird and its sparse attention mechanism have been open-sourced by Google Research, making it easy to get started with training and fine-tuning your own long-context language models.

The original TensorFlow implementation is available on GitHub at google-research/bigbird. This repository includes the core BigBird model along with pre-training and fine-tuning code for tasks like question answering and summarization.

BigBird has also been integrated into the popular Hugging Face Transformers library, which provides a unified API for working with state-of-the-art NLP models. You can load a pre-trained BigBird model with just a few lines of code:

from transformers import BigBirdModel, BigBirdTokenizer

model_name = ‘google/bigbird-roberta-base‘
tokenizer = BigBirdTokenizer.from_pretrained(model_name)
model = BigBirdModel.from_pretrained(model_name)

input_ids = tokenizer(‘This is a long paragraph to encode.‘, return_tensors=‘pt‘)[‘input_ids‘]
sequence_outputs = model(input_ids)[0]  # retrieve the final layer‘s representation

This will load the bigbird-roberta-base model, which has been pre-trained on a large corpus of web pages and books. You can then fine-tune this model on your own long-context datasets using the standard Hugging Face Trainer class.

For more advanced use cases, you can also train a BigBird model from scratch using the BigBirdForPreTraining class. This allows you to control the vocabulary size, number of attention heads, hidden dimension, and other architecture hyperparameters. You can also experiment with different pre-training objectives like masked language modeling, next sentence prediction, and sentence order prediction.

The Frontier Ahead

As the amount of natural language data continues to grow exponentially, the ability to efficiently process long documents and contexts will become increasingly important. BigBird represents an important step towards language models that can extract knowledge and insights from large, unstructured text corpora.

However, much work remains to be done to realize the full potential of sparse attention mechanisms and long-context language understanding. Some key challenges and directions for future research include:

• Dynamic sparsity patterns: BigBird uses a fixed attention pattern, but the optimal sparsity structure may depend on the content and task. Models that can learn and adapt their attention patterns on the fly could be even more efficient and effective.

• Controlling attention span: For some tasks, attending to the full context may not be necessary or even desirable. Models that can dynamically adjust their attention span based on the query and content could be more interpretable and controllable.

• Knowledge-enhanced language models: Integrating structured knowledge from knowledge bases and knowledge graphs could help language models reason more effectively over long contexts. Techniques like entity linking, fact injection, and knowledge-guided attention could help ground language understanding in real-world knowledge.

• Multi-scale representation learning: Text exhibits structure at multiple scales, from individual words to sentences, paragraphs, and sections. Models that can learn hierarchical representations at different granularities could capture richer semantics and structures in long documents.

• Efficient training and inference: While BigBird can handle much longer sequences than standard transformers, training and inference are still expensive for very long contexts. Techniques like progressive attention, sparse backpropagation, and memory-efficient attention could help scale up to even longer sequences.

Tackling these challenges will require novel model architectures, training techniques, and benchmarks. But the potential payoff is immense: language models that can truly understand and reason over the vast amount of unstructured knowledge contained in books, articles, and web pages.

BigBird is an exciting step in this direction, demonstrating the potential of sparse attention mechanisms to expand the capabilities of language models. As the research community continues to push the frontiers of efficient transformers and long-context language understanding, we can expect to see rapid progress towards more knowledgeable, insightful, and intelligent language systems.

Acknowledgments

The author would like to thank the Google Research team for open-sourcing BigBird and its pre-trained models. Gratitude also goes to the Hugging Face team for integrating BigBird into their wonderful Transformers library.

References

• Zaheer, M., Guruganesh, G., Dubey, A., Ainslie, J., Alberti, C., Ontanon, S., … & Ahmed, A. (2020). Big bird: Transformers for longer sequences. arXiv preprint arXiv:2007.14062.

• Child, R., Gray, S., Radford, A., & Sutskever, I. (2019). Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509.

• Beltagy, I., Peters, M. E., & Cohan, A. (2020). Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150.

• Kitaev, N., Kaiser, Ł., & Levskaya, A. (2020). Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451.