【轻量化复刻大模型 / KIMI K3】KDA 模块交互式演示
0x00 文件信息
注：
- 项目github.com/RossiLH/Mod… 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
- 涉及 KIMI K3，DeepSeek V4，PI0.5，PI0.7 等模型。后续可能会跟踪新模型。
- 全 CPU 复刻，不涉及 GPU，也不涉及大规模训练（这是另外一个庞杂深邃的领域）和硬件。
- 模型的组件尽量是原始大模型的缩小版，如果组件没有开源版本，则用近似组件替代。
- 会解释每个组件在解决一个什么问题、这个问题难在哪、为什么只能这么解、解法的原理是什么，以及组件之间如何咬合成一个整体。重点不是记名词，而是理解一条主线。
- 每个主要组件都有可直接运行的 Gradio demo——无需 GPU。另有若干可视化脚本（plot_*.py），可生成静态图。
- 所有文件行号基于本地源码快照；源码后续变化时，行号可能移动，但函数名和调用关系仍可作为定位入口。
- 为了保证每个文章可以尽量独立成文，所以有些内容会在不同文章中都存在，可能在系列层面上存在冗余，还请大家理解。
0x01 KDA核心数学
KDA (Kimi Delta Attention) 是 K3 的线性注意力机制，核心思想是用固定大小状态矩阵 + delta‑rule 纠错 + 有界衰减实现O(1) 状态、O(T) 并行的序列建模。
1.1 ShortConv: 因果 depthwise 短卷积
K3 的 MLA 层是 NoPE (无位置编码)，ShortConv 给每个通道注入局部位置信息：
out[t, c] = Σ_{lag=0}^{K‑1} w[c, θ, lag] · x[t − lag, c] + bias[c]
- 因果性：只看当前 + 过去 K‑1 个 token (左侧补零)
- 恒等初始化：weight[:, θ, ‑1] = 1, 其余为 0 → 初始为 no‑op, 不扰动 KDA 衰减动力学
- 单步 decode: 维护 K‑1 长度的滑动窗口, O(K) 逐 token 更新
1.2 SiTU‑GLU: 有界激活函数 (Eq. 12)
softcap(x, β) = β · tanh(x / β)
SiTU‑GLU(g, u) = softcap(g, β_g) · σ(g) · softcap(u, β_u)
- 有界性: |SiTU‑GLU| ≤ β_g · β_u = 4 × 25 = 100 < FP8_E4M3_MAX = 448
- 饱和区梯度: softcap 用 tanh 截断, 饱和区仍有非零梯度 (tanh' > 0), 不会像硬 clamp 那样掐断训练
- 对比 SwiGLU: SwiGLU(g, u) = silu(g) · u, 无界, 仅作消融基线
1.3 KDA 递推方程 (Eq. 1)
KDA 维护一个固定大小状态矩阵 S ∈ R^{Dk×Dv}, 逐 token 更新：
状态递推 (Delta Rule) :
S_t = Diag(a_t) · S_{t‑1} + β_t · k_t · (v_t − S_{t‑1}^T k_t)^T
读出：
o_t = S_t^T · q_t
其中：
- q, k ∈ R^{Dk}: 查询和键, L2 归一化
- v ∈ R^{Dv}: 值
- a_t = exp(g_t) ∈ (0, 1]: 逐 channel 衰减因子 (保留旧记忆)
- β_t ∈ (0, 1): 写入强度 (sigmoid 输出)
- v_t − S_{t‑1}^T k_t: delta‑rule 纠错项—只写入"当前检索与目标的误差", 而非盲目覆盖
delta‑rule 的直觉: 当同一 key 反复出现、value 逐渐改变时, S 中该 key 对应的检索值会收敛到最新 value (而非被覆盖或遗忘)。
1.4 有界衰减 (Eq. 5)
g = g_min · sigmoid(exp(A_h)) · z z = W_a^up · W_a^down · x + b_a
- g ∈ (g_min, 0), 严格有下界 → a = exp(g) ∈ (exp(g_min), 1)
- 为什么必须有界: 无界门控 (如 softplus) 在 chunkwise 并行中, 累积保留因子 1/γ = exp(‑Σg) 可能溢出到 NaN
- safe_chunk_size(g_min, dtype) = ⌊log(dtype_max / headroom) / (‑g_min)⌋, 保证数值安全
1.5 分块并行 (Eqs. 3‑4)
块内并行、块间串行, 用 UT 变换 (三角求解) 实现块内 O(T·D²) 并行：
块内: O = (Γ·Q)·S_θ + A·Ṽ A = Tril[(Q·Γ)(K/Γ)^T]
UT变换: Ṽ = (I+M)^{‑1}·Diag(β)·V M = tril(Diag(β)·K·K^T, −1)
块末:: S_C = Diag(γ^{<1→C>})·S_θ + Σ_r (γ^{<r+1→C>})·k_r·Ṽ_r^T
1.6 KCP 前缀扫描 (§5.1.2)
KDA 状态跨 segment 不可加 (S_{seg1} + S_{seg2} ≠ S_{full}), 必须用前缀扫描重组：
S_running = M_seg · S_running + S̃_seg
通信量固定 O(Dk·(Dk+Dv)), 与序列长度无关 (对比 softmax KV 的 O(T·Dk·Dv))。
1.7 KimiDeltaAttention 完整层
q, k, v = ShortConv(W_{q/k/v} · x)
q = L2Norm(SiLU(q)), k = L2Norm(SiLU(k)), v = SiLU(v)
g = g_min · sigmoid(exp(A_h) · (W_a^up · W_a^down · x + b_a)) # 有界衰减
β = sigmoid(W_β · x) # 写入强度
(o, S) = kda_recurrent / kda_chunkwise / kda_step # 核心递推
output = W_o · [sigmoid(W_g · x) ⊙ HeadRMSNorm(o)] # 输出门
1.8 3:1 KDA/MLA 调度
K3 每 4 层中 3 层用 KDA (线性注意力, O(1) 状态) + 1 层用 MLA (softmax 注意力, 压缩 latent):
层 0: KDA 层 1: MLA 层 2: KDA 层 3: MLA ...
KDA 层负责长程线性回忆 (固定状态), MLA 层负责精确近程检索 (KV cache)。
0x02 KDA 实现了哪些功能
2.1 覆盖的源模块 (5 个)
2.2 覆盖的测试（90 个测试函数）
2.3 15 个 Tab 功能清单
0x03 演示实现思路
3.1 整体架构
gradio_kda_demo.py
├── sys.path 配置（向上 5 层找到 src/）
├── 源模块导入（kda_ref / kda / short_conv / situ_glu / model）
├── 辅助函数（_make_kda_inputs / _make_delayed_task / _fit_conv）
├── 13 个 Tab 函数（每个返回 Tuple[Figure, str]）
└── build_demo() + __main__（Gradio 界面构建 + 启动）
3.2 核心设计决策
决策 1：所有函数返回 (fig, text) 元组
每个 Tab 函数的签名统一为 func(...) -> Tuple[plt.Figure, str]：
- fig：Matplotlib Figure 对象，由 Gradio 的 gr.Plot 组件原生渲染
- text：结论字符串，由 gr.Textbox 组件显示
这样做的好处：
- Gradio 6 的 gr.Plot.postprocess 自动处理 Figure→base64 转换
- 不需要手动 plt.savefig 或 fig2arr
- 函数可以独立测试（不依赖 Gradio 运行时）
决策 2：直接调用源模块，不复制代码
所有可视化函数直接 import 并调用 src/models/llm/kimi_k3/modules/ 下的源模块：
- kda_recurrent / kda_step / kda_chunkwise 提供真实计算
- KimiDeltaAttention 提供完整的生产层
- ShortConv / SiTUGLU 提供组件级验证
唯一"复制"的代码是 Tab 8 (v1 vs v2) 中的手写单步 KDA——因为 kda_ref.py 的 kda_step 只支持 v2，需要手动实现 v1 的 softplus 门控路径来对比。
决策 3：Gradio 独有实验内联定义
13 个 Tab 中有 7 个是 Gradio 独有的（不在现有 plot 脚本中）：
- Tab 5: KDA 原理 4 格——组合了 S 热图、单步拆解、delta rule、门控
- Tab 6: 递推 vs 分块——直接对比两种计算路径
- Tab 8: v1 vs v2 门控——手动实现两种门控做数值对比
- Tab 10: 模块形状——实时构造 Kimi3Block 并统计参数
- Tab 11: 单步匹配——逐 token 调用 kda_step 验证与递推一致
- Tab 12: 诊断指标——调用 build_kda_diagnostics 展示运行时指标
- Tab 13: 3:1 调度——实时构造 Kimi3Model 并可视化层调度
这些实验没有现成的 plot 脚本，因此直接在 gradio 文件中定义。
3.3 测试→Tab 映射逻辑
每个测试函数都被映射到最能展示其验证内容的 Tab：
映射原则：
- 形状/缓存测试 → Tab 10（模块形状与缓存）
- 数值等价性测试 → Tab 6（递推 vs 分块）或 Tab 11（单步匹配）
- 机制验证测试 → 对应原理 Tab（如 KCP→Tab 9，衰减→Tab 7）
- 模型集成测试 → Tab 13（3:1 调度）
- 梯度测试 → Tab 12（诊断指标）
示例映射：
test_shape_and_no_bias → Tab 1（感受野）
test_identity_initialization → Tab 2（恒等初始化）
test_causality → Tab 3（因果性）
test_output_is_bounded → Tab 4（激活曲线）
test_gradients_survive_saturation → Tab 5（饱和区梯度）
test_swiglu_baseline_explodes → Tab 6（离群点回归）
test_chunkwise_equals_recurrence → Tab 8（递推 vs 分块）
test_step_matches_sequence → Tab 13（单步匹配递推）
test_kcp_prefix_scan → Tab 11（KCP 前缀扫描）
test_module_forward_shapes → Tab 12（形状与缓存）
test_unbounded_decay_would_overflow → Tab 9（有界 vs 无界衰减）
test_kimi3_model → Tab 15（3:1 调度）
3.4 环境与路径
# 向上 5 层找到项目根目录（文件在 experiments/models/llm/kimi_k3/modules/）
_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
_p = _THIS_DIR
for _ in range(5):
_p = os.path.dirname(_p)
_SRC = os.path.join(_p, "src") # → Model-Viz/src/
sys.path.insert(0, _SRC)
注意：向上层数取决于文件深度。linear_attn 的 demo 只需 4 层，kimi_k3 需要 5 层。
0x04 每个 Tab 的使用说明
Tab 1: ShortConv 感受野 vs 延迟误差
目的 验证 ShortConv 的局部感受野不足以捕获长程依赖 (论文 §2.1.1)。当 kernel_size < delay 时, MSE ≈ 信号方差 (任务失败)。
实现：
- 生成"延迟线"任务：输入是随机正弦信号，目标是延迟 D 步后的副本
- 训练不同 kernel 宽度（2, 4, 8, 12, 20）的 ShortConv 拟合该任务
- 绘制 kernel 宽度 vs 最终 MSE 的折线图
- 标注 delay=8 的垂直线——kernel <= delay 时 MSE ≈ 信号方差（失败）
关键代码路径：
_make_delayed_task() → 生成 (x, y)
_fit_conv() → Adam 训练 3000 步，返回最终 MSE
tab_receptive_field() → 绘图 + 返回结论
UI 组件与参数影响
输出图解读
tab-1
- 红色折线图: X 轴 = kernel_size, Y 轴 = 延迟任务重建 MSE (log scale)
- 灰色虚线: delay=8 的位置 —kernel ≤ delay 时 MSE 高 (无法捕获)
- 绿色方块: 当前滑块选择的 kernel_size 对应的 MSE
- 灰色水平点线: MSE=1.0 (信号方差基线)
- 你该看什么: kernel 从 2 增大到 20 时, MSE 应先高后低, 在 kernel ≈ delay 处有拐点
进一步解释：
拐点不是几何拐点，而是对数轴上 MSE 下降速率变化最大处。实测 MSE 序列 1.5e-2, 9.5e-3, 1.9e-3, 2.5e-4, 8.8e-5（kernel = 2, 4, 8, 12, 20），相邻下降倍数 1.6x, 4.9x, 7.7x, 2.8x → 最大下降发生在 kernel=8→12 区间，故拐点在 kernel=12。
为什么重要
ShortConv 是 K3 中 NoPE 位置编码的替代方案, kernel_size 决定了它能捕获多长的局部模式。此 Tab 量化展示"局部感受野的局限"—delay=8 的任务需要 kernel 至少 9 才能重建。
Tab 2: ShortConv 恒等初始化
目的：展示恒等初始化 (conv 初始为 no‑op) 不扰动 KDA 衰减动力学, 训练起步更平稳。
实现：
- 恒等初始化：reset_parameters 将 weight 最后一个 tap 设为 1，其余为 0
- 随机初始化：weight.normal_(std=0.5)
- 两种初始化各训练 1500 步，记录每 100 步的 MSE
- 绘制训练曲线对比
输出图解读
tab-2
- 绿色圆点线: 恒等初始化训练曲线 (初始为 no‑op, 起步 MSE 低)
- 红色方块线: 随机初始化训练曲线 (初始扰动大, 起步更抖)
- Y 轴 log scale: 让训练后期的细微差异可见
为什么重要
KDA 的衰减动力学依赖 ShortConv 不在初始阶段引入大幅扰动。恒等初始化保证 conv(x) ≈ x, 不干扰 KDA 的状态矩阵 S 的收敛轨迹。
Tab 3: ShortConv 因果性
目的
验证 ShortConv 只看当前+过去 token, 修改未来输入不影响过去输出 (论文 §2.1.1)。
tab-3
输出图解读
- 蓝色圆点线: 原始输出
- 橙色方块线: 未来被扰动后的输出 (t≥10 处输入 +100)
- 灰色虚线: 扰动起点 (第 10 位)
- 你该看什么: 两条线在 t<10 处几乎完全重合 (因果性成立), t≥10 处分开 (扰动生效)
进一步解释：
- 扰动 = 把输入 x 在 t=10..19 这 10 个时间步上整体 +100（代码 xp[:, 10:] += 100.0），前 10 步保持原样。
- 扰动起点 = 被修改的第一个时间步索引，即 t=10。图中 axvline(9) 画在索引 9（未扰动段右边界）作为视觉分隔线。文字已重写说明。
为什么重要
因果性是自回归生成的硬性要求—decode 时未来 token 尚未生成, 卷积绝不能偷看。此 Tab 直接可视化"修改未来不影响过去"。
Tab 4: SiTU‑GLU 激活曲线对比
目的
展示 SiTU‑GLU 的有界性 (|f(x)| ≤ β_g·β_u = 100), 对比 SwiGLU 的无界增长。
实现：
- 在 [x_min, x_max] 范围内计算 SiTU-GLU 和 SwiGLU 的激活值
- 绘制两条曲线 + 上界水平线
- 输出最大绝对值对比
UI 组件与参数影响
输出图解读
tab-4
- 红色线: SwiGLU (无界, |x| 大时线性增长)
- 绿色线: SiTU‑GLU (有界, 被 ±100 水平线截住)
- 灰色水平虚线: 上界 ±100
调节建议
- 想看 SwiGLU 溢出幅度 → 增大 x_max 到 40, SwiGLU 会远超 100
- 想看 SiTU‑GLU 饱和平滑性 → 增大 x_max 到 50, SiTU‑GLU 平稳趋近 100
Tab 5: SiTU‑GLU 饱和区梯度
目的
展示 softcap 在饱和区保持非零梯度 (优于硬 clamp), 不会掐断训练。
实现：
- 计算 softcap(z, β) 的梯度 d/dz
- 计算 clamp(z, max=β) 的梯度 d/dz
- 对比：softcap 梯度在饱和区 > 0，硬 clamp 梯度 = 0
输出图解读
tab-5
- 绿色线: softcap 梯度 d/dz (饱和区仍 > 0, 光滑衰减)
- 红色线: 硬 clamp 梯度 d/dz (饱和区 = 0, 梯度消失)
- 你该看什么: |z| > β=4 时, softcap 梯度虽小但非零, clamp 梯度精确为 0
为什么重要
FP8 训练中激活值经常进入饱和区。如果梯度在那里消失 (硬 clamp), 对应参数就停止学习。softcap 的 tanh 截断保证梯度永远非零, 饱和区的参数仍能被微调。
Tab 6: SiTU‑GLU 离群点回归
目的
展示 SiTU‑GLU 对大离群激活的保守行为—有界截断使 loss 略高但数值稳定, SwiGLU 无界通过但低精度下溢出。
输出图解读
tab-6
左图 MSE 曲线，右图离群样本预测散点（直观展示 SwiGLU 过冲、SiTU-GLU 光滑截平）。
右图直观展示：
- SwiGLU（红圆） ：预测范围 [21, 178]，震荡剧烈——同一个真实 y≈200，预测从 21 到 178 跨度 157
- SiTU-GLU（绿方） ：预测范围 [75, 194]，更集中在 100 附近（截平线的存在让它把"做不到的"压缩到上界）
虽然两者都拟合不到 200，但 SiTU-GLU 的预测方差更小，离 y=x 这条完美线整体偏差反而更近。
右图亮点：直观展示 SwiGLU 预测震荡剧烈（[21,178]，跨度 157），SiTU-GLU 预测更集中（[75,194]，跨度 119），且整体离 y=x 完美线更近。绿色截平线 beta_g*beta_u=100 也画出来了。
你该看什么：左图中，绿色（SiTU-GLU）曲线在红色（SwiGLU）下方——SiTU-GLU 的 MSE 反而更低。这不是 bug，而是 SiTU-GLU 设计的真实优势：
数学事实
- SiTU-GLU 的上界是 beta_g × beta_u = 4 × 25 = 100（不是 25），所以 200 离群点确实超出范围（被截到 ~100）
- 但 SwiGLU 无界，会去"硬扛" 200 离群点，梯度被离群点拉爆，连带把 98% 正常样本（值域 ±3）的拟合也带歪
- SiTU-GLU 在 100 处光滑截平（tanh' 在饱和区仍非零），相当于对离群点自动做梯度裁剪，让优化器专注拟合正常样本
数值分解（训练 1800 步后）
SiTU-GLU 赢在两头：
- 离群点拟合误差更小（6168 vs 7647）：因为它只到 100，反而离 200 比 SwiGLU 的"过冲 + 震荡"更近
- 正常样本拟合误差几乎相同（124 vs 120）：截平没伤害正常区间
物理意义
- FP8 安全：SiTU-GLU 激活 ≤100 < FP8 E4M3 上限 448，不会溢出
- 鲁棒性：自动抑制离群点对梯度的主导，等价于内置 gradient clipping
- 不是"有界的代价" ：在这个任务上 SiTU-GLU 反而更优，因为它防止了 SwiGLU 的"离群点过拟合 + 正常样本被牺牲"
为什么重要
真实训练中 2% 的离群激活可能把 SwiGLU 推到 FP8 溢出 (>448), 导致 NaN。SiTU‑GLU 用 softcap 截断到 100 以内, 牺牲少量拟合精度换取数值安全。
Tab 7: KDA 原理 (4格)
目的
展示 KDA 的三个核心概念: 固定大小记忆 (S 矩阵)、delta‑rule 纠错、有界衰减门控。
实现（2×2 布局）：
- 左上：调用 kda_recurrent 拿最终状态 S，绘制热图
- 右上：手动构造 t=0 状态，展示单步写入量 β·k·(v - Sᵀk)ᵀ
- 左下：构造"同一 key 反复出现、v 逐渐改变"的场景，展示检索值收敛
- 右下：绘制 g = g_min·sigmoid(z) 的门控曲线
UI 组件与参数影响
输出图解读
tab-7
2×2 布局：
左上: 状态矩阵 S 热图
- RdBu_r 色图, 展示 S ∈ R^{Dk×Dv} 的最终状态
- 你该看什么: S 是固定大小 (不随 T 增长), 这是 KDA O(1) 状态的核心
右上 —— 单步三件事拆解：每来一个 token，状态只做三件事：
- 衰减 Diag(alpha) S：旧记忆先按通道乘一个 alpha=exp(g)<=1 的保留因子（旧信息"忘掉一点"）；
- 检索 S^T k：用当前 key 在记忆里"查一下"已存了什么；
- 写入 beta * k (v - S^T k)^T：只把误差 v - 检索到的记忆 写进状态（delta rule 的核心）。
左下: delta rule 纠错
- 同一 key 反复出现、v 逐渐改变时, 检索值的收敛轨迹
- 你该看什么: 检索值逐步收敛到最新 v, 而非被覆盖或遗忘
右下: 有界 log‑decay 门控
- g = g_min·sigmoid(z) 的曲线, 永远在 (g_min, 0) 之间
- 红色虚线: 下界 g_min
- 你该看什么: g 严格有下界 → a = exp(g) 有正下界 → chunkwise 不会溢出
Tab 8: 递推 vs 分块等价
目的
验证 Eq.1 (逐 token 递推) 与 Eq.3‑4 (块内并行) 逐位一致—两种计算路径的数值等价性。
实现：
- 同一组输入分别送入 kda_recurrent 和 kda_chunkwise
- 对比输出和状态的 max abs 误差
- 绘制输出对比曲线、逐位误差柱状图、状态差异热图
UI 组件与参数影响
输出图解读
tab-8
1×3 布局：
左: 输出逐位对比
- 蓝色圆点 = 递推输出, 橙色方块 = 分块输出
- 理想情况: 两条线完全重合
中: 逐位误差柱状图
- 绿色柱 = 每个时间步的 abs 误差 (Y 轴 log scale)
- 标题显示 max 误差 (应 < 1e‑10)
- 你该看什么: 误差在机器精度级别 (~1e‑15), 说明两种计算路径逐位一致
右: 状态差异热图
- hot 色图展示 |S_rec − S_chk|的逐元素差异
- 理想情况: 全黑 (差异为 0)
进一步解释：三张图都是同一份误差（量级都在 1e-16，fp64 机器精度）：
- 左图：直接叠加原始曲线，1e-16 远小于线宽 → 视觉重合（不是没误差，是看不见）
- 中图：把每个时间步的 |递推-分块|用 log 轴画，能看到 1e-16~1e-18 浮点噪声
- 右图：把状态矩阵 S (16×16) 每个元素的 abs 差画成热图，看误差的空间分布
为什么重要
递推是 ground truth (O(T) 串行), 分块是生产路径 (块内并行)。如果两者不等价, 生产推理结果就是错的。此 Tab 是 KDA 实现正确性的主锚点。
Tab 9: 有界 vs 无界衰减
目的 展示无界 softplus 门控在 chunkwise 计算中会溢出到 NaN, 有界 sigmoid 门控保证数值安全 (论文 §4.2.7)。
实现（2×2 布局）：
- 左上：理论 1/gamma 上界 vs tile 大小（对数坐标）
- 右上：tile=16 实测——有界有限，无界 NaN
- 左下：safe_chunk_size vs g_min 曲线
- 右下：chunk=64 失败 case
输出图解读
tab-9
2×2 布局：
左上： 1/γ 上界 vs tile 大小 (log-log)
- 有界版： 1/γ ≤ exp(-g_min·tile)，数据无关硬上界
- 红色虚线: BF16 上限 (3.4e38)
- 灰色点线: 论文 tile=16
- 关键: tile=16 时 1/γ≈1e35，刚好低于 BF16 上限；tile=32 时约 1e68，远超 BF16 上限
右上： tile=16 实测 (fp32)
- 固定 chunk=16，换门控 → 测试"门控是否决定溢出"。蓝色 = 输出有限，红色 = 输出 NaN（柱顶白字标 NaN）。softplus z+6 时 1/γ→inf，红柱。
- bounded(K3) 1/γ≈6.1e22 有限； softplus(z+6) 1/γ→inf → 输出 NaN
- 用 fp32 是因为 bf16 值域上限与 fp32 几乎相同（都是 e38 量级），溢出行为一致； fp32 精度更高，能更清楚地看到 1/γ 从有限到 inf 的跃迁
左下： safe_chunk_size vs g_min
- 圆点实线 = bf16, 方块虚线 = fp32
- 两条曲线完全重合——因为 bf16.max≈3.39e38 与 fp32.max≈3.40e38 几乎相同（指数都是 38）
- bf16 和 fp32 的真正差别在精度（尾数 7 位 vs 23 位），不在值域上限； 所以"bf16 更早溢出"是误解
右下： bounded chunk=16 vs 无界 chunk=64 (fp32)
- bounded 用 chunk=16 vs 无界 softplus 用更大 chunk=64 → 测试"无界门控调大 chunk 能不能救"。即使 chunk=64 也无济于事（红柱）。
- 无界 softplus 即使 chunk=64: 1/γ→inf → NaN
- 有界衰减的关键不是"避免 dtype 上限"， 而是让 1/γ 有数据无关的硬上界——这是 chunkwise 能稠密化的前提
为什么重要 这是 K3 选择 sigmoid 门控而非 softplus 的根本原因—数值安全性。无界门控在 chunkwise 并行中, 累积保留因子 1/γ 可能溢出, 导致 NaN 训练崩溃。
Tab 10: v1 vs v2 门控
目的 对比 KDA v1 (softplus 无界) 和 v2 (sigmoid 有界) 的门控行为和数值差异。
实现：
- 绘制两条门控曲线 g(z)
- 绘制 alpha 随 A 变化的范围
- 手动实现单步 KDA（因为 kda_step 只支持 v2），逐 token 跑 v1/v2 对比
- 输出状态范数和数值差异
为什么需要手写单步：kda_ref.py 的 kda_step 内部硬
本文由 AI图鉴自动抓取并翻译，版权归原作者所有；如需查看原文请点下方链接。