ALiBi: Attention with Linear Biases — Press 2021's Simple Solution and Extrapolation Advantage

import torch
import math
 
def get_alibi_slopes(n_heads):
    """Geometric sequence of slopes for ALiBi."""
    def get_slopes_power_of_2(n):
        start = 2 ** (-(2 ** -(math.log2(n) - 3)))
        ratio = start
        return [start * ratio ** i for i in range(n)]
    
    if math.log2(n_heads).is_integer():
        return get_slopes_power_of_2(n_heads)
    else:
        # Handle non-power-of-2 heads (BLOOM extension)
        closest_power = 2 ** math.floor(math.log2(n_heads))
        slopes = get_slopes_power_of_2(closest_power)
        slopes += get_slopes_power_of_2(2 * closest_power)[0::2][:n_heads - closest_power]
        return slopes
 
 
def get_alibi_bias(n_heads, seq_len, device='cpu'):
    """Build ALiBi bias matrix."""
    slopes = torch.tensor(get_alibi_slopes(n_heads), device=device)
    # Distance matrix: bias[i,j] = j - i (negative for past)
    distances = torch.arange(seq_len, device=device).unsqueeze(0) - \
                torch.arange(seq_len, device=device).unsqueeze(1)
    # [n_heads, seq, seq]
    bias = slopes.view(-1, 1, 1) * distances.unsqueeze(0).float()
    return bias
 
 
class ALiBiAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.qkv_proj = torch.nn.Linear(d_model, 3 * d_model)
        self.out_proj = torch.nn.Linear(d_model, d_model)
    
    def forward(self, x, alibi_bias=None):
        batch, seq, _ = x.shape
        qkv = self.qkv_proj(x).view(batch, seq, 3, self.n_heads, self.d_head)
        Q, K, V = qkv.unbind(2)
        Q, K, V = Q.transpose(1, 2), K.transpose(1, 2), V.transpose(1, 2)
        
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_head)
        
        # Add ALiBi bias
        if alibi_bias is not None:
            scores = scores + alibi_bias
        
        # Causal mask
        mask = torch.triu(torch.ones(seq, seq), diagonal=1).bool().to(x.device)
        scores = scores.masked_fill(mask, float('-inf'))
        
        weights = torch.softmax(scores, dim=-1)
        out = weights @ V
        out = out.transpose(1, 2).contiguous().view(batch, seq, -1)
        return self.out_proj(out)
 
 
# Usage
model = ALiBiAttention(d_model=4096, n_heads=32)
alibi_bias = get_alibi_bias(32, 2048)  # [32, 2048, 2048]
print(f"Slopes (first 4): {get_alibi_slopes(32)[:4]}")
print(f"ALiBi bias shape: {alibi_bias.shape}")