Spaces:
Running on Zero
Running on Zero
File size: 14,066 Bytes
5f7afef | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | {
"slug": "alibi-attention",
"title": "手撕带线性偏置的注意力机制 (ALiBi Attention)",
"difficulty": "Hard",
"category": "Attention",
"tags": [
"常见手撕",
"ALiBi",
"位置编码",
"外推性",
"Attention"
],
"description_md": "### 题目描述\n**ALiBi (Attention with Linear Biases)** 抛弃了传统显式的绝对/旋转位置编码嵌入，直接在 Softmax 之前的注意力分数上施加一个与相对距离成正比的固定静态惩罚矩阵：\n$$\\text{softmax}\\left( \\frac{QK^T}{\\sqrt{d}} - m \\cdot (i - j) \\right)$$\n其中每个注意力头拥有固定的衰减斜率 $m$，使得模型具备极强的**长度外推性（Length Extrapolation）**。\n\n### 函数签名\n```python\ndef alibi_attention(\n q: torch.Tensor,\n k: torch.Tensor,\n v: torch.Tensor,\n slopes: torch.Tensor,\n) -> torch.Tensor:\n```\n\n### 示例\n```python\n# 输入\nq = [[[[0.33, 0.67], [0.43, 0.57]], [[0.45, 0.55], [0.47, 0.53]]], [[[0.47, 0.53], [0.48, 0.52]], [[0.48, 0.52], [0.48, 0.52]]]]\nk = [[[[3.0, 4.0], [-1.0, 0.5]], [[1.5, -0.5], [2.5, 0.0]]], [[[-1.5, 3.5], [1.0, 2.0]], [[3.0, 4.0], [-1.0, 0.5]]]]\nv = [[[[-1.0, 0.5], [1.5, -0.5]], [[2.5, 0.0], [-1.5, 3.5]]], [[[1.0, 2.0], [3.0, 4.0]], [[-1.0, 0.5], [1.5, -0.5]]]]\nslopes = [1.5, -0.5]\n\n# 输出 (Tensor, shape=[2, 2, 2, 2], dtype=float32)\n[[[[-0.3717, 0.2487], [-0.3882, 0.2553]], [[0.4873, 1.7611], [0.4803, 1.7673]]], [[[2.7086, 3.7086], [2.7156, 3.7156]], [[-0.8967, 0.4587], [-0.8967, 0.4587]]]]\n```\n\n### 参数规格与约束\n- **`q`**：`torch.Tensor`，测试规格 `shape=[2, 4, 8, 16], dtype=float32`\n- **`k`**：`torch.Tensor`，测试规格 `shape=[2, 4, 8, 16], dtype=float32`\n- **`v`**：`torch.Tensor`，测试规格 `shape=[2, 4, 8, 16], dtype=float32`\n- **`slopes`**：`torch.Tensor`，测试规格 `shape=[4], dtype=float32`\n- **返回值**：`torch.Tensor`，输出规格 `shape=[2, 4, 8, 16], dtype=float32`\n- **精度容差**：`atol=0.0001, rtol=0.0001`",
"entry_point": "alibi_attention",
"starter_code": "import torch\nimport torch.nn.functional as F\nimport math\n\ndef alibi_attention(\n q: torch.Tensor,\n k: torch.Tensor,\n v: torch.Tensor,\n slopes: torch.Tensor,\n) -> torch.Tensor:\n \"\"\"\n ALiBi 注意力实现。\n q, k, v: [B, H, S, D]\n slopes: [H]\n \"\"\"\n # 请在此编写你的代码\n pass\n",
"reference_code": "import torch\nimport torch.nn.functional as F\nimport math\n\ndef alibi_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, slopes: torch.Tensor) -> torch.Tensor:\n B, H, S, D = q.shape\n scores = (q @ k.transpose(-2, -1)) / math.sqrt(D)\n positions = torch.arange(S)\n relative = positions.unsqueeze(1) - positions.unsqueeze(0)\n bias = -relative.float().unsqueeze(0) * slopes.unsqueeze(1).unsqueeze(2)\n scores = scores + bias.unsqueeze(0)\n attn = F.softmax(scores, dim=-1)\n return attn @ v\n",
"forbidden_modules": [
"os",
"sys",
"subprocess",
"shutil"
],
"forbidden_calls": [
"MultiheadAttention"
],
"time_limit": 3.0,
"memory_limit_mb": 512,
"atol": 0.0001,
"rtol": 0.0001,
"test_code": "import torch\ncases = [\n (\"标准 4 头测试 (B=2, H=4, S=8, D=16)\", 2, 4, 8, 16),\n (\"长序列 8 头测试 (B=1, H=8, S=16, D=32)\", 1, 8, 16, 32),\n (\"单 Token 步边界测试 (B=2, H=2, S=1, D=8)\", 2, 2, 1, 8),\n (\"奇数序列长度测试 (B=1, H=4, S=7, D=16)\", 1, 4, 7, 16),\n]\nfor idx, (name, B, H, S, D) in enumerate(cases, 1):\n q = torch.randn(B, H, S, D)\n k = torch.randn(B, H, S, D)\n v = torch.randn(B, H, S, D)\n slopes = torch.tensor([2.0 ** (-8.0 * i / H) for i in range(1, H + 1)])\n\n t0 = time.time()\n try:\n exp = ref_obj(q, k, v, slopes)\n act = user_obj(q, k, v, slopes)\n match, err, max_diff = compare_outputs(act, exp, atol=atol, rtol=rtol)\n passed = match\n msg = err\n except Exception as e:\n passed = False\n msg = f\"Runtime Exception: {str(e)}\"\n act = None\n exp = None\n max_diff = None\n\n cases_results.append({\n \"case_id\": idx,\n \"case_name\": name,\n \"passed\": passed,\n \"time_ms\": round((time.time() - t0) * 1000, 2),\n \"message\": msg,\n \"user_output_summary\": summarize_obj(act) if act is not None else None,\n \"expected_summary\": summarize_obj(exp) if exp is not None else None,\n \"max_diff\": max_diff,\n })\n",
"is_featured": true,
"editorial_md": "#### 💡 核心算法原理与数学推导\n**手撕带线性偏置的注意力机制 (ALiBi Attention)**（`alibi_attention`）聚焦于 **Attention** 中的关键计算逻辑：\n\n**ALiBi (Attention with Linear Biases)** 抛弃了传统显式的绝对/旋转位置编码嵌入，直接在 Softmax 之前的注意力分数上施加一个与相对距离成正比的固定静态惩罚矩阵：\n$$\\text{softmax}\\left( \\frac{QK^T}{\\sqrt{d}} - m \\cdot (i - j) \\right)$$\n其中每个注意力头拥有固定的衰减斜率 $m$，使得模型具备极强的**长度外推性（Length Extrapolation）**。\n\n在实现中，函数接收参数 `q` (`torch.Tensor`)、`k` (`torch.Tensor`)、`v` (`torch.Tensor`)、`slopes` (`torch.Tensor`)，目标是计算出类型与形状严格符合 `torch.Tensor` 契约的结果。\n\n#### 📐 输入输出张量变换与步骤拆解\n- **输入基准规格**：`q` 规格为 `torch.Tensor(shape=[2, 4, 8, 16], dtype=float32)`，`k` 规格为 `torch.Tensor(shape=[2, 4, 8, 16], dtype=float32)`，`v` 规格为 `torch.Tensor(shape=[2, 4, 8, 16], dtype=float32)`，`slopes` 规格为 `torch.Tensor(shape=[4], dtype=float32)`。\n- **期望输出规格**：`torch.Tensor(shape=[2, 4, 8, 16], dtype=float32)`。\n\n**具体推导与计算步骤如下：**\n1. 计算中间量 `B, H, S, D`：通过表达式 `q.shape` 完成张量或数值变换。\n2. 计算中间量 `scores`：通过表达式 `(q @ k.transpose(-2, -1)) / math.sqrt(D)` 完成张量或数值变换。\n3. 计算中间量 `positions`：通过表达式 `torch.arange(S)` 完成张量或数值变换。\n4. 计算中间量 `relative`：通过表达式 `positions.unsqueeze(1) - positions.unsqueeze(0)` 完成张量或数值变换。\n5. 计算中间量 `bias`：通过表达式 `-relative.float().unsqueeze(0) * slopes.unsqueeze(1).unsqueeze(2)` 完成张量或数值变换。\n6. 计算中间量 `scores`：通过表达式 `scores + bias.unsqueeze(0)` 完成张量或数值变换。\n7. 计算中间量 `attn`：通过表达式 `F.softmax(scores, dim=-1)` 完成张量或数值变换。\n8. 返回最终结果：输出 `attn @ v`，严格对齐函数签名要求的返回类型与维度。\n\n#### 📊 复杂度分析\n- **时间复杂度**：$O(N \\cdot M \\cdot D)$，其中 $N, M$ 为矩阵行/列样本规模、$D$ 为内积特征维度，由矩阵乘法算子主导。\n- **空间复杂度**：$O(N \\cdot M)$，用于存储输出矩阵及中间内积张量。\n\n#### ⚠️ 关键实现细节与易错点\n- **除零保护与数值平滑 (`eps`)**：在执行分母归一化、标准差缩放或对数运算时，必须按公式位置准确加上 `eps`（例如在平方根内 `sqrt(var + eps)` 还是分母外 `+ eps`），防止全零或常量输入触发 `NaN` / `Inf`。\n- **归约维度 (`dim`) 与广播对齐 (`keepdim`)**：调用 `sum`、`mean`、`max`、`norm` 或 `softmax` 时需明确指定目標维度（通常为最后一维 `dim=-1` 或特征维 `dim=0`），并在需要与原张量逐元素相减/相除时开启 `keepdim=True`。\n- **内存连续性与维度重排顺序**：在执行 `transpose` 或 `permute` 交换多头/序列维度后，张量在底层内存中变为非连续布局；若后续调用 `.view(...)` 合并维度，必须先调用 `.contiguous()` 或改用 `.reshape(...)`。\n\n#### 💻 官方标准参考实现\n```python\nimport torch\nimport torch.nn.functional as F\nimport math\n\ndef alibi_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, slopes: torch.Tensor) -> torch.Tensor:\n B, H, S, D = q.shape\n scores = (q @ k.transpose(-2, -1)) / math.sqrt(D)\n positions = torch.arange(S)\n relative = positions.unsqueeze(1) - positions.unsqueeze(0)\n bias = -relative.float().unsqueeze(0) * slopes.unsqueeze(1).unsqueeze(2)\n scores = scores + bias.unsqueeze(0)\n attn = F.softmax(scores, dim=-1)\n return attn @ v\n```",
"difficulty_breakdown": {
"concept": 4,
"implementation": 4,
"math_derivation": 4,
"boundary_complexity": 4
},
"test_cases": [
{
"name": "标准 4 头测试 (B=2, H=4, S=8, D=16)",
"input_desc": "真实参数契约用例 (q, k, v, slopes)",
"params": [
{
"name": "q",
"type": "torch.Tensor",
"value": "torch.tensor([[[[0.33, 0.67], [0.43, 0.57]], [[0.45, 0.55], [0.47, 0.53]]], [[[0.47, 0.53], [0.48, 0.52]], [[0.48, 0.52], [0.48, 0.52]]]], dtype=torch.float32)"
},
{
"name": "k",
"type": "torch.Tensor",
"value": "torch.tensor([[[[3.0, 4.0], [-1.0, 0.5]], [[1.5, -0.5], [2.5, 0.0]]], [[[-1.5, 3.5], [1.0, 2.0]], [[3.0, 4.0], [-1.0, 0.5]]]], dtype=torch.float32)"
},
{
"name": "v",
"type": "torch.Tensor",
"value": "torch.tensor([[[[-1.0, 0.5], [1.5, -0.5]], [[2.5, 0.0], [-1.5, 3.5]]], [[[1.0, 2.0], [3.0, 4.0]], [[-1.0, 0.5], [1.5, -0.5]]]], dtype=torch.float32)"
},
{
"name": "slopes",
"type": "torch.Tensor",
"value": "torch.tensor([1.5, -0.5], dtype=torch.float32)"
}
],
"raw_inputs": {
"q": "torch.tensor([[[[0.33, 0.67], [0.43, 0.57]], [[0.45, 0.55], [0.47, 0.53]]], [[[0.47, 0.53], [0.48, 0.52]], [[0.48, 0.52], [0.48, 0.52]]]], dtype=torch.float32)",
"k": "torch.tensor([[[[3.0, 4.0], [-1.0, 0.5]], [[1.5, -0.5], [2.5, 0.0]]], [[[-1.5, 3.5], [1.0, 2.0]], [[3.0, 4.0], [-1.0, 0.5]]]], dtype=torch.float32)",
"v": "torch.tensor([[[[-1.0, 0.5], [1.5, -0.5]], [[2.5, 0.0], [-1.5, 3.5]]], [[[1.0, 2.0], [3.0, 4.0]], [[-1.0, 0.5], [1.5, -0.5]]]], dtype=torch.float32)",
"slopes": "torch.tensor([1.5, -0.5], dtype=torch.float32)"
}
},
{
"name": "长序列 8 头测试 (B=1, H=8, S=16, D=32)",
"input_desc": "真实参数契约用例 (q, k, v, slopes)",
"params": [
{
"name": "q",
"type": "torch.Tensor",
"value": "torch.randn(1, 8, 16, 32)"
},
{
"name": "k",
"type": "torch.Tensor",
"value": "torch.randn(1, 8, 16, 32)"
},
{
"name": "v",
"type": "torch.Tensor",
"value": "torch.randn(1, 8, 16, 32)"
},
{
"name": "slopes",
"type": "torch.Tensor",
"value": "torch.tensor([0.5, 0.25, 0.125, 0.0625, 0.0312, 0.0156, 0.0078, 0.0039], dtype=torch.float32)"
}
],
"raw_inputs": {
"q": "torch.randn(1, 8, 16, 32)",
"k": "torch.randn(1, 8, 16, 32)",
"v": "torch.randn(1, 8, 16, 32)",
"slopes": "torch.tensor([0.5, 0.25, 0.125, 0.0625, 0.0312, 0.0156, 0.0078, 0.0039], dtype=torch.float32)"
}
},
{
"name": "单 Token 步边界测试 (B=2, H=2, S=1, D=8)",
"input_desc": "真实参数契约用例 (q, k, v, slopes)",
"params": [
{
"name": "q",
"type": "torch.Tensor",
"value": "torch.tensor([[[[0.1593, 0.1944, -0.9836, -0.2747, -0.2545, -0.4611, 0.3362, 0.8286]], [[-0.0239, -0.5277, -1.0591, -0.9485, 1.004, 0.5197, -1.1292, -0.3924]]], [[[-0.2783, 0.304, 0.3854, 1.489, -0.9158, -0.1873, -1.0439, -0.0129]], [[-0.967, -0.1866, -1.0899, 0.9881, -0.8024, 0.0332, -0.752, 0.7219]]]], dtype=torch.float32)"
},
{
"name": "k",
"type": "torch.Tensor",
"value": "torch.tensor([[[[0.2591, -0.648, 0.3581, -1.4733, 0.3322, -1.6188, 2.6514, 0.2979]], [[0.361, -0.4477, 1.1588, -0.906, 0.1041, 0.3926, -0.7857, -0.2947]]], [[[-1.2036, 1.4919, 0.1389, -0.4275, -1.8924, -1.0036, 0.313, 1.4533]], [[-0.5578, 0.5746, 2.4484, -0.1331, -0.7109, -1.1014, -0.2638, -0.8466]]]], dtype=torch.float32)"
},
{
"name": "v",
"type": "torch.Tensor",
"value": "torch.tensor([[[[0.8731, 1.2488, 0.9059, -1.9169, 1.2062, 0.1234, -0.2116, 0.158]], [[1.2571, -0.109, 0.7571, 1.6256, 0.2195, -0.5969, 0.7425, -0.0007]]], [[[0.8032, -0.4695, 0.0964, 1.0252, -0.7467, 0.3237, -0.4603, 0.4073]], [[1.7341, -0.6813, -0.8275, 0.1656, -0.6037, -0.1836, -1.1753, 0.4287]]]], dtype=torch.float32)"
},
{
"name": "slopes",
"type": "torch.Tensor",
"value": "torch.tensor([0.0625, 0.0039], dtype=torch.float32)"
}
],
"raw_inputs": {
"q": "torch.tensor([[[[0.1593, 0.1944, -0.9836, -0.2747, -0.2545, -0.4611, 0.3362, 0.8286]], [[-0.0239, -0.5277, -1.0591, -0.9485, 1.004, 0.5197, -1.1292, -0.3924]]], [[[-0.2783, 0.304, 0.3854, 1.489, -0.9158, -0.1873, -1.0439, -0.0129]], [[-0.967, -0.1866, -1.0899, 0.9881, -0.8024, 0.0332, -0.752, 0.7219]]]], dtype=torch.float32)",
"k": "torch.tensor([[[[0.2591, -0.648, 0.3581, -1.4733, 0.3322, -1.6188, 2.6514, 0.2979]], [[0.361, -0.4477, 1.1588, -0.906, 0.1041, 0.3926, -0.7857, -0.2947]]], [[[-1.2036, 1.4919, 0.1389, -0.4275, -1.8924, -1.0036, 0.313, 1.4533]], [[-0.5578, 0.5746, 2.4484, -0.1331, -0.7109, -1.1014, -0.2638, -0.8466]]]], dtype=torch.float32)",
"v": "torch.tensor([[[[0.8731, 1.2488, 0.9059, -1.9169, 1.2062, 0.1234, -0.2116, 0.158]], [[1.2571, -0.109, 0.7571, 1.6256, 0.2195, -0.5969, 0.7425, -0.0007]]], [[[0.8032, -0.4695, 0.0964, 1.0252, -0.7467, 0.3237, -0.4603, 0.4073]], [[1.7341, -0.6813, -0.8275, 0.1656, -0.6037, -0.1836, -1.1753, 0.4287]]]], dtype=torch.float32)",
"slopes": "torch.tensor([0.0625, 0.0039], dtype=torch.float32)"
}
}
]
} |