codestar-backend / data /problems /adamax-infinity-norm-optimizer-step.json
github-actions
Sync backend from GitHub 47e8067
5f7afef
Raw History Blame Contribute Delete
14.2 kB
{
"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```"
}