Sparse Attention（Longformer・BigBird）の理論と実装 — 長系列を効率的に処理する

import numpy as np
import matplotlib.pyplot as plt

# Full Attentionの注意行列を可視化
np.random.seed(42)
n = 32  # 系列長（小さな例で示す）
d_k = 8

# ランダムなQuery, Keyを生成
Q = np.random.randn(n, d_k)
K = np.random.randn(n, d_k)

# 注意スコアの計算
scores = Q @ K.T / np.sqrt(d_k)

# softmaxの適用
exp_scores = np.exp(scores - np.max(scores, axis=1, keepdims=True))
attention = exp_scores / np.sum(exp_scores, axis=1, keepdims=True)

fig, axes = plt.subplots(1, 2, figsize=(14, 6))

# 注意行列（密行列）
im0 = axes[0].imshow(attention, cmap='viridis', aspect='equal')
axes[0].set_title('Full Attention Matrix (Dense)', fontsize=14)
axes[0].set_xlabel('Key position')
axes[0].set_ylabel('Query position')
plt.colorbar(im0, ax=axes[0], fraction=0.046)

# マスクパターン（全て1 = 全接続）
mask = np.ones((n, n))
axes[1].imshow(mask, cmap='Blues', aspect='equal', vmin=0, vmax=1)
axes[1].set_title('Full Attention Mask (All Connected)', fontsize=14)
axes[1].set_xlabel('Key position')
axes[1].set_ylabel('Query position')

plt.tight_layout()
plt.savefig('full_attention_matrix.png', dpi=150, bbox_inches='tight')
plt.show()

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

def create_local_mask(seq_len, window_size):
    """
    Local Attention（スライディングウィンドウ）のマスクを生成する。
    各トークンは前後 window_size // 2 のトークンに注目する。
    """
    mask = torch.zeros(seq_len, seq_len)
    half_w = window_size // 2
    for i in range(seq_len):
        start = max(0, i - half_w)
        end = min(seq_len, i + half_w + 1)
        mask[i, start:end] = 1.0
    return mask

def create_dilated_mask(seq_len, window_size, dilation):
    """
    Dilated Attention（拡張スライディングウィンドウ）のマスクを生成する。
    dilation間隔でサンプリングしながら window_size 個のトークンに注目する。
    """
    mask = torch.zeros(seq_len, seq_len)
    half_w = window_size // 2
    for i in range(seq_len):
        for k in range(-half_w, half_w + 1):
            j = i + k * dilation
            if 0 <= j < seq_len:
                mask[i, j] = 1.0
    return mask

def create_global_mask(seq_len, global_indices):
    """
    Global Attention のマスクを生成する。
    global_indices で指定された位置のトークンは全トークンに注目し、
    全トークンからも注目される。
    """
    mask = torch.zeros(seq_len, seq_len)
    for g in global_indices:
        mask[g, :] = 1.0  # グローバルトークン → 全トークン
        mask[:, g] = 1.0  # 全トークン → グローバルトークン
    return mask

def create_random_mask(seq_len, num_random, seed=42):
    """
    Random Attention のマスクを生成する。
    各トークンがランダムに num_random 個のトークンに注目する。
    """
    rng = np.random.RandomState(seed)
    mask = torch.zeros(seq_len, seq_len)
    for i in range(seq_len):
        # 自分以外からランダムにサンプル
        candidates = list(range(seq_len))
        candidates.remove(i)
        chosen = rng.choice(candidates, size=min(num_random, len(candidates)), replace=False)
        mask[i, chosen] = 1.0
        mask[i, i] = 1.0  # 自分自身は常に含める
    return mask

# 動作確認: 各マスクのサイズを表示
seq_len = 32
print(f"Local mask:   {create_local_mask(seq_len, window_size=5).sum().item():.0f} connections")
print(f"Dilated mask: {create_dilated_mask(seq_len, window_size=5, dilation=2).sum().item():.0f} connections")
print(f"Global mask:  {create_global_mask(seq_len, global_indices=[0]).sum().item():.0f} connections")
print(f"Random mask:  {create_random_mask(seq_len, num_random=3).sum().item():.0f} connections")
print(f"Full mask:    {(seq_len * seq_len):.0f} connections")

import numpy as np
import torch
import matplotlib.pyplot as plt

# 前のセクションの関数（create_local_mask 等）を使用

seq_len = 64

