【大模型】- 原生稀疏注意力

宋 算法
  2026-08-19 09:17 2026-08-19 09:17 创建   2026-08-21 16:20:11 2026-08-21 16:20:11 更新      2k 字  10 分钟

原生稀疏注意力

高效处理长序列的注意力机制

类型： 学习　｜　语言： Python　｜　🏷 前置：《差分注意力 v2》（本系列第 15 篇）

学习目标

  • 理解稀疏注意力的原理
  • 实现原生稀疏注意力机制
  • 优化长序列处理
  • 比较不同稀疏注意力方法
  • 分析稀疏注意力的权衡

稀疏注意力概述

稀疏注意力通过只计算部分 token 对之间的注意力，减少计算复杂度。

标准注意力的问题

  • 计算复杂度：O(n²)，其中 n 是序列长度
  • 内存使用：O(n²)，存储注意力矩阵
  • 长序列处理：难以处理长序列

稀疏注意力的优势

  • 降低复杂度：O(n√n) 或 O(n log n)
  • 减少内存：只存储部分注意力
  • 加速推理：减少计算量

原生稀疏注意力实现

滑动窗口注意力

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import torch
import torch.nn as nn
import torch.nn.functional as F

class SlidingWindowAttention(nn.Module):
    def __init__(self, dim, n_heads, window_size=256):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.window_size = window_size
        
        self.wq = nn.Linear(dim, dim, bias=False)
        self.wk = nn.Linear(dim, dim, bias=False)
        self.wv = nn.Linear(dim, dim, bias=False)
        self.wo = nn.Linear(dim, dim, bias=False)
    
    def forward(self, x):
        batch_size, seq_len, _ = x.shape
        
        q = self.wq(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
        k = self.wk(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
        v = self.wv(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
        
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        
        # 创建滑动窗口掩码
        mask = self._create_sliding_window_mask(seq_len, self.window_size)
        mask = mask.unsqueeze(0).unsqueeze(0)
        
        # 计算注意力
        attn = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn = attn.masked_fill(mask == 0, float("-inf"))
        attn = F.softmax(attn, dim=-1)
        
        output = torch.matmul(attn, v)
        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, seq_len, -1)
        
        return self.wo(output)
    
    def _create_sliding_window_mask(self, seq_len, window_size):
        """创建滑动窗口掩码"""
        mask = torch.zeros(seq_len, seq_len)
        
        for i in range(seq_len):
            start = max(0, i - window_size // 2)
            end = min(seq_len, i + window_size // 2 + 1)
            mask[i, start:end] = 1
        
        return mask

分块稀疏注意力

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class ChunkedSparseAttention(nn.Module):
    def __init__(self, dim, n_heads, chunk_size=64):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.chunk_size = chunk_size
        
        self.wq = nn.Linear(dim, dim, bias=False)
        self.wk = nn.Linear(dim, dim, bias=False)
        self.wv = nn.Linear(dim, dim, bias=False)
        self.wo = nn.Linear(dim, dim, bias=False)
        
        # 局部注意力
        self.local_attn = nn.MultiheadAttention(dim, n_heads, batch_first=True)
        
        # 全局注意力（每 chunk 的第一个 token）
        self.global_attn = nn.MultiheadAttention(dim, n_heads, batch_first=True)
    
    def forward(self, x):
        batch_size, seq_len, _ = x.shape
        
        # 分块
        n_chunks = seq_len // self.chunk_size
        
        # 局部注意力
        x_chunked = x.view(batch_size * n_chunks, self.chunk_size, -1)
        local_out, _ = self.local_attn(x_chunked, x_chunked, x_chunked)
        local_out = local_out.view(batch_size, seq_len, -1)
        
        # 全局注意力（每个 chunk 的第一个 token）
        global_tokens = x[:, ::self.chunk_size, :]  # (batch, n_chunks, dim)
        global_out, _ = self.global_attn(global_tokens, global_tokens, global_tokens)
        
        # 广播全局信息
        global_out = global_out.repeat_interleave(self.chunk_size, dim=1)
        global_out = global_out[:, :seq_len, :]
        
        # 融合
        output = local_out + global_out
        
        return self.wo(output)

稀疏全局注意力

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
class SparseGlobalAttention(nn.Module):
    def __init__(self, dim, n_heads, n_global_tokens=32):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.n_global_tokens = n_global_tokens
        
        self.wq = nn.Linear(dim, dim, bias=False)
        self.wk = nn.Linear(dim, dim, bias=False)
        self.wv = nn.Linear(dim, dim, bias=False)
        self.wo = nn.Linear(dim, dim, bias=False)
        
        # 可学习的全局 token
        self.global_tokens = nn.Parameter(torch.randn(1, n_global_tokens, dim))
    
    def forward(self, x):
        batch_size, seq_len, _ = x.shape
        
        # 扩展全局 token
        global_tokens = self.global_tokens.expand(batch_size, -1, -1)
        
        # 拼接
        x_cat = torch.cat([global_tokens, x], dim=1)
        
        q = self.wq(x_cat).view(batch_size, -1, self.n_heads, self.head_dim)
        k = self.wk(x_cat).view(batch_size, -1, self.n_heads, self.head_dim)
        v = self.wv(x_cat).view(batch_size, -1, self.n_heads, self.head_dim)
        
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        
        # 计算注意力（全局 token 关注所有，局部 token 只关注全局）
        attn = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        
        # 创建掩码：局部 token 只能关注全局 token
        mask = torch.zeros(seq_len + self.n_global_tokens, seq_len + self.n_global_tokens)
        mask[:self.n_global_tokens, :] = 1  # 全局 token 关注所有
        mask[self.n_global_tokens:, :self.n_global_tokens] = 1  # 局部 token 关注全局
        
        attn = attn.masked_fill(mask == 0, float("-inf"))
        attn = F.softmax(attn, dim=-1)
        
        output = torch.matmul(attn, v)
        
        # 只取局部 token 的输出
        output = output[:, :, self.n_global_tokens:, :]
        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, seq_len, -1)
        
        return self.wo(output)

高级稀疏模式

局部敏感哈希（LSH）注意力

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
class LSHAttention(nn.Module):
    def __init__(self, dim, n_heads, n_hashes=4):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.n_hashes = n_hashes
        
        self.wq = nn.Linear(dim, dim, bias=False)
        self.wk = nn.Linear(dim, dim, bias=False)
        self.wv = nn.Linear(dim, dim, bias=False)
        self.wo = nn.Linear(dim, dim, bias=False)
        
        # 哈希函数
        self.hash_vectors = nn.Parameter(torch.randn(n_heads, self.head_dim))
    
    def forward(self, x):
        batch_size, seq_len, _ = x.shape
        
        q = self.wq(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
        k = self.wk(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
        v = self.wv(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
        
        # 计算哈希
        q_hash = torch.matmul(q, self.hash_vectors.unsqueeze(0).unsqueeze(2))
        k_hash = torch.matmul(k, self.hash_vectors.unsqueeze(0).unsqueeze(2))
        
        # 排序
        q_sorted, q_indices = torch.sort(q_hash, dim=1)
        k_sorted, k_indices = torch.sort(k_hash, dim=1)
        
        # 只计算相同桶内的注意力
        output = torch.zeros_like(v)
        
        for i in range(seq_len):
            # 找到相同的桶
            same_bucket = (q_hash[:, i, :, :] == k_hash).any(dim=1)
            
            # 计算注意力
            attn = torch.matmul(q[:, i:i+1, :, :], k.transpose(-2, -1))
            attn = attn.masked_fill(~same_bucket.unsqueeze(1).unsqueeze(2), float("-inf"))
            attn = F.softmax(attn, dim=-1)
            
            out = torch.matmul(attn, v)
            output[:, i, :, :] = out[:, 0, :, :]
        
        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, seq_len, -1)
        
        return self.wo(output)

随机稀疏注意力

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class RandomSparseAttention(nn.Module):
    def __init__(self, dim, n_heads, sparsity=0.1):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.sparsity = sparsity
        
        self.wq = nn.Linear(dim, dim, bias=False)
        self.wk = nn.Linear(dim, dim, bias=False)
        self.wv = nn.Linear(dim, dim, bias=False)
        self.wo = nn.Linear(dim, dim, bias=False)
    
    def forward(self, x):
        batch_size, seq_len, _ = x.shape
        
        q = self.wq(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
        k = self.wk(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
        v = self.wv(x).view(batch_size, seq_len, self.n_heads, self.head_dim)
        
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        
        # 计算完整注意力
        attn = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        
        # 随机稀疏化
        n_sparse = int(seq_len * self.sparsity)
        sparse_mask = torch.zeros(batch_size, self.n_heads, seq_len, seq_len)
        
        for b in range(batch_size):
            for h in range(self.n_heads):
                indices = torch.randperm(seq_len)[:n_sparse]
                sparse_mask[b, h, :, indices] = 1
        
        attn = attn.masked_fill(sparse_mask == 0, float("-inf"))
        attn = F.softmax(attn, dim=-1)
        
        output = torch.matmul(attn, v)
        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, seq_len, -1)
        
        return self.wo(output)

性能比较

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def benchmark_sparse_attention():
    """基准测试稀疏注意力"""
    seq_lengths = [512, 1024, 2048, 4096]
    methods = {
        "standard": MultiHeadAttention(768, 8),
        "sliding_window": SlidingWindowAttention(768, 8, window_size=256),
        "chunked": ChunkedSparseAttention(768, 8, chunk_size=64),
        "lsh": LSHAttention(768, 8),
        "random": RandomSparseAttention(768, 8, sparsity=0.1),
    }
    
    results = {}
    
    for seq_len in seq_lengths:
        print(f"\n序列长度: {seq_len}")
        x = torch.randn(1, seq_len, 768)
        
        for name, model in methods.items():
            try:
                start = time.time()
                _ = model(x)
                time_taken = time.time() - start
                
                # 内存使用
                memory = torch.cuda.max_memory_allocated() / 1024 / 1024 if torch.cuda.is_available() else 0
                
                results[f"{name}_{seq_len}"] = {
                    "time": time_taken,
                    "memory": memory
                }
                
                print(f"  {name}: {time_taken:.4f}s, {memory:.2f}MB")
            
            except Exception as e:
                print(f"  {name}: 失败 - {e}")
    
    return results

选择稀疏策略

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def select_sparse_strategy(seq_len, memory_constraint, accuracy_requirement):
    """根据约束选择稀疏策略"""
    
    if seq_len <= 1024:
        # 短序列：使用标准注意力
        return "standard"
    
    elif seq_len <= 4096:
        # 中等序列：使用滑动窗口
        if memory_constraint == "high":
            return "sliding_window"
        else:
            return "chunked"
    
    else:
        # 长序列：使用稀疏全局注意力
        if accuracy_requirement == "high":
            return "lsh"
        else:
            return "random"

权衡分析

计算复杂度

方法 时间复杂度 空间复杂度 适用场景
标准注意力 O(n²) O(n²) 短序列
滑动窗口 O(nw) O(nw) 局部依赖
分块稀疏 O(n²/c) O(nc) 结构化数据
LSH O(n log n) O(n) 长序列
随机稀疏 O(ns) O(ns) 快速近似

精度-效率权衡

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def analyze_tradeoff(method_name, seq_len):
    """分析精度-效率权衡"""
    
    tradeoffs = {
        "standard": {"accuracy": 1.0, "efficiency": 1.0},
        "sliding_window": {"accuracy": 0.95, "efficiency": 2.0},
        "chunked": {"accuracy": 0.92, "efficiency": 3.0},
        "lsh": {"accuracy": 0.88, "efficiency": 5.0},
        "random": {"accuracy": 0.85, "efficiency": 4.0},
    }
    
    tradeoff = tradeoffs.get(method_name, {"accuracy": 0.8, "efficiency": 2.0})
    
    print(f"方法: {method_name}")
    print(f"序列长度: {seq_len}")
    print(f"相对精度: {tradeoff['accuracy']:.2%}")
    print(f"效率提升: {tradeoff['efficiency']:.2f}x")
    print(f"有效精度: {tradeoff['accuracy'] * tradeoff['efficiency']:.2f}")
    
    return tradeoff

总结

原生稀疏注意力通过只计算部分 token 对之间的注意力，显著降低了计算复杂度。不同方法适用于不同场景：滑动窗口适合局部依赖，LSH 适合长序列，随机稀疏适合快速近似。选择合适的方法需要权衡精度和效率。

下一步

下一课将介绍多 token 预测，加速生成过程。

📚 本文改编自 AI Engineering from Scratch（MIT License · 作者 Rohit Ghumare），中文内容来自官方中文镜像。原课程共 503 课 · 20 阶段 · 免费开源，教程网站见 aiengineeringfromscratch.com。

  • 标题: 【大模型】- 原生稀疏注意力
  • 作者: 宋
  • 创建于 : 2026-08-19 09:17:00
  • 更新于 : 2026-08-21 16:20:11
  • 链接: https://sxl-space.tk/2026/08/19/010_LLM/010_LLM-17-NativeSparseAttention/
  • 版权声明: 版权所有 © 宋，禁止转载。
目录
【大模型】- 原生稀疏注意力
  1. 1. 原生稀疏注意力
  2. 2. 学习目标
  3. 3. 稀疏注意力概述
    1. 3.1. 标准注意力的问题
    2. 3.2. 稀疏注意力的优势
  4. 4. 原生稀疏注意力实现
    1. 4.1. 滑动窗口注意力
    2. 4.2. 分块稀疏注意力
    3. 4.3. 稀疏全局注意力
  5. 5. 高级稀疏模式
    1. 5.1. 局部敏感哈希（LSH）注意力
    2. 5.2. 随机稀疏注意力
  6. 6. 性能比较
  7. 7. 选择稀疏策略
  8. 8. 权衡分析
    1. 8.1. 计算复杂度
    2. 8.2. 精度-效率权衡
  9. 9. 总结
  10. 10. 下一步
  11. 11. 下一课将介绍多 token 预测，加速生成过程。