Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Global Tokens › Implementation Strategies

Michael BrenndoerferUpdated June 23, 20256 min read

Language AI HandbookAll chapters

Part 17 of 67Efficient Attention

Chapter 120Global Tokens › Implementation Strategies

Reading section 9 of 15 · Implementation Strategies · 6 minSection 9 of 15 · 6 min

Section 9 of 15 · 6 min read

Global Tokens › Implementation StrategiesOne section · Sections in this chapter9 / 15 ↓

Sections in this chapter15

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

Implementing global tokens efficiently requires careful attention to the attention computation itself. The naive approach computes global and local attention separately, then merges results. A more efficient approach uses a single sparse attention kernel that handles both patterns in one pass.

The reason efficiency matters here is not just academic. The whole point of global tokens is to add minimal overhead to an already-efficient sliding window attention. If the implementation of global attention adds significant constant-factor overhead, it can undermine the efficiency gains of the sparse pattern. Production implementations in frameworks like HuggingFace Transformers handle this by computing global attention using dense matrix operations for the small set of global positions, while using specialized CUDA kernels for the local window attention for the remaining positions. The two results are then combined.

For a pure-Python reference implementation, the dense approach is simpler and still instructive. The key step is applying the attention mask before softmax: setting non-attended positions to negative infinity ensures that those positions contribute zero weight in the softmax output, effectively excluding them from the weighted average.

In[18]:

Code

def global_local_attention(
    query, key, value, window_size, global_positions, d_k
):
    """
    Compute global-local attention efficiently.

    Args:
        query, key, value: Arrays of shape (seq_len, d_model)
        window_size: Size of local attention window
        global_positions: List of positions with global attention
        d_k: Dimension for scaling

    Returns:
        output: Attended values of shape (seq_len, d_model)
        attention_weights: Sparse attention matrix
    """
    seq_len = query.shape[0]

    # Create attention mask
    mask = create_longformer_attention_mask(
        seq_len, window_size, global_positions
    )

    # Compute raw attention scores
    scores = query @ key.T / np.sqrt(d_k)

    # Apply mask: set non-attended positions to -inf
    masked_scores = np.where(mask, scores, -np.inf)

    # Softmax (row-wise)
    exp_scores = np.exp(
        masked_scores - masked_scores.max(axis=1, keepdims=True)
    )
    # Handle -inf: exp(-inf) = 0
    exp_scores = np.nan_to_num(exp_scores, nan=0.0, posinf=0.0, neginf=0.0)
    attention_weights = exp_scores / (
        exp_scores.sum(axis=1, keepdims=True) + 1e-9
    )

    # Compute output
    output = attention_weights @ value

    return output, attention_weights

The scoring step computes the full $n \times n$ matrix of dot products, then zeroes out the positions that do not appear in the mask. The formula is:

A_{ij} = \begin{cases} \dfrac{\mathbf{q}_i \cdot \mathbf{k}_j}{\sqrt{d_k}} & \text{if } \text{mask}[i, j] = 1 \\ -\infty & \text{otherwise} \end{cases}

where:

$A_{ij}$ : the attention score from query position $i$ to key position $j$
$\mathbf{q}_i$ : the query vector at position $i$
$\mathbf{k}_j$ : the key vector at position $j$
$d_k$ : the key dimensionality (used for scaling to prevent vanishing gradients)
$\text{mask}[i, j]$ : 1 if position $i$ is allowed to attend to position $j$ , 0 otherwise

Setting masked positions to $-\infty$ before the row-wise softmax is the standard trick: since $e^{-\infty} = 0$ , those positions contribute zero to the softmax denominator and produce zero attention weight. The numerical stability step of subtracting the row maximum before exponentiation (the log-sum-exp trick) prevents overflow for the non-masked positions.

The mask determines which attention scores to compute. Setting masked positions to negative infinity before softmax effectively zeros their contribution to the weighted average. The global tokens see the full sequence in their rows, while local tokens see only their windows plus global positions.

Let us verify the implementation with a small example:

In[19]:

Code

# Test the implementation
seq_len = 16
d_model = 64
d_k = 64
window_size = 4
global_positions = [0, 8]  # CLS and mid-sequence global token

# Random queries, keys, values
Q = np.random.randn(seq_len, d_model)
K = np.random.randn(seq_len, d_model)
V = np.random.randn(seq_len, d_model)

output, weights = global_local_attention(
    Q, K, V, window_size, global_positions, d_k
)

Out[20]:

Console

Input shape: (16, 64)
Output shape: (16, 64)
Attention weights shape: (16, 16)

Active connections per position:
  Position  0 (global): 16 connections
  Position  1 (local ):  5 connections
  Position  2 (local ):  6 connections
  Position  3 (local ):  7 connections
  Position  4 (local ):  7 connections
  Position  5 (local ):  7 connections
  Position  6 (local ):  6 connections
  Position  7 (local ):  6 connections
  Position  8 (global): 16 connections
  Position  9 (local ):  6 connections
  Position 10 (local ):  6 connections
  Position 11 (local ):  7 connections
  Position 12 (local ):  7 connections
  Position 13 (local ):  7 connections
  Position 14 (local ):  6 connections
  Position 15 (local ):  5 connections

Global positions connect to all 16 tokens. Local positions connect to their window (4-5 tokens) plus the 2 global tokens, totaling around 6-7 connections each.

Notice the asymmetry in connection counts. Global tokens at positions 0 and 8 have 16 connections each, one per token. Local tokens have 6-7 connections. This asymmetry is intentional and algorithmically motivated: the global tokens need to see everything to serve as effective relays, while local tokens only need to see their neighborhood and the global hubs. The total number of connections is much smaller than the $16 \times 16 = 256$ required by full attention.

Out[21]:

Visualization

Heatmap of attention weights showing dense rows for global tokens and sparse diagonal patterns for local tokens. — Attention weights from global-local attention on a 16-token sequence with global tokens at positions 0 and 8 (window size 4). Global token rows show distributed attention across all positions. Local token rows show concentrated attention within their windows, with additional connections visible at columns 0 and 8 where they attend to the global tokens. The red dashed lines mark the global token positions.

The attention weight heatmap shows the characteristic pattern. The global token rows (0 and 8) have distributed attention across all positions. The local token rows show concentrated attention within their windows, with visible bumps at columns 0 and 8 where they attend to global tokens.

In practice, the attention weights in the global rows are not uniform. The global token uses the standard scaled dot-product attention mechanism to decide how much to attend to each position, so some positions receive more attention than others based on their semantic relevance. A global CLS token does not treat every word equally: it attends more to content words and salient phrases than to function words and punctuation. This learned non-uniformity is what makes global tokens more powerful than a simple average pooling operation would be.