# 各マスクを生成
local_mask = create_local_mask(seq_len, window_size=7)
dilated_mask = create_dilated_mask(seq_len, window_size=7, dilation=3)
global_mask = create_global_mask(seq_len, global_indices=[0, 1])
random_mask = create_random_mask(seq_len, num_random=4)

fig, axes = plt.subplots(2, 3, figsize=(18, 12))

# (a) Full Attention
full_mask = torch.ones(seq_len, seq_len)
axes[0, 0].imshow(full_mask.numpy(), cmap='Blues', aspect='equal', vmin=0, vmax=1)
axes[0, 0].set_title(f'(a) Full Attention\n({int(full_mask.sum())} connections)', fontsize=12)
axes[0, 0].set_xlabel('Key position')
axes[0, 0].set_ylabel('Query position')

# (b) Local Attention
axes[0, 1].imshow(local_mask.numpy(), cmap='Blues', aspect='equal', vmin=0, vmax=1)
axes[0, 1].set_title(f'(b) Local Attention (w=7)\n({int(local_mask.sum())} connections)', fontsize=12)
axes[0, 1].set_xlabel('Key position')
axes[0, 1].set_ylabel('Query position')

# (c) Dilated Attention
axes[0, 2].imshow(dilated_mask.numpy(), cmap='Blues', aspect='equal', vmin=0, vmax=1)
axes[0, 2].set_title(f'(c) Dilated Attention (w=7, d=3)\n({int(dilated_mask.sum())} connections)', fontsize=12)
axes[0, 2].set_xlabel('Key position')
axes[0, 2].set_ylabel('Query position')

# (d) Global Attention
axes[1, 0].imshow(global_mask.numpy(), cmap='Blues', aspect='equal', vmin=0, vmax=1)
axes[1, 0].set_title(f'(d) Global Attention (g=2)\n({int(global_mask.sum())} connections)', fontsize=12)
axes[1, 0].set_xlabel('Key position')
axes[1, 0].set_ylabel('Query position')

# (e) Random Attention
axes[1, 1].imshow(random_mask.numpy(), cmap='Blues', aspect='equal', vmin=0, vmax=1)
axes[1, 1].set_title(f'(e) Random Attention (r=4)\n({int(random_mask.sum())} connections)', fontsize=12)
axes[1, 1].set_xlabel('Key position')
axes[1, 1].set_ylabel('Query position')

# (f) BigBird (Local + Global + Random)
bigbird_mask = torch.clamp(local_mask + global_mask + random_mask, 0, 1)
axes[1, 2].imshow(bigbird_mask.numpy(), cmap='Blues', aspect='equal', vmin=0, vmax=1)
axes[1, 2].set_title(f'(f) BigBird Combined\n({int(bigbird_mask.sum())} connections)', fontsize=12)
axes[1, 2].set_xlabel('Key position')
axes[1, 2].set_ylabel('Query position')

