Spaces:
Running on Zero
Running on Zero
File size: 11,804 Bytes
5f7afef | 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 | {
"slug": "adaln-zero-diffusion-modulation",
"title": "DiT adaLN-Zero 自适应层归一化与残差门控",
"difficulty": "Medium",
"difficulty_breakdown": {
"concept": 3,
"implementation": 3,
"math_derivation": 3,
"boundary_complexity": 3
},
"category": "DL",
"tags": [
"DiT",
"Diffusion",
"Normalization",
"TorchLeet"
],
"description_md": "### 题目描述\n实现 Diffusion Transformer (DiT) 核心的自适应层归一化（adaLN-Zero）调制与残差门控运算。输入隐藏状态 `x`（`[B, S, D]`）以及条件嵌入映射输出的缩放系数 `scale`（`[B, D]`）、平移系数 `shift`（`[B, D]`）与残差门控因子 `gate`（`[B, D]`）。先对 `x` 进行标准 LayerNorm 归一化 $\\hat{x} = \\frac{x - \\mu}{\\sigma + \\epsilon}$，按维度广播调制为 $y = \\hat{x} \\odot (1 + \\gamma) + \\beta$，最后乘以门控因子 $y_{out} = y \\odot \\alpha$，返回形状 `[B, S, D]` 的张量。\n\n### 函数签名\n```python\ndef adaln_zero_modulate(\n x: torch.Tensor,\n scale: torch.Tensor,\n shift: torch.Tensor,\n gate: torch.Tensor,\n) -> torch.Tensor:\n```\n\n### 示例\n```python\n# 输入\nx = [[[1.0, 2.0], [3.0, 4.0]], [[-1.0, 0.5], [1.5, -0.5]]]\nscale = [[1.0, 2.0], [0.5, 1.5]]\nshift = [[0.5, -0.5], [1.0, 0.0]]\ngate = [[1.5, -0.5], [2.5, 0.0]]\n\n# 输出 (Tensor, shape=[2, 2, 2], dtype=float32)\n[[[-2.25, -1.25], [-2.25, -1.25]], [[-1.25, 0.0], [6.25, 0.0]]]\n```\n\n### 参数规格与约束\n- **`x`**：`torch.Tensor`，测试规格 `shape=[2, 4, 8], dtype=float32`\n- **`scale`**：`torch.Tensor`，测试规格 `shape=[2, 8], dtype=float32`\n- **`shift`**：`torch.Tensor`，测试规格 `shape=[2, 8], dtype=float32`\n- **`gate`**：`torch.Tensor`，测试规格 `shape=[2, 8], dtype=float32`\n- **返回值**：`torch.Tensor`，输出规格 `shape=[2, 4, 8], dtype=float32`\n- **精度容差**：`atol=0.0001, rtol=0.0001`\n> 题目来源：改编自 [TorchLeet](https://github.com/Exorust/TorchLeet)（MIT License），已改写为函数式评测接口与中文题面。",
"entry_point": "adaln_zero_modulate",
"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 adaln_zero_modulate(\n x: torch.Tensor,\n scale: torch.Tensor,\n shift: torch.Tensor,\n gate: torch.Tensor,\n) -> 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 adaln_zero_modulate(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:\n mean = x.mean(dim=-1, keepdim=True)\n var = x.var(dim=-1, keepdim=True, unbiased=False)\n x_norm = (x - mean) / torch.sqrt(var + 1e-6)\n modulated = x_norm * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)\n return modulated * gate.unsqueeze(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 (\"门控为 0 时输出全为 0 (Zero-init 特性)\", torch.randn(2, 4, 8), torch.randn(2, 8), torch.randn(2, 8), torch.zeros(2, 8)),\n (\"常数门控与调制\", torch.ones(2, 3, 4), torch.zeros(2, 4), torch.ones(2, 4), torch.ones(2, 4)),\n (\"批处理 B=4, S=16, D=32\", torch.randn(4, 16, 32), torch.randn(4, 32), torch.randn(4, 32), torch.randn(4, 32)),\n (\"B=1, S=8, D=16\", torch.randn(1, 8, 16), torch.randn(1, 16), torch.randn(1, 16), torch.randn(1, 16)),\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": "门控为 0 时输出全为 0 (Zero-init 特性)",
"input_desc": "真实参数契约用例 (x, scale, shift, gate)",
"params": [
{
"name": "x",
"type": "torch.Tensor",
"value": "torch.tensor([[[1.0, 2.0], [3.0, 4.0]], [[-1.0, 0.5], [1.5, -0.5]]], dtype=torch.float32)"
},
{
"name": "scale",
"type": "torch.Tensor",
"value": "torch.tensor([[1.0, 2.0], [0.5, 1.5]], dtype=torch.float32)"
},
{
"name": "shift",
"type": "torch.Tensor",
"value": "torch.tensor([[0.5, -0.5], [1.0, 0.0]], dtype=torch.float32)"
},
{
"name": "gate",
"type": "torch.Tensor",
"value": "torch.tensor([[1.5, -0.5], [2.5, 0.0]], dtype=torch.float32)"
}
],
"raw_inputs": {
"x": "torch.tensor([[[1.0, 2.0], [3.0, 4.0]], [[-1.0, 0.5], [1.5, -0.5]]], dtype=torch.float32)",
"scale": "torch.tensor([[1.0, 2.0], [0.5, 1.5]], dtype=torch.float32)",
"shift": "torch.tensor([[0.5, -0.5], [1.0, 0.0]], dtype=torch.float32)",
"gate": "torch.tensor([[1.5, -0.5], [2.5, 0.0]], dtype=torch.float32)"
}
},
{
"name": "常数门控与调制",
"input_desc": "真实参数契约用例 (x, scale, shift, gate)",
"params": [
{
"name": "x",
"type": "torch.Tensor",
"value": "torch.tensor([[[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]], [[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]], dtype=torch.float32)"
},
{
"name": "scale",
"type": "torch.Tensor",
"value": "torch.tensor([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]], dtype=torch.float32)"
},
{
"name": "shift",
"type": "torch.Tensor",
"value": "torch.tensor([[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]], dtype=torch.float32)"
},
{
"name": "gate",
"type": "torch.Tensor",
"value": "torch.tensor([[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]], dtype=torch.float32)"
}
],
"raw_inputs": {
"x": "torch.tensor([[[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]], [[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]], dtype=torch.float32)",
"scale": "torch.tensor([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]], dtype=torch.float32)",
"shift": "torch.tensor([[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]], dtype=torch.float32)",
"gate": "torch.tensor([[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]], dtype=torch.float32)"
}
},
{
"name": "批处理 B=4, S=16, D=32",
"input_desc": "真实参数契约用例 (x, scale, shift, gate)",
"params": [
{
"name": "x",
"type": "torch.Tensor",
"value": "torch.randn(4, 16, 32)"
},
{
"name": "scale",
"type": "torch.Tensor",
"value": "torch.randn(4, 32)"
},
{
"name": "shift",
"type": "torch.Tensor",
"value": "torch.randn(4, 32)"
},
{
"name": "gate",
"type": "torch.Tensor",
"value": "torch.randn(4, 32)"
}
],
"raw_inputs": {
"x": "torch.randn(4, 16, 32)",
"scale": "torch.randn(4, 32)",
"shift": "torch.randn(4, 32)",
"gate": "torch.randn(4, 32)"
}
}
],
"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**DiT adaLN-Zero 自适应层归一化与残差门控**（`adaln_zero_modulate`）聚焦于 **DL** 中的关键计算逻辑：\n\n实现 Diffusion Transformer (DiT) 核心的自适应层归一化（adaLN-Zero）调制与残差门控运算。输入隐藏状态 `x`（`[B, S, D]`）以及条件嵌入映射输出的缩放系数 `scale`（`[B, D]`）、平移系数 `shift`（`[B, D]`）与残差门控因子 `gate`（`[B, D]`）。先对 `x` 进行标准 LayerNorm 归一化 $\\hat{x} = \\frac{x - \\mu}{\\sigma + \\epsilon}$，按维度广播调制为 $y = \\hat{x} \\odot (1 + \\gamma) + \\beta$，最后乘以门控因子 $y_{out} = y \\odot \\alpha$，返回形状 `[B, S, D]` 的张量。\n\n在实现中，函数接收参数 `x` (`torch.Tensor`)、`scale` (`torch.Tensor`)、`shift` (`torch.Tensor`)、`gate` (`torch.Tensor`)，目标是计算出类型与形状严格符合 `torch.Tensor` 契约的结果。\n\n#### 📐 输入输出张量变换与步骤拆解\n- **输入基准规格**：`x` 规格为 `torch.Tensor(shape=[2, 4, 8], dtype=float32)`，`scale` 规格为 `torch.Tensor(shape=[2, 8], dtype=float32)`，`shift` 规格为 `torch.Tensor(shape=[2, 8], dtype=float32)`，`gate` 规格为 `torch.Tensor(shape=[2, 8], dtype=float32)`。\n- **期望输出规格**：`torch.Tensor(shape=[2, 4, 8], dtype=float32)`。\n\n**具体推导与计算步骤如下：**\n1. 计算中间量 `mean`：通过表达式 `x.mean(dim=-1, keepdim=True)` 完成张量或数值变换。\n2. 计算中间量 `var`：通过表达式 `x.var(dim=-1, keepdim=True, unbiased=False)` 完成张量或数值变换。\n3. 计算中间量 `x_norm`：通过表达式 `(x - mean) / torch.sqrt(var + 1e-6)` 完成张量或数值变换。\n4. 计算中间量 `modulated`：通过表达式 `x_norm * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)` 完成张量或数值变换。\n5. 返回最终结果：输出 `modulated * gate.unsqueeze(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 adaln_zero_modulate(x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:\n mean = x.mean(dim=-1, keepdim=True)\n var = x.var(dim=-1, keepdim=True, unbiased=False)\n x_norm = (x - mean) / torch.sqrt(var + 1e-6)\n modulated = x_norm * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)\n return modulated * gate.unsqueeze(1)\n```"
} |