{
  "slug": "adam-optimizer-step",
  "title": "手撕 Adam 优化器单步更新 (Adam Optimizer Step)",
  "difficulty": "Hard",
  "category": "ML",
  "tags": [
    "常见手撕",
    "优化算法",
    "Adam",
    "动量",
    "梯度下降"
  ],
  "description_md": "### 题目描述\nAdam 综合利用梯度的一阶动量与二阶动量，并进行偏差校正：\n$$m_t = \\beta_1 m_{t-1} + (1 - \\beta_1) g_t$$\n$$v_t = \\beta_2 v_{t-1} + (1 - \\beta_2) g_t^2$$\n$$\\hat{m}_t = \\frac{m_t}{1 - \\beta_1^t}, \\quad \\hat{v}_t = \\frac{v_t}{1 - \\beta_2^t}$$\n$$\\theta_t = \\theta_{t-1} - \\frac{\\alpha}{\\sqrt{\\hat{v}_t} + \\epsilon} \\hat{m}_t$$\n\n### 函数签名\n```python\ndef adam_step(\n    param: torch.Tensor,\n    grad: torch.Tensor,\n    m: torch.Tensor,\n    v: torch.Tensor,\n    t: int,\n    lr: float = 1e-3,\n    beta1: float = 0.9,\n    beta2: float = 0.999,\n    eps: float = 1e-8,\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n```\n\n### 示例\n```python\n# 输入\nparam = [1.0, 2.0, -1.0]\ngrad = [0.1, -0.2, 0.05]\nm = [0.05, -0.1, 0.02]\nv = [1.5, -0.5, 2.5]\nt = 1\nlr = 1e-3\nbeta1 = 0.9\nbeta2 = 0.999\neps = 1e-8\n\n# 输出\n(\n  [1.0, nan, -1.0],  # shape=[3], dtype=float32\n  [0.055, -0.11, 0.023],  # shape=[3], dtype=float32\n  [1.4985, -0.4995, 2.4975]  # shape=[3], dtype=float32\n)\n```\n\n### 参数规格与约束\n- **`param`**：`torch.Tensor`，测试规格 `shape=[3], dtype=float32`\n- **`grad`**：`torch.Tensor`，测试规格 `shape=[3], dtype=float32`\n- **`m`**：`torch.Tensor`，测试规格 `shape=[3], dtype=float32`\n- **`v`**：`torch.Tensor`，测试规格 `shape=[3], dtype=float32`\n- **`t`**：`int`\n- **`lr`**：`float`（默认 `1e-3`）\n- **`beta1`**：`float`（默认 `0.9`）\n- **`beta2`**：`float`（默认 `0.999`）\n- **`eps`**：`float`（默认 `1e-8`）\n- **返回值**：`tuple[torch.Tensor, torch.Tensor, torch.Tensor]`，输出规格 `Tuple` (`torch.Tensor(shape=[3], dtype=float32)`, `torch.Tensor(shape=[3], dtype=float32)`, `torch.Tensor(shape=[3], dtype=float32)`)\n- **精度容差**：`atol=0.0001, rtol=0.0001`",
  "entry_point": "adam_step",
  "starter_code": "import torch\n\ndef adam_step(\n    param: torch.Tensor,\n    grad: torch.Tensor,\n    m: torch.Tensor,\n    v: torch.Tensor,\n    t: int,\n    lr: float = 1e-3,\n    beta1: float = 0.9,\n    beta2: float = 0.999,\n    eps: float = 1e-8,\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n    \"\"\"\n    执行单步 Adam 梯度更新并返回 (param, m, v)\n    \"\"\"\n    # 请在此编写代码\n    pass\n",
  "reference_code": "import torch\n\ndef adam_step(\n    param: torch.Tensor,\n    grad: torch.Tensor,\n    m: torch.Tensor,\n    v: torch.Tensor,\n    t: int,\n    lr: float = 1e-3,\n    beta1: float = 0.9,\n    beta2: float = 0.999,\n    eps: float = 1e-8,\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n    new_m = beta1 * m + (1.0 - beta1) * grad\n    new_v = beta2 * v + (1.0 - beta2) * (grad ** 2)\n    m_hat = new_m / (1.0 - beta1 ** t)\n    v_hat = new_v / (1.0 - beta2 ** t)\n    new_param = param - lr * m_hat / (torch.sqrt(v_hat) + eps)\n    return new_param, new_m, new_v\n",
  "forbidden_modules": [
    "os",
    "sys",
    "subprocess",
    "shutil",
    "socket",
    "requests"
  ],
  "forbidden_calls": [],
  "time_limit": 3.0,
  "memory_limit_mb": 512,
  "atol": 0.0001,
  "rtol": 0.0001,
  "test_code": "import torch\ncases = [\n    (\"第 1 步初始偏差修正更新 (t=1)\", 1),\n    (\"第 5 步过渡期更新 (t=5)\", 5),\n    (\"第 10 步稳定更新 (t=10)\", 10),\n    (\"第 100 步长程渐近更新 (t=100)\", 100),\n]\nfor idx, (name, step_t) in enumerate(cases, 1):\n    p = torch.tensor([1.0, 2.0, -1.0])\n    g = torch.tensor([0.1, -0.2, 0.05])\n    m = torch.tensor([0.05, -0.1, 0.02])\n    v = torch.tensor([0.01, 0.04, 0.002])\n    t0 = time.time()\n    try:\n        exp = ref_obj(p, g, m, v, step_t)\n        act = user_obj(p, g, m, v, step_t)\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, 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",
  "is_featured": true,
  "editorial_md": "#### 💡 核心算法原理与数学推导\n**手撕 Adam 优化器单步更新 (Adam Optimizer Step)**（`adam_step`）聚焦于 **ML** 中的关键计算逻辑：\n\nAdam 综合利用梯度的一阶动量与二阶动量，并进行偏差校正：\n$$m_t = \\beta_1 m_{t-1} + (1 - \\beta_1) g_t$$\n$$v_t = \\beta_2 v_{t-1} + (1 - \\beta_2) g_t^2$$\n$$\\hat{m}_t = \\frac{m_t}{1 - \\beta_1^t}, \\quad \\hat{v}_t = \\frac{v_t}{1 - \\beta_2^t}$$\n$$\\theta_t = \\theta_{t-1} - \\frac{\\alpha}{\\sqrt{\\hat{v}_t} + \\epsilon} \\hat{m}_t$$\n\n在实现中，函数接收参数 `param` (`torch.Tensor`)、`grad` (`torch.Tensor`)、`m` (`torch.Tensor`)、`v` (`torch.Tensor`)、`t` (`int`)、`lr` (`float`)、`beta1` (`float`)、`beta2` (`float`)、`eps` (`float`)，目标是计算出类型与形状严格符合 `tuple[torch.Tensor, torch.Tensor, torch.Tensor]` 契约的结果。\n\n#### 📐 输入输出张量变换与步骤拆解\n- **输入基准规格**：`param` 规格为 `torch.Tensor(shape=[3], dtype=float32)`，`grad` 规格为 `torch.Tensor(shape=[3], dtype=float32)`，`m` 规格为 `torch.Tensor(shape=[3], dtype=float32)`，`v` 规格为 `torch.Tensor(shape=[3], dtype=float32)`，`t` 规格为 `int` (示例值: `1`)。\n- **期望输出规格**：`Tuple` (`torch.Tensor(shape=[3], dtype=float32)`, `torch.Tensor(shape=[3], dtype=float32)`, `torch.Tensor(shape=[3], dtype=float32)`)。\n\n**具体推导与计算步骤如下：**\n1. 计算中间量 `new_m`：通过表达式 `beta1 * m + (1.0 - beta1) * grad` 完成张量或数值变换。\n2. 计算中间量 `new_v`：通过表达式 `beta2 * v + (1.0 - beta2) * (grad ** 2)` 完成张量或数值变换。\n3. 计算中间量 `m_hat`：通过表达式 `new_m / (1.0 - beta1 ** t)` 完成张量或数值变换。\n4. 计算中间量 `v_hat`：通过表达式 `new_v / (1.0 - beta2 ** t)` 完成张量或数值变换。\n5. 计算中间量 `new_param`：通过表达式 `param - lr * m_hat / (torch.sqrt(v_hat) + eps)` 完成张量或数值变换。\n6. 返回最终结果：输出 `new_param, new_m, new_v`，严格对齐函数签名要求的返回类型与维度。\n\n#### 📊 复杂度分析\n- **时间复杂度**：$O(N)$（线性于输入张量总元素个数 $\\prod d_i$），所有运算均为向量化逐元素算子或沿指定维度的归约计算。\n- **空间复杂度**：$O(N)$，需要与输入或输出张量同规模的显存空间保存计算结果。\n\n#### ⚠️ 关键实现细节与易错点\n- **除零保护与数值平滑 (`eps`)**：在执行分母归一化、标准差缩放或对数运算时，必须按公式位置准确加上 `eps`（例如在平方根内 `sqrt(var + eps)` 还是分母外 `+ eps`），防止全零或常量输入触发 `NaN` / `Inf`。\n- **边界输入鲁棒性**：注意单样本（`Batch=1`）、单步序列（`Seq=1`）、零值输入或极端超参数下的计算行为，不要在代码中硬编码特定测试用例的维度常数。\n- **纯向量化算子优先**：尽量使用 PyTorch 原生向量化张量算子替代低效的 Python 逐元素循环，以兼顾代码简洁性与执行性能。\n\n#### 💻 官方标准参考实现\n```python\nimport torch\n\ndef adam_step(\n    param: torch.Tensor,\n    grad: torch.Tensor,\n    m: torch.Tensor,\n    v: torch.Tensor,\n    t: int,\n    lr: float = 1e-3,\n    beta1: float = 0.9,\n    beta2: float = 0.999,\n    eps: float = 1e-8,\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n    new_m = beta1 * m + (1.0 - beta1) * grad\n    new_v = beta2 * v + (1.0 - beta2) * (grad ** 2)\n    m_hat = new_m / (1.0 - beta1 ** t)\n    v_hat = new_v / (1.0 - beta2 ** t)\n    new_param = param - lr * m_hat / (torch.sqrt(v_hat) + eps)\n    return new_param, new_m, new_v\n```",
  "difficulty_breakdown": {
    "concept": 4,
    "implementation": 5,
    "math_derivation": 4,
    "boundary_complexity": 4
  },
  "test_cases": [
    {
      "name": "第 1 步初始偏差修正更新 (t=1)",
      "input_desc": "真实参数契约用例 (param, grad, m, v, t, lr, beta1, beta2, eps)",
      "params": [
        {
          "name": "param",
          "type": "torch.Tensor",
          "value": "torch.tensor([1.0, 2.0, -1.0], dtype=torch.float32)"
        },
        {
          "name": "grad",
          "type": "torch.Tensor",
          "value": "torch.tensor([0.1, -0.2, 0.05], dtype=torch.float32)"
        },
        {
          "name": "m",
          "type": "torch.Tensor",
          "value": "torch.tensor([0.05, -0.1, 0.02], dtype=torch.float32)"
        },
        {
          "name": "v",
          "type": "torch.Tensor",
          "value": "torch.tensor([1.5, -0.5, 2.5], dtype=torch.float32)"
        },
        {
          "name": "t",
          "type": "int",
          "value": "1"
        },
        {
          "name": "lr",
          "type": "float",
          "value": "1e-3"
        },
        {
          "name": "beta1",
          "type": "float",
          "value": "0.9"
        },
        {
          "name": "beta2",
          "type": "float",
          "value": "0.999"
        },
        {
          "name": "eps",
          "type": "float",
          "value": "1e-8"
        }
      ],
      "raw_inputs": {
        "param": "torch.tensor([1.0, 2.0, -1.0], dtype=torch.float32)",
        "grad": "torch.tensor([0.1, -0.2, 0.05], dtype=torch.float32)",
        "m": "torch.tensor([0.05, -0.1, 0.02], dtype=torch.float32)",
        "v": "torch.tensor([1.5, -0.5, 2.5], dtype=torch.float32)",
        "t": "1",
        "lr": "1e-3",
        "beta1": "0.9",
        "beta2": "0.999",
        "eps": "1e-8"
      }
    },
    {
      "name": "第 5 步过渡期更新 (t=5)",
      "input_desc": "真实参数契约用例 (param, grad, m, v, t, lr, beta1, beta2, eps)",
      "params": [
        {
          "name": "param",
          "type": "torch.Tensor",
          "value": "torch.tensor([1.0, 2.0, -1.0], dtype=torch.float32)"
        },
        {
          "name": "grad",
          "type": "torch.Tensor",
          "value": "torch.tensor([0.1, -0.2, 0.05], dtype=torch.float32)"
        },
        {
          "name": "m",
          "type": "torch.Tensor",
          "value": "torch.tensor([0.05, -0.1, 0.02], dtype=torch.float32)"
        },
        {
          "name": "v",
          "type": "torch.Tensor",
          "value": "torch.tensor([0.01, 0.04, 0.002], dtype=torch.float32)"
        },
        {
          "name": "t",
          "type": "int",
          "value": "5"
        },
        {
          "name": "lr",
          "type": "float",
          "value": "1e-3"
        },
        {
          "name": "beta1",
          "type": "float",
          "value": "0.9"
        },
        {
          "name": "beta2",
          "type": "float",
          "value": "0.999"
        },
        {
          "name": "eps",
          "type": "float",
          "value": "1e-8"
        }
      ],
      "raw_inputs": {
        "param": "torch.tensor([1.0, 2.0, -1.0], dtype=torch.float32)",
        "grad": "torch.tensor([0.1, -0.2, 0.05], dtype=torch.float32)",
        "m": "torch.tensor([0.05, -0.1, 0.02], dtype=torch.float32)",
        "v": "torch.tensor([0.01, 0.04, 0.002], dtype=torch.float32)",
        "t": "5",
        "lr": "1e-3",
        "beta1": "0.9",
        "beta2": "0.999",
        "eps": "1e-8"
      }
    },
    {
      "name": "第 10 步稳定更新 (t=10)",
      "input_desc": "真实参数契约用例 (param, grad, m, v, t, lr, beta1, beta2, eps)",
      "params": [
        {
          "name": "param",
          "type": "torch.Tensor",
          "value": "torch.tensor([1.0, 2.0, -1.0], dtype=torch.float32)"
        },
        {
          "name": "grad",
          "type": "torch.Tensor",
          "value": "torch.tensor([0.1, -0.2, 0.05], dtype=torch.float32)"
        },
        {
          "name": "m",
          "type": "torch.Tensor",
          "value": "torch.tensor([0.05, -0.1, 0.02], dtype=torch.float32)"
        },
        {
          "name": "v",
          "type": "torch.Tensor",
          "value": "torch.tensor([0.01, 0.04, 0.002], dtype=torch.float32)"
        },
        {
          "name": "t",
          "type": "int",
          "value": "10"
        },
        {
          "name": "lr",
          "type": "float",
          "value": "1e-3"
        },
        {
          "name": "beta1",
          "type": "float",
          "value": "0.9"
        },
        {
          "name": "beta2",
          "type": "float",
          "value": "0.999"
        },
        {
          "name": "eps",
          "type": "float",
          "value": "1e-8"
        }
      ],
      "raw_inputs": {
        "param": "torch.tensor([1.0, 2.0, -1.0], dtype=torch.float32)",
        "grad": "torch.tensor([0.1, -0.2, 0.05], dtype=torch.float32)",
        "m": "torch.tensor([0.05, -0.1, 0.02], dtype=torch.float32)",
        "v": "torch.tensor([0.01, 0.04, 0.002], dtype=torch.float32)",
        "t": "10",
        "lr": "1e-3",
        "beta1": "0.9",
        "beta2": "0.999",
        "eps": "1e-8"
      }
    }
  ]
}