plt.tight_layout()
plt.savefig('attention_patterns.png', dpi=150, bbox_inches='tight')
plt.show()

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class LongformerSparseAttention(nn.Module):
    """
    Longformer風のSparse Attention。
    Local Attention（スライディングウィンドウ）+ Global Attention を組み合わせる。
    """
    def __init__(self, d_model, num_heads, window_size, num_global_tokens=1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.window_size = window_size
        self.num_global_tokens = num_global_tokens

        # 通常トークン用のQ, K, V射影
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)

        # グローバルトークン用の別のQ, K, V射影（Longformerの特徴）
        self.W_Q_global = nn.Linear(d_model, d_model)
        self.W_K_global = nn.Linear(d_model, d_model)
        self.W_V_global = nn.Linear(d_model, d_model)

        # 出力射影
        self.W_O = nn.Linear(d_model, d_model)

    def _create_combined_mask(self, seq_len, device):
        """Local + Global の統合マスクを生成する"""
        # Local Attention マスク
        mask = torch.zeros(seq_len, seq_len, device=device)
        half_w = self.window_size // 2
        for i in range(seq_len):
            start = max(0, i - half_w)
            end = min(seq_len, i + half_w + 1)
            mask[i, start:end] = 1.0

        # Global Attention マスク（先頭 num_global_tokens 個）
        g = self.num_global_tokens
        mask[:g, :] = 1.0  # グローバルトークン → 全トークン
        mask[:, :g] = 1.0  # 全トークン → グローバルトークン

        return mask

    def forward(self, x):
        """
        x: (batch_size, seq_len, d_model)
        returns: (batch_size, seq_len, d_model)
        """
        B, N, D = x.shape
        g = self.num_global_tokens

        # マスクの生成
        mask = self._create_combined_mask(N, x.device)  # (N, N)

        # 通常トークンのQ, K, V
        Q = self.W_Q(x)  # (B, N, D)
        K = self.W_K(x)
        V = self.W_V(x)

        # グローバルトークンは別の射影を使用
        Q_global = self.W_Q_global(x[:, :g, :])  # (B, g, D)
        K_global = self.W_K_global(x[:, :g, :])
        V_global = self.W_V_global(x[:, :g, :])

        # グローバルトークンの射影で上書き
        Q = Q.clone()
        K = K.clone()
        V = V.clone()
        Q[:, :g, :] = Q_global
        K[:, :g, :] = K_global
        V[:, :g, :] = V_global

        # Multi-Head 形状に変換: (B, num_heads, N, d_k)
        Q = Q.view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(B, N, self.num_heads, self.d_k).transpose(1, 2)

        # Scaled Dot-Product Attention with Sparse Mask
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)  # (B, H, N, N)

        # マスクの適用: mask=0 の位置を -inf にする
        mask_expanded = mask.unsqueeze(0).unsqueeze(0)  # (1, 1, N, N)
        scores = scores.masked_fill(mask_expanded == 0, float('-inf'))

        attn_weights = F.softmax(scores, dim=-1)  # (B, H, N, N)
        attn_weights = attn_weights.masked_fill(mask_expanded == 0, 0.0)

        # 重み付き和
        output = torch.matmul(attn_weights, V)  # (B, H, N, d_k)

        # ヘッドの結合
        output = output.transpose(1, 2).contiguous().view(B, N, D)
        output = self.W_O(output)

        return output, attn_weights

# 動作確認
torch.manual_seed(42)
d_model, num_heads, window_size = 64, 4, 5
model = LongformerSparseAttention(d_model, num_heads, window_size, num_global_tokens=2)

x = torch.randn(2, 32, d_model)  # バッチ2、系列長32
output, attn_weights = model(x)
print(f"Input shape:  {x.shape}")
print(f"Output shape: {output.shape}")
print(f"Attention weights shape: {attn_weights.shape}")
print(f"Non-zero attention entries: {(attn_weights[0, 0] > 0).sum().item()} / {32*32}")

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class BigBirdSparseAttention(nn.Module):
    """
    BigBird風のSparse Attention。
    Local + Global + Random の3パターンを組み合わせる。
    """
    def __init__(self, d_model, num_heads, window_size,
                 num_global_tokens=1, num_random=3, random_seed=42):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.window_size = window_size
        self.num_global_tokens = num_global_tokens
        self.num_random = num_random
        self.random_seed = random_seed

        # Q, K, V射影（BigBird-ITCでは単一の射影を使用）
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)

    def _create_bigbird_mask(self, seq_len, device):
        """Local + Global + Random の統合マスクを生成"""
        mask = torch.zeros(seq_len, seq_len, device=device)

        # 1. Local Attention
        half_w = self.window_size // 2
        for i in range(seq_len):
            start = max(0, i - half_w)
            end = min(seq_len, i + half_w + 1)
            mask[i, start:end] = 1.0

        # 2. Global Attention（先頭 g 個のトークン）
        g = self.num_global_tokens
        mask[:g, :] = 1.0
        mask[:, :g] = 1.0

        # 3. Random Attention
        rng = np.random.RandomState(self.random_seed)
        for i in range(seq_len):
            candidates = [j for j in range(seq_len) if mask[i, j] == 0]
            if len(candidates) > 0:
                num_to_select = min(self.num_random, len(candidates))
                chosen = rng.choice(candidates, size=num_to_select, replace=False)
                mask[i, chosen] = 1.0

        return mask

    def forward(self, x):
        """
        x: (batch_size, seq_len, d_model)
        returns: (batch_size, seq_len, d_model)
        """
        B, N, D = x.shape

        # マスク生成
        mask = self._create_bigbird_mask(N, x.device)

        # Q, K, V の計算
        Q = self.W_Q(x).view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_K(x).view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_V(x).view(B, N, self.num_heads, self.d_k).transpose(1, 2)

        # Scaled Dot-Product Attention with Sparse Mask
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)

        mask_expanded = mask.unsqueeze(0).unsqueeze(0)
        scores = scores.masked_fill(mask_expanded == 0, float('-inf'))

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = attn_weights.masked_fill(mask_expanded == 0, 0.0)

        output = torch.matmul(attn_weights, V)
        output = output.transpose(1, 2).contiguous().view(B, N, D)
        output = self.W_O(output)

        return output, attn_weights

