Spaces:
Running on Zero
Running on Zero
File size: 14,199 Bytes
56bd691 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 | {
"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```"
} |