Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

BigBird › Implementing BigBird Attention

Michael BrenndoerferUpdated June 24, 20256 min read

Language AI HandbookAll chapters

Part 17 of 67Efficient Attention

Chapter 122BigBird › Implementing BigBird Attention

Reading section 4 of 12 · Implementing BigBird Attention · 6 minSection 4 of 12 · 6 min

Section 4 of 12 · 6 min read

BigBird › Implementing BigBird AttentionOne section · Sections in this chapter4 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Now that we understand the three components and their theoretical justification, let's implement a complete BigBird attention layer. This implementation will help solidify the concepts and show how the sparse pattern translates into actual computation.

The implementation follows the standard attention pipeline with one key modification: we apply a sparse mask before the softmax to zero out attention scores for non-attended positions. This mask encodes our three-component pattern: local windows, global tokens, and random connections.

In[11]:

Code

def softmax(x, axis=-1):
    """Numerically stable softmax."""
    exp_x = np.exp(x - x.max(axis=axis, keepdims=True))
    return exp_x / exp_x.sum(axis=axis, keepdims=True)


class BigBirdAttention:
    """
    BigBird sparse attention implementation.

    Combines sliding window, global tokens, and random attention
    for O(n) complexity with strong theoretical guarantees.
    """

    def __init__(
        self, d_model, window_size=64, num_global=2, num_random=3, seed=None
    ):
        """
        Initialize BigBird attention.

        Args:
            d_model: Model dimension
            window_size: One-sided sliding window size
            num_global: Number of global tokens at sequence start
            num_random: Number of random attention connections per token
            seed: Random seed for reproducibility
        """
        self.d_model = d_model
        self.window_size = window_size
        self.num_global = num_global
        self.num_random = num_random
        self.seed = seed

        self.scale = 1.0 / np.sqrt(d_model)

        # Initialize projections with Xavier initialization
        self.W_Q = np.random.randn(d_model, d_model) * np.sqrt(
            2.0 / (2 * d_model)
        )
        self.W_K = np.random.randn(d_model, d_model) * np.sqrt(
            2.0 / (2 * d_model)
        )
        self.W_V = np.random.randn(d_model, d_model) * np.sqrt(
            2.0 / (2 * d_model)
        )
        self.W_O = np.random.randn(d_model, d_model) * np.sqrt(
            2.0 / (2 * d_model)
        )

    def _create_attention_mask(self, seq_len):
        """Create the BigBird attention pattern."""
        return create_bigbird_attention_pattern(
            seq_len,
            self.window_size,
            self.num_global,
            self.num_random,
            seed=self.seed,
        )

    def __call__(self, X):
        """
        Apply BigBird attention.

        Args:
            X: Input tensor of shape (seq_len, d_model)

        Returns:
            output: Attended representations (seq_len, d_model)
            attention_weights: Sparse attention weights (seq_len, seq_len)
        """
        seq_len = X.shape[0]

        # Project to Q, K, V
        Q = X @ self.W_Q
        K = X @ self.W_K
        V = X @ self.W_V

        # Compute full attention scores
        scores = (Q @ K.T) * self.scale

        # Apply BigBird mask
        mask = self._create_attention_mask(seq_len)
        masked_scores = np.where(mask > 0, scores, -1e9)

        # Softmax and aggregate
        attention_weights = softmax(masked_scores, axis=-1)

        # Zero out weights for masked positions
        attention_weights = attention_weights * mask

        # Compute output
        output = attention_weights @ V
        output = output @ self.W_O

        return output, attention_weights

In[12]:

Code

# Test the BigBird attention implementation
d_model = 64
seq_len = 256

# Create input sequence
X = np.random.randn(seq_len, d_model)

# Initialize BigBird attention
bigbird = BigBirdAttention(
    d_model=d_model, window_size=8, num_global=4, num_random=5, seed=42
)

# Apply attention
output, weights = bigbird(X)

Out[13]:

Console

BigBird attention test:
  Input shape: (256, 64)
  Output shape: (256, 64)
  Attention weights shape: (256, 256)

Attention statistics:
  Non-zero weights per row: 29.3
  Sparsity: 88.5%

Expected connections per token: 2*8 + 1 + 4 + 5 = 26

The output shape matches the input, confirming that BigBird attention preserves sequence dimensions. Each token attends to approximately 26 positions on average, which aligns with our theoretical expectation of $2w + 1 + g + r = 17 + 4 + 5 = 26$ . The 89.7% sparsity means we're computing less than 11% of the attention scores that full attention would require.

Out[14]:

Visualization

Bar chart showing regular tokens clustered around 26 attended positions and a separate four-token bar at 256 for global tokens. — Distribution of attention connections per token position. Global tokens (positions 0-3) attend to all 256 positions, creating the separate bar at the right. Regular tokens attend to a consistent number of positions determined by window size, global tokens, and random connections. The narrow distribution for regular tokens reflects the bounded fan-out that ensures linear complexity.

The histogram reveals the bimodal nature of BigBird's connection pattern. Regular tokens cluster tightly around the expected value of $2w + 1 + g + r$ . This reflects the consistent sparse pattern. The few global tokens (positions 0-3) attend to all positions, appearing at the far right. This structure ensures that most of the computation is bounded while still maintaining global connectivity through the hub tokens.

The implementation shows that each token attends to roughly $2w + g + r$ other tokens, independent of sequence length. More precisely, for a token at position $i$ (excluding the first $g$ global tokens), the number of attended positions is:

|\text{Attend}(i)| = \underbrace{(2w + 1)}_{\text{local window}} + \underbrace{g}_{\text{global tokens}} + \underbrace{r}_{\text{random connections}}

where some positions may overlap (e.g., if a random connection happens to fall within the local window or on a global token). The global tokens themselves attend to all $n$ positions, but since there are only $g$ of them, this adds $O(g \cdot n) = O(n)$ total connections. This bounded fan-out per token is the source of BigBird's linear complexity.