# 動作確認
torch.manual_seed(42)
d_model, num_heads, window_size = 64, 4, 5
model = BigBirdSparseAttention(
    d_model, num_heads, window_size,
    num_global_tokens=2, num_random=3
)

x = torch.randn(2, 32, d_model)
output, attn_weights = model(x)
print(f"Input shape:  {x.shape}")
print(f"Output shape: {output.shape}")
print(f"Attention weights shape: {attn_weights.shape}")
print(f"Non-zero attention entries: {(attn_weights[0, 0] > 0).sum().item()} / {32*32}")

import torch
import numpy as np
import matplotlib.pyplot as plt

# マスク生成関数（前のセクションと同じ）を使用
# ここではインラインで再定義して自己完結させる

def create_longformer_mask(seq_len, window_size, dilation, global_indices):
    """Longformer風の統合マスク（Local + Dilated + Global）"""
    mask = torch.zeros(seq_len, seq_len)
    half_w = window_size // 2

    # Local Attention
    for i in range(seq_len):
        start = max(0, i - half_w)
        end = min(seq_len, i + half_w + 1)
        mask[i, start:end] = 1.0

    # Dilated Attention
    for i in range(seq_len):
        for k in range(-half_w, half_w + 1):
            j = i + k * dilation
            if 0 <= j < seq_len:
                mask[i, j] = 1.0

    # Global Attention
    for g in global_indices:
        mask[g, :] = 1.0
        mask[:, g] = 1.0

    return mask

def create_bigbird_mask(seq_len, window_size, global_indices, num_random, seed=42):
    """BigBird風の統合マスク（Local + Global + Random）"""
    mask = torch.zeros(seq_len, seq_len)
    half_w = window_size // 2

    # Local Attention
    for i in range(seq_len):
        start = max(0, i - half_w)
        end = min(seq_len, i + half_w + 1)
        mask[i, start:end] = 1.0

    # Global Attention
    for g in global_indices:
        mask[g, :] = 1.0
        mask[:, g] = 1.0

    # Random Attention
    rng = np.random.RandomState(seed)
    for i in range(seq_len):
        candidates = [j for j in range(seq_len) if mask[i, j] == 0]
        if len(candidates) > 0:
            chosen = rng.choice(candidates,
                                size=min(num_random, len(candidates)),
                                replace=False)
            mask[i, chosen] = 1.0

    return mask

seq_len = 64
global_indices = [0, 1]

longformer_mask = create_longformer_mask(seq_len, window_size=7, dilation=3,
                                          global_indices=global_indices)
bigbird_mask = create_bigbird_mask(seq_len, window_size=7,
                                    global_indices=global_indices,
                                    num_random=5)

fig, axes = plt.subplots(1, 3, figsize=(18, 6))

# Full Attention
full_mask = torch.ones(seq_len, seq_len)
axes[0].imshow(full_mask.numpy(), cmap='Blues', aspect='equal', vmin=0, vmax=1)
axes[0].set_title(f'Full Attention\n({int(full_mask.sum())} connections)', fontsize=13)
axes[0].set_xlabel('Key position')
axes[0].set_ylabel('Query position')

# Longformer
axes[1].imshow(longformer_mask.numpy(), cmap='Blues', aspect='equal', vmin=0, vmax=1)
sparsity_lf = 1 - longformer_mask.sum().item() / (seq_len * seq_len)
axes[1].set_title(f'Longformer\n({int(longformer_mask.sum())} connections, '
                   f'sparsity={sparsity_lf:.1%})', fontsize=13)
axes[1].set_xlabel('Key position')
axes[1].set_ylabel('Query position')

# BigBird
axes[2].imshow(bigbird_mask.numpy(), cmap='Blues', aspect='equal', vmin=0, vmax=1)
sparsity_bb = 1 - bigbird_mask.sum().item() / (seq_len * seq_len)
axes[2].set_title(f'BigBird\n({int(bigbird_mask.sum())} connections, '
                   f'sparsity={sparsity_bb:.1%})', fontsize=13)
axes[2].set_xlabel('Key position')
axes[2].set_ylabel('Query position')

plt.tight_layout()
plt.savefig('longformer_vs_bigbird.png', dpi=150, bbox_inches='tight')
plt.show()

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import time
import matplotlib.pyplot as plt

