Spaces:
Running on Zero
Running on Zero
| { | |
| "slug": "based-taylor-linear-attention", | |
| "title": "Based 二阶泰勒展开特征映射线性注意力", | |
| "difficulty": "Medium", | |
| "difficulty_breakdown": { | |
| "concept": 3, | |
| "implementation": 3, | |
| "math_derivation": 3, | |
| "boundary_complexity": 3 | |
| }, | |
| "category": "Attention", | |
| "tags": [ | |
| "Based", | |
| "LinearAttention", | |
| "TaylorExpansion", | |
| "labml.ai" | |
| ], | |
| "description_md": "### 题目描述\n实现 Based (Arora et al., 2024) 近似 $\\exp(q^\\top k / \\sqrt{d})$ 的二阶泰勒特征映射 $\\phi(x)$：给定向量 `x`（形状 `[B, S, D]`），令缩放后 $\\tilde{x} = x / D^{1/4}$，拼接零阶常数项 `1.0`（形状 `[B, S, 1]`）、一阶项 `$x / \\sqrt{D}$`（形状 `[B, S, D]`）以及二阶外积展平项 $(\\tilde{x} \\otimes \\tilde{x}) / \\sqrt{2}$（形状 `[B, S, D*D]`），返回形状 `[B, S, 1 + D + D*D]` 的特征张量。\n\n### 函数签名\n```python\ndef based_taylor_feature_map(x: torch.Tensor) -> torch.Tensor:\n```\n\n### 示例\n```python\n# 输入\nx = [[[1.0, 2.0], [0.0, -1.0]]]\n\n# 输出 (Tensor, shape=[1, 2, 7], dtype=float32)\n[[[1.0, 0.7071, 1.4142, 0.5, 1.0, 1.0, 2.0], [1.0, 0.0, -0.7071, 0.0, 0.0, 0.0, 0.5]]]\n```\n\n### 参数规格与约束\n- **`x`**：`torch.Tensor`，测试规格 `shape=[1, 2, 2], dtype=float32`\n- **返回值**：`torch.Tensor`，输出规格 `shape=[1, 2, 7], dtype=float32`\n- **精度容差**：`atol=0.0001, rtol=0.0001`\n> 题目来源：改编自 [labml.ai annotated_deep_learning_paper_implementations](https://github.com/labmlai/annotated_deep_learning_paper_implementations)（MIT License），已改写为函数式评测接口与中文题面。", | |
| "entry_point": "based_taylor_feature_map", | |
| "starter_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom typing import List, Tuple, Dict, Optional\n\ndef based_taylor_feature_map(x: torch.Tensor) -> torch.Tensor:\n # 请在此处编写你的实现代码\n pass\n", | |
| "reference_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom typing import List, Tuple, Dict, Optional\n\ndef based_taylor_feature_map(x: torch.Tensor) -> torch.Tensor:\n b, s, d = x.shape\n t0 = torch.ones(b, s, 1, dtype=x.dtype)\n t1 = x / math.sqrt(d)\n x_scaled = x / (d ** 0.25)\n t2 = torch.einsum('bsi,bsj->bsij', x_scaled, x_scaled).reshape(b, s, d * d) / math.sqrt(2.0)\n return torch.cat([t0, t1, t2], dim=-1)\n", | |
| "test_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom typing import List, Tuple, Dict, Optional\n\ncases = [\n (\"手算 B=1,S=2,D=2\", torch.tensor([[[1.0, 2.0], [0.0, -1.0]]])),\n (\"零向量仅保留常数项 1\", torch.zeros(2, 3, 2)),\n (\"B=2,S=4,D=3\", torch.randn(2, 4, 3)),\n (\"B=3,S=2,D=4\", torch.randn(3, 2, 4)),\n]\nfor idx, (name, *args) in enumerate(cases, 1):\n t0 = time.time()\n try:\n exp = ref_obj(*args)\n act = user_obj(*args)\n match, err, max_diff = compare_outputs(act, exp, atol=atol, rtol=rtol)\n passed, msg = match, err\n except Exception as e:\n passed, msg = False, f\"Runtime Exception: {str(e)}\"\n act, exp, max_diff = None, None, None\n cases_results.append({\n \"case_id\": idx, \"case_name\": name, \"passed\": passed,\n \"time_ms\": round((time.time() - t0) * 1000, 2), \"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", | |
| "test_cases": [ | |
| { | |
| "name": "手算 B=1,S=2,D=2", | |
| "input_desc": "真实参数契约用例 (x)", | |
| "params": [ | |
| { | |
| "name": "x", | |
| "type": "torch.Tensor", | |
| "value": "torch.tensor([[[1.0, 2.0], [0.0, -1.0]]], dtype=torch.float32)" | |
| } | |
| ], | |
| "raw_inputs": { | |
| "x": "torch.tensor([[[1.0, 2.0], [0.0, -1.0]]], dtype=torch.float32)" | |
| } | |
| }, | |
| { | |
| "name": "零向量仅保留常数项 1", | |
| "input_desc": "真实参数契约用例 (x)", | |
| "params": [ | |
| { | |
| "name": "x", | |
| "type": "torch.Tensor", | |
| "value": "torch.tensor([[[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]], dtype=torch.float32)" | |
| } | |
| ], | |
| "raw_inputs": { | |
| "x": "torch.tensor([[[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]], dtype=torch.float32)" | |
| } | |
| }, | |
| { | |
| "name": "B=2,S=4,D=3", | |
| "input_desc": "真实参数契约用例 (x)", | |
| "params": [ | |
| { | |
| "name": "x", | |
| "type": "torch.Tensor", | |
| "value": "torch.tensor([[[1.9269, 1.4873, 0.9007], [-2.1055, 0.6784, -1.2345], [-0.0431, -1.6047, 0.3559], [-0.6866, -0.4934, 0.2415]], [[-1.1109, 0.0915, -2.3169], [-0.2168, -0.3097, -0.3957], [0.8034, -0.6216, -0.592], [-0.0631, -0.8286, 0.3309]]], dtype=torch.float32)" | |
| } | |
| ], | |
| "raw_inputs": { | |
| "x": "torch.tensor([[[1.9269, 1.4873, 0.9007], [-2.1055, 0.6784, -1.2345], [-0.0431, -1.6047, 0.3559], [-0.6866, -0.4934, 0.2415]], [[-1.1109, 0.0915, -2.3169], [-0.2168, -0.3097, -0.3957], [0.8034, -0.6216, -0.592], [-0.0631, -0.8286, 0.3309]]], dtype=torch.float32)" | |
| } | |
| } | |
| ], | |
| "time_limit": 5.0, | |
| "memory_limit_mb": 512, | |
| "atol": 0.0001, | |
| "rtol": 0.0001, | |
| "forbidden_modules": [ | |
| "os", | |
| "sys", | |
| "subprocess", | |
| "socket", | |
| "shutil", | |
| "pathlib", | |
| "ctypes", | |
| "multiprocessing" | |
| ], | |
| "forbidden_calls": [ | |
| "eval", | |
| "exec", | |
| "compile", | |
| "open", | |
| "__import__", | |
| "input", | |
| "breakpoint" | |
| ], | |
| "is_featured": false, | |
| "editorial_md": "#### 💡 核心算法原理与数学推导\n**Based 二阶泰勒展开特征映射线性注意力**（`based_taylor_feature_map`）聚焦于 **Attention** 中的关键计算逻辑：\n\n实现 Based (Arora et al., 2024) 近似 $\\exp(q^\\top k / \\sqrt{d})$ 的二阶泰勒特征映射 $\\phi(x)$：给定向量 `x`（形状 `[B, S, D]`），令缩放后 $\\tilde{x} = x / D^{1/4}$，拼接零阶常数项 `1.0`（形状 `[B, S, 1]`）、一阶项 `$x / \\sqrt{D}$`（形状 `[B, S, D]`）以及二阶外积展平项 $(\\tilde{x} \\otimes \\tilde{x}) / \\sqrt{2}$（形状 `[B, S, D*D]`），返回形状 `[B, S, 1 + D + D*D]` 的特征张量。\n\n在实现中，函数接收参数 `x` (`torch.Tensor`)，目标是计算出类型与形状严格符合 `torch.Tensor` 契约的结果。\n\n#### 📐 输入输出张量变换与步骤拆解\n- **输入基准规格**：`x` 规格为 `torch.Tensor(shape=[1, 2, 2], dtype=float32)`。\n- **期望输出规格**：`torch.Tensor(shape=[1, 2, 7], dtype=float32)`。\n\n**具体推导与计算步骤如下：**\n1. 计算中间量 `b, s, d`：通过表达式 `x.shape` 完成张量或数值变换。\n2. 计算中间量 `t0`：通过表达式 `torch.ones(b, s, 1, dtype=x.dtype)` 完成张量或数值变换。\n3. 计算中间量 `t1`：通过表达式 `x / math.sqrt(d)` 完成张量或数值变换。\n4. 计算中间量 `x_scaled`：通过表达式 `x / (d ** 0.25)` 完成张量或数值变换。\n5. 计算中间量 `t2`：通过表达式 `torch.einsum('bsi,bsj->bsij', x_scaled, x_scaled).reshape(b, s, d * d) / math.sqrt(2.0)` 完成张量或数值变换。\n6. 返回最终结果：输出 `torch.cat([t0, t1, t2], dim=-1)`，严格对齐函数签名要求的返回类型与维度。\n\n#### 📊 复杂度分析\n- **时间复杂度**：$O(N \\cdot D^2 + D^3)$，主要由矩阵乘积与线性代数求逆/奇异值分解开销主导。\n- **空间复杂度**：$O(N \\cdot D + D^2)$，用于存储核矩阵/协方差矩阵及分解结果张量。\n\n#### ⚠️ 关键实现细节与易错点\n- **除零保护与数值平滑 (`eps`)**：在执行分母归一化、标准差缩放或对数运算时，需加上 `eps` 或执行 `clamp_min` 防止除零产生 `NaN` / `Inf`。\n- **归约维度 (`dim`) 与广播对齐 (`keepdim`)**：调用归约算子时需明确指定目标轴（如 `dim=-1`），并在需要与原张量逐元素运算时开启 `keepdim=True`。\n- **向量化算子优先**：优先使用 PyTorch 原生向量化张量算子，避免低效的逐元素循环。\n\n#### 💻 官方标准参考实现\n```python\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport math\nfrom typing import List, Tuple, Dict, Optional\n\ndef based_taylor_feature_map(x: torch.Tensor) -> torch.Tensor:\n b, s, d = x.shape\n t0 = torch.ones(b, s, 1, dtype=x.dtype)\n t1 = x / math.sqrt(d)\n x_scaled = x / (d ** 0.25)\n t2 = torch.einsum('bsi,bsj->bsij', x_scaled, x_scaled).reshape(b, s, d * d) / math.sqrt(2.0)\n return torch.cat([t0, t1, t2], dim=-1)\n```" | |
| } |