{
  "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```"
}