class FullAttention(nn.Module):
    """標準的なFull Attention"""
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.W_O = nn.Linear(d_model, d_model)

    def forward(self, x):
        B, N, D = x.shape
        Q = self.W_Q(x).view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_K(x).view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_V(x).view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        attn = F.softmax(scores, dim=-1)
        output = torch.matmul(attn, V)
        output = output.transpose(1, 2).contiguous().view(B, N, D)
        return self.W_O(output)

def measure_performance(seq_lengths, d_model=64, num_heads=4, window_size=32,
                        num_trials=5):
    """各系列長でFull AttentionとSparse Attentionの計算時間を計測"""
    full_times = []
    sparse_times = []
    full_memory = []
    sparse_memory = []

    for n in seq_lengths:
        # Full Attention
        model_full = FullAttention(d_model, num_heads)
        x = torch.randn(1, n, d_model)

        # ウォームアップ
        with torch.no_grad():
            _ = model_full(x)

        # 計測
        times = []
        for _ in range(num_trials):
            start = time.perf_counter()
            with torch.no_grad():
                _ = model_full(x)
            times.append(time.perf_counter() - start)
        full_times.append(np.median(times))

        # メモリ（注意行列のサイズ）
        full_memory.append(n * n * num_heads * 4 / (1024 ** 2))  # MB

        # Sparse Attention（BigBird風）
        # 各トークンの接続数: window_size + num_global + num_random
        connections_per_token = min(window_size + 2 + 3, n)
        sparse_memory.append(n * connections_per_token * num_heads * 4 / (1024 ** 2))

        # BigBird風のSparse Attention計測
        model_sparse = BigBirdSparseAttention(
            d_model, num_heads, window_size,
            num_global_tokens=2, num_random=3
        )
        with torch.no_grad():
            _ = model_sparse(x)

        times = []
        for _ in range(num_trials):
            start = time.perf_counter()
            with torch.no_grad():
                output, _ = model_sparse(x)
            times.append(time.perf_counter() - start)
        sparse_times.append(np.median(times))

    return full_times, sparse_times, full_memory, sparse_memory

# 計測
seq_lengths = [64, 128, 256, 512, 1024, 2048]
full_times, sparse_times, full_mem, sparse_mem = measure_performance(seq_lengths)

fig, axes = plt.subplots(1, 2, figsize=(14, 6))

# 計算時間の比較
axes[0].plot(seq_lengths, full_times, 'o-', label='Full Attention', linewidth=2, markersize=8)
axes[0].plot(seq_lengths, sparse_times, 's-', label='Sparse Attention (BigBird)', linewidth=2, markersize=8)
axes[0].set_xlabel('Sequence Length', fontsize=12)
axes[0].set_ylabel('Time (seconds)', fontsize=12)
axes[0].set_title('Computation Time Comparison', fontsize=14)
axes[0].legend(fontsize=11)
axes[0].set_xscale('log', base=2)
axes[0].set_yscale('log')
axes[0].grid(True, alpha=0.3)

# メモリ使用量の比較（理論値）
axes[1].plot(seq_lengths, full_mem, 'o-', label='Full Attention ($O(n^2)$)', linewidth=2, markersize=8)
axes[1].plot(seq_lengths, sparse_mem, 's-', label='Sparse Attention ($O(n)$)', linewidth=2, markersize=8)
axes[1].set_xlabel('Sequence Length', fontsize=12)
axes[1].set_ylabel('Attention Matrix Memory (MB)', fontsize=12)
axes[1].set_title('Attention Matrix Memory Comparison', fontsize=14)
axes[1].legend(fontsize=11)
axes[1].set_xscale('log', base=2)
axes[1].set_yscale('log')
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('performance_comparison.png', dpi=150, bbox_inches='tight')
plt.show()

# 数値を表形式で出力
print(f"\n{'Seq Length':>10} | {'Full Time (ms)':>14} | {'Sparse Time (ms)':>16} | "
      f"{'Full Mem (MB)':>13} | {'Sparse Mem (MB)':>15} | {'Speedup':>8}")
print("-" * 95)
for i, n in enumerate(seq_lengths):
    speedup = full_times[i] / sparse_times[i] if sparse_times[i] > 0 else float('inf')
    print(f"{n:>10} | {full_times[i]*1000:>14.2f} | {sparse_times[i]*1000:>16.2f} | "
          f"{full_mem[i]:>13.3f} | {sparse_mem[i]:>15.3f} | {speedup:>7.2f}x")