{ "slug": "adamax-infinity-norm-optimizer-step", "title": "AdaMax 无穷范数指数移动上界自适应优化器单步", "difficulty": "Easy", "difficulty_breakdown": { "concept": 2, "implementation": 2, "math_derivation": 2, "boundary_complexity": 2 }, "category": "Optim", "tags": [ "Optimizer", "AdaMax", "InfinityNorm", "labml.ai" ], "description_md": "### 题目描述\n实现 Kingma & Ba (2015) 在 Adam 论文中提出的无穷范数变体 AdaMax 单步：更新一阶矩 $m_t = \\beta_1 m_{t-1} + (1 - \\beta_1) g_t$，无穷范数上界无需偏差修正，直接按逐元素最大值递推 $u_t = \\max(\\beta_2 u_{t-1},\\ |g_t|)$；对一阶矩做偏差修正 $\\hat{m}_t = m_t / (1 - \\beta_1^t)$，更新参数 $\\theta_t = \\theta_{t-1} - \\text{lr} \\cdot \\hat{m}_t / (u_t + \\epsilon)$，返回 `(theta_t, m_t, u_t)`。\n\n### 函数签名\n```python\ndef adamax_step(\n param: torch.Tensor,\n grad: torch.Tensor,\n m: torch.Tensor,\n u: torch.Tensor,\n step: int,\n lr: float = 2e-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]\ngrad = [0.5, -1.0]\nm = [0.0, 0.0]\nu = [0.0, 0.0]\nstep = 1\nlr = 2e-3\nbeta1 = 0.9\nbeta2 = 0.999\neps = 1e-8\n\n# 输出\n(\n [0.998, 2.002], # shape=[2], dtype=float32\n [0.05, -0.1], # shape=[2], dtype=float32\n [0.5, 1.0] # shape=[2], dtype=float32\n)\n```\n\n### 参数规格与约束\n- **`param`**：`torch.Tensor`，测试规格 `shape=[2], dtype=float32`\n- **`grad`**：`torch.Tensor`，测试规格 `shape=[2], dtype=float32`\n- **`m`**：`torch.Tensor`，测试规格 `shape=[2], dtype=float32`\n- **`u`**：`torch.Tensor`，测试规格 `shape=[2], dtype=float32`\n- **`step`**：`int`\n- **`lr`**：`float`（默认 `2e-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=[2], dtype=float32)`, `torch.Tensor(shape=[2], dtype=float32)`, `torch.Tensor(shape=[2], 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": "adamax_step", "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 adamax_step(\n param: torch.Tensor,\n grad: torch.Tensor,\n m: torch.Tensor,\n u: torch.Tensor,\n step: int,\n lr: float = 2e-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 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 adamax_step(param: torch.Tensor, grad: torch.Tensor, m: torch.Tensor, u: torch.Tensor, step: int, lr: float = 2e-3, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n m_t = beta1 * m + (1.0 - beta1) * grad\n u_t = torch.maximum(beta2 * u, grad.abs())\n m_hat = m_t / (1.0 - beta1 ** step)\n new_param = param - lr * m_hat / (u_t + eps)\n return new_param, m_t, u_t\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 (\"首步 step=1\", torch.tensor([1.0, 2.0]), torch.tensor([0.5, -1.0]), torch.zeros(2), torch.zeros(2), 1),\n (\"历史 u_prev 大于当前梯度\", torch.ones(3), torch.tensor([0.1, 0.1, 0.1]), torch.zeros(3), torch.tensor([2.0, 2.0, 2.0]), 5),\n (\"当前梯度刷新无穷范数上界\", torch.randn(4), torch.full((4,), 5.0), torch.zeros(4), torch.ones(4), 3),\n (\"2D 参数 step=10\", torch.randn(2, 4), torch.randn(2, 4), torch.randn(2, 4) * 0.1, torch.rand(2, 4), 10),\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": "首步 step=1", "input_desc": "真实参数契约用例 (param, grad, m, u, step, lr, beta1, beta2, eps)", "params": [ { "name": "param", "type": "torch.Tensor", "value": "torch.tensor([1.0, 2.0], dtype=torch.float32)" }, { "name": "grad", "type": "torch.Tensor", "value": "torch.tensor([0.5, -1.0], dtype=torch.float32)" }, { "name": "m", "type": "torch.Tensor", "value": "torch.tensor([0.0, 0.0], dtype=torch.float32)" }, { "name": "u", "type": "torch.Tensor", "value": "torch.tensor([0.0, 0.0], dtype=torch.float32)" }, { "name": "step", "type": "int", "value": "1" }, { "name": "lr", "type": "float", "value": "2e-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], dtype=torch.float32)", "grad": "torch.tensor([0.5, -1.0], dtype=torch.float32)", "m": "torch.tensor([0.0, 0.0], dtype=torch.float32)", "u": "torch.tensor([0.0, 0.0], dtype=torch.float32)", "step": "1", "lr": "2e-3", "beta1": "0.9", "beta2": "0.999", "eps": "1e-8" } }, { "name": "历史 u_prev 大于当前梯度", "input_desc": "真实参数契约用例 (param, grad, m, u, step, lr, beta1, beta2, eps)", "params": [ { "name": "param", "type": "torch.Tensor", "value": "torch.tensor([1.0, 1.0, 1.0], dtype=torch.float32)" }, { "name": "grad", "type": "torch.Tensor", "value": "torch.tensor([0.1, 0.1, 0.1], dtype=torch.float32)" }, { "name": "m", "type": "torch.Tensor", "value": "torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32)" }, { "name": "u", "type": "torch.Tensor", "value": "torch.tensor([2.0, 2.0, 2.0], dtype=torch.float32)" }, { "name": "step", "type": "int", "value": "5" }, { "name": "lr", "type": "float", "value": "2e-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, 1.0, 1.0], dtype=torch.float32)", "grad": "torch.tensor([0.1, 0.1, 0.1], dtype=torch.float32)", "m": "torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32)", "u": "torch.tensor([2.0, 2.0, 2.0], dtype=torch.float32)", "step": "5", "lr": "2e-3", "beta1": "0.9", "beta2": "0.999", "eps": "1e-8" } }, { "name": "当前梯度刷新无穷范数上界", "input_desc": "真实参数契约用例 (param, grad, m, u, step, lr, beta1, beta2, eps)", "params": [ { "name": "param", "type": "torch.Tensor", "value": "torch.tensor([0.3367, 0.1288, 0.2345, 0.2303], dtype=torch.float32)" }, { "name": "grad", "type": "torch.Tensor", "value": "torch.tensor([5.0, 5.0, 5.0, 5.0], dtype=torch.float32)" }, { "name": "m", "type": "torch.Tensor", "value": "torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float32)" }, { "name": "u", "type": "torch.Tensor", "value": "torch.tensor([1.0, 1.0, 1.0, 1.0], dtype=torch.float32)" }, { "name": "step", "type": "int", "value": "3" }, { "name": "lr", "type": "float", "value": "2e-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([0.3367, 0.1288, 0.2345, 0.2303], dtype=torch.float32)", "grad": "torch.tensor([5.0, 5.0, 5.0, 5.0], dtype=torch.float32)", "m": "torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float32)", "u": "torch.tensor([1.0, 1.0, 1.0, 1.0], dtype=torch.float32)", "step": "3", "lr": "2e-3", "beta1": "0.9", "beta2": "0.999", "eps": "1e-8" } } ], "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**AdaMax 无穷范数指数移动上界自适应优化器单步**（`adamax_step`）聚焦于 **Optim** 中的关键计算逻辑：\n\n实现 Kingma & Ba (2015) 在 Adam 论文中提出的无穷范数变体 AdaMax 单步：更新一阶矩 $m_t = \\beta_1 m_{t-1} + (1 - \\beta_1) g_t$，无穷范数上界无需偏差修正，直接按逐元素最大值递推 $u_t = \\max(\\beta_2 u_{t-1},\\ |g_t|)$；对一阶矩做偏差修正 $\\hat{m}_t = m_t / (1 - \\beta_1^t)$，更新参数 $\\theta_t = \\theta_{t-1} - \\text{lr} \\cdot \\hat{m}_t / (u_t + \\epsilon)$，返回 `(theta_t, m_t, u_t)`。\n\n在实现中，函数接收参数 `param` (`torch.Tensor`)、`grad` (`torch.Tensor`)、`m` (`torch.Tensor`)、`u` (`torch.Tensor`)、`step` (`int`)、`lr` (`float`)、`beta1` (`float`)、`beta2` (`float`)、`eps` (`float`)，目标是计算出类型与形状严格符合 `Tuple[torch.Tensor, torch.Tensor, torch.Tensor]` 契约的结果。\n\n#### 📐 输入输出张量变换与步骤拆解\n- **输入基准规格**：`param` 规格为 `torch.Tensor(shape=[2], dtype=float32)`，`grad` 规格为 `torch.Tensor(shape=[2], dtype=float32)`，`m` 规格为 `torch.Tensor(shape=[2], dtype=float32)`，`u` 规格为 `torch.Tensor(shape=[2], dtype=float32)`，`step` 规格为 `int` (示例值: `1`)。\n- **期望输出规格**：`Tuple` (`torch.Tensor(shape=[2], dtype=float32)`, `torch.Tensor(shape=[2], dtype=float32)`, `torch.Tensor(shape=[2], dtype=float32)`)。\n\n**具体推导与计算步骤如下：**\n1. 计算中间量 `m_t`：通过表达式 `beta1 * m + (1.0 - beta1) * grad` 完成张量或数值变换。\n2. 计算中间量 `u_t`：通过表达式 `torch.maximum(beta2 * u, grad.abs())` 完成张量或数值变换。\n3. 计算中间量 `m_hat`：通过表达式 `m_t / (1.0 - beta1 ** step)` 完成张量或数值变换。\n4. 计算中间量 `new_param`：通过表达式 `param - lr * m_hat / (u_t + eps)` 完成张量或数值变换。\n5. 返回最终结果：输出 `new_param, m_t, u_t`，严格对齐函数签名要求的返回类型与维度。\n\n#### 📊 复杂度分析\n- **时间复杂度**：$O(N)$（线性于输入张量总元素个数 $\\prod d_i$），所有运算均为向量化逐元素算子或归约计算。\n- **空间复杂度**：$O(N)$，需要与输入或输出张量同规模的显存空间保存计算结果。\n\n#### ⚠️ 关键实现细节与易错点\n- **除零保护与数值平滑 (`eps`)**：在执行分母归一化、标准差缩放或对数运算时，需加上 `eps` 或执行 `clamp_min` 防止除零产生 `NaN` / `Inf`。\n- **边界输入鲁棒性**：注意单样本（`Batch=1`）、零值输入或极端超参数下的计算行为，禁止硬编码特定测试用例常量。\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 adamax_step(param: torch.Tensor, grad: torch.Tensor, m: torch.Tensor, u: torch.Tensor, step: int, lr: float = 2e-3, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-8) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n m_t = beta1 * m + (1.0 - beta1) * grad\n u_t = torch.maximum(beta2 * u, grad.abs())\n m_hat = m_t / (1.0 - beta1 ** step)\n new_param = param - lr * m_hat / (u_t + eps)\n return new_param, m_t, u_t\n```" }