2. 架构

2.1. 概述

DeepSeek-V4.1-Flash 是一个多模态 MoE Transformer，以图像和文本为输入，并以自回归方式生成文本。其语言主干包含 40 个因果 Transformer 层，由 20 层因果编码器和后续 20 层解码器组成。除前两层仅使用 SWA 外，每一层均同时包含全局注意力和滑动窗口注意力。视觉编码器和 MLP 投影器将图像转换为视觉嵌入，并与文本嵌入共同处理；多模态数据从语言模型预训练一开始即被纳入训练。

总体上，DeepSeek-V4.1-Flash 包含 552B 主干参数和 196B Engram 参数；预填充阶段每个 token 激活 8B 参数，解码阶段激活 16B 参数。图 3 给出了其总体架构。

因果编码器-解码器（CED）架构和 Compressed Sparse Attention 2（CSA2）分别针对长上下文推理中的不同成本来源进行优化。

CED 从编码器输出构建解码器的全局键值（KV）缓存，使大多数提示 token 可以绕过完整的解码器计算，同时保留逐层局部滑动窗口注意力。这几乎将预填充计算量减半，从而降低智能体工作负载中处理新增或未缓存输入的成本。CSA2 则跨层共享全局 KV 以降低缓存存储，并复用稀疏选择结果以减少索引开销。在解码器中，分层稀疏索引器将后续索引器的搜索范围限制到前序索引器选出的候选池，从而进一步减少每个查询需要打分的条目数。

我们保留 DeepSeekMoE（Dai et al., 2024）的共享专家和细粒度路由专家，并针对图像与文本 token 引入模态特定的负载均衡（Wang et al., 2024a）。Single-Pass mHC（Xie et al., 2026）修改残差流混合方式，使内核融合更加高效；Engram（Cheng et al., 2026b）加入稀疏访问的条件记忆。主干预训练期间不使用 MTP 模块，而改用 DSpark（Cheng et al., 2026a）进行推测解码。DSpark 在主干预训练结束后单独训练。此外，我们将主 KV 缓存压缩为 FP4，以进一步降低存储开销。以下各节将介绍这些组件及相应优化。

2.1.1. 多模态架构

多模态输入路径由视觉编码器和 MLP 投影器组成。对于每张输入图像，视觉编码器会生成一个空间视觉特征网格。随后执行 $3 \times 3$ pixel-unshuffle 操作，将每个局部邻域沿通道维重新排列，在降低空间分辨率之后，由 MLP 投影器将特征映射到语言主干的隐藏维度。最终得到的视觉嵌入被插入输入嵌入序列中对应的图像 token 位置，并与文本嵌入一起由语言主干处理。

DeepSeek-ViT。 我们从头训练名为 DeepSeek-ViT 的视觉编码器，使其原生支持不同分辨率的图像。DeepSeek-ViT 以 Vision Transformer（Dosovitskiy et al., 2021）为基础，并做出若干修改。为适配任意分辨率输入，我们用 2D-RoPE 替换标准绝对位置嵌入。为使 ViT 更接近 LLM 的设计原则，我们将 patch embedding 层中的卷积替换为线性投影，以兼容 Muon 优化器。归一化采用 RMSNorm（Zhang and Sennrich, 2019），激活函数采用 SwiGLU（Shazeer, 2020）。在视觉特征输入 LLM 之前，我们使用 $3 \times 3$ 下采样的 pixel-unshuffle，将视觉 token 数减少 9 倍，从而有效支持最高约 $1344 \times 1344$ 像素的输入分辨率。

面向 MoE 的多模态无辅助损失负载均衡。 图像 token 与文本 token 的表示分布不同，因此可能产生不同的专家路由偏好。如果只对二者的总体负载进行均衡，可能掩盖特定模态内部的不平衡。为解决这一问题，我们扩展无辅助损失负载均衡（Wang et al., 2024a），分别为文本和图像 token 维护按专家划分的校正偏置。路由时，每个 token 使用与其模态对应的校正偏置进行专家选择，而原始路由分数仍用于加权所选专家的输出。每个训练步骤后，两组偏置根据各自的专家负载独立更新。该设计可以在各模态内部平衡专家利用率，并有助于实现稳定、高效的多模态训练。

2.2. 因果编码器-解码器（CED）

在智能体工作流中，频繁的工具调用会产生大量预填充请求；一旦 KV 缓存未命中，就会带来很高的计算开销。为缓解这一预填充瓶颈，我们提出受 YoCo（Sun et al., 2024）启发的因果编码器-解码器（CED）架构。YoCo 允许上半部分网络层直接共享下半部分网络层产生的 KV 缓存，从而降低预填充计算。CED 在此基础上进行一系列结构改进，同时增强总体 KV 缓存容量和 KV 生成的计算深度。因此，在保持与基线相当性能的同时，CED 能够将预填充计算量降低近一半。

对于全局注意力，CED 将 Transformer 的底部 $L/2$ 层视为因果编码器。对于上半部分网络层，即解码器中的 $l>L/2$，KV 条目不再由各层自身的隐藏状态 $H_l$ 产生，而是使用逐层独立的投影权重 $W_{KV}^{l}$ 与 $W_Z^l$，直接从第 $L/2$ 层的隐藏状态 $H_{L/2}$ 投影：

$$C_l = H_{L/2}W_{KV}^{l}, \qquad Z_l = H_{L/2}W_Z^{l}, \qquad l>\frac{L}{2}. \tag{1}$$

其中，$C$ 和 $Z$ 分别表示 KV 条目及其对应的压缩权重。借助这一设计，CED 在预填充阶段只需计算前半部分网络层，即可用极低计算成本获得上层的全局 KV 缓存。

对于滑动窗口注意力（SWA），CED 在所有层上继续维持传统的逐层计算。具体而言，对于任意第 $l$ 层，局部 key 和 value 直接由该层当前隐藏状态 $H_l$ 得到。这一设计有效增加了局部 KV 生成的计算深度，但维持逐层计算意味着必须执行 SWA replay。在预填充阶段，为解码器计算 SWA KV 缓存需要额外处理 $n_{\text{win}}\times L/2$ 个 token，其中 $n_{\text{win}}$ 表示窗口大小。对于每轮提示较短的多轮交互，这部分解码器计算开销不可忽略。

已有工作（Chen et al., 2025）表明，SWA 的实际有效感受野远小于理论上的 $n_{\text{win}}\times L/2$。受此启发，我们提出 Decoder SWA Bounded Replay：SWA 计算只对提示最后的 $n_{\text{win}}$ 个 token 进行预填充，从而显著降低计算成本。详细内容见第 3.2.2 节。

总体而言，对于序列长度 $N\gg n_{\text{win}}$，CED 将预填充复杂度从

$$O(NL)$$

降低为

$$O\left(\frac{NL}{2}+n_{\text{win}}\times\frac{L}{2}\right) \approx O\left(\frac{NL}{2}\right),$$

即总体计算量近似减半。

2.3. Compressed Sparse Attention 2（CSA2）

服务长上下文必须同时控制 KV 缓存存储与注意力计算成本。这些成本可以沿三个相乘的维度降低：条目大小维度，其中 GQA（Ainslie et al., 2023）减少 KV head 数量，MLA（DeepSeek-AI, 2024）在多个 head 之间共享一个较小的 latent；序列维度，其中每 $m$ 个 token 被压缩为一个条目，例如 DeepSeek-V4（DeepSeek-AI, 2026b）中的 CSA 和 HCA；以及层维度，其中部分层复用其他层的缓存（Brandon et al., 2024）和选择结果，而不是保留自身的缓存，或者干脆由效率更高的层替代。

已有工作证明沿层维度压缩是有效的：IndexCache（Bai et al., 2026）跨层复用 Top-K 索引，以减少索引器计算；YOIO（Sun et al., 2026b）只计算一次稀疏路由，并在所有层间共享；HySparse（Gao et al., 2026）则允许稀疏层复用稠密层的 KV 缓存。不过，单纯索引复用无法节省主 KV 存储；全网络共享路由会限制性能；混合架构仍保留完整注意力层。更重要的是，这些方法都没有同时覆盖上述三个乘法维度。

原始图 4

图 4｜CSA2 的三种运行模式。 三种模式在获取主 KV、索引器 K 和 Top-K 索引的方式上不同。绿色模块表示当前层计算得到的量；黄色模块表示从最近的 Full Mode 层复用的主 KV 与索引器 K；红色模块表示从最近一个产生索引的层（Full 或 Reindex Mode）复用的 Top-K 索引。三种模式都会在当前层计算主 Q 和 SWA KV。

CSA2 联合利用上述三个维度：跨层共享主 KV 与索引器 K，并允许不同层复用 Top-K 索引，同时将缓存共享与索引复用解耦。它还将这些复用策略与简化的压缩器以及分层稀疏索引器结合；后者在解码器中缩小后续索引层的搜索空间。

与 CSA 类似，CSA2 包含轻量索引器，使用索引器 Q 和索引器 K 对主 KV 条目打分，并为每个查询选择 Top-K 条目。每个 Q 同时关注这些被选条目和该层本地的滑动窗口 KV（SWA KV）。CSA2 还将未压缩的主 KV 作为压缩率为 1 的特殊情况。

同时，CSA2 简化了压缩器和索引器。在 CSA 中，压缩率为 $m$ 时，每个主 KV 条目由原始 KV 缓存中的 $2m$ 个条目生成，相邻压缩条目的来源存在重叠；压缩过程中还使用绝对位置嵌入来编码这 $2m$ 个条目的位置。CSA2 移除了这种重叠和绝对位置嵌入。此外，CSA2 通过对主 KV 条目进行投影得到索引器 K，替代 CSA 中从隐藏状态单独执行的压缩路径。这两项设计都简化了实现并提高训练效率。

第 2.3.1 与 2.3.2 节将分别介绍跨层复用策略和分层稀疏索引器。

2.3.1. 跨层 KV 与索引复用

每个 CSA2 层都会被静态指定为 Full、Reindex 或 Reuse 三种模式之一。在三种模式中，每层都计算自身的 query 和 SWA KV，并将其与选中的主 KV 条目共同用于生成新的注意力输出。区别在于主 KV、索引器 K 与 Top-K 索引的来源。图 4 展示了三种模式。

Full Mode。 当前层计算自身主 KV 与索引器 Q，再由主 KV 投影得到索引器 K，并运行索引器生成新的 Top-K 索引。因此，它执行完整的 CSA2 计算路径，其组件职责与 DeepSeek-V4 中完整的 CSA 层相同。

原始图 5

图 5｜分层稀疏索引器。 每个方格代表一个位置；绿色方格表示被选中的索引，蓝色矩形表示依据块内最大索引器分数选中的块。解码器中第一个处于 Full Mode 的 CSA2 层选择自身的 Top-512 索引，并根据所选块构建共享候选池，供后续层使用。处于 Reindex Mode 的 CSA2 层随后从该候选池中选择各自的 Top-512 索引。

Reindex Mode。 当前层复用前序层最近可用的主 KV 及其对应的索引器 K。索引器计算自身 query，对复用的 key 重新打分，并生成新的 Top-K 索引。由此，主 KV 与索引器 K 可以跨层共享，同时稀疏选择仍可在不同层发生变化。

Reuse Mode。 当前层复用最近可用的主 KV，以及前序 Full 或 Reindex Mode 层针对该主 KV 计算出的最新 Top-K 索引。该层直接利用这一选择执行注意力计算，无需计算索引器 Q，也无需评估索引分数。

共享主 KV 与索引器 K 能减少缓存存储，而复用 Top-K 索引则避免额外的索引器计算。Reindex Mode 在保持缓存共享的同时，允许不同层改变所选条目。当 CSA2 与 CED 结合时，解码器中指定为 Full Mode 的层从第 $L/2$ 层（即因果编码器最后一层）的隐藏状态计算自身全局 KV。Reindex 和 Reuse Mode 保持不变。

2.3.2. 分层稀疏索引器

跨层索引复用减少了索引器执行次数，但剩余索引器仍需对因果可见的完整上下文进行打分。对于极长上下文，这一成本依然是主要计算瓶颈。此前工作通过先对池化后的块表示打分并裁剪，再进行 token 级索引，从而引入索引器稀疏性（Xu et al., 2026b）。

我们发现，在解码器中，较浅层索引器的信息天然可用于限制较深层索引器的候选范围，而且无需增加额外状态。因此，我们提出 Hierarchical Sparse Indexer，仅用于 CED 的解码器，以减少解码时的重复打分。对于每个查询，第一个被指定为 Full Mode 的层构建候选池，之后需要重新索引的层将该候选池作为搜索域。若候选池大小固定，则更深层索引器的单查询成本会从随上下文长度线性增长变为常数级。

该机制在训练中是显式感知的，并在后训练阶段引入：训练与推理时采用完全相同的候选范围限制，因此较深层索引器会在与实际推理相同的搜索域下得到优化。图 5 展示了这一过程。

第一个 Full Mode 层会对所有因果可见的主 KV 位置打分，并生成自身注意力使用的 Top-K 索引。同时，它还进行块级候选选择：每个块的分数取该块内各位置索引分数的最大值，然后选择分数最高的若干块。接着，将这些块覆盖的位置汇总成一个大于最终 Top-K 集合的候选池。例如，选择 2,048 个块、每块包含 8 个位置，即得到 16,384 个候选位置。该候选池规定后续索引器的搜索范围；最终 Top-K 选择则决定每层真正读取哪些主 KV 条目。

后续 Reindex Mode 层只对与相应 query 对应的候选位置打分，并在候选池中选择自身 Top-K 条目。Reuse Mode 层不执行新的索引，而是使用针对其所复用主 KV 已计算出的最新 Top-K 索引。因此，候选池在多个索引层间共享，而各层最终选择可以不同。

在候选池大小固定时，后续每个索引器对每个查询需要打分的位置数与上下文长度无关并有固定上界。第一个 Full Mode 层仍会扫描完整因果可见范围。因此，分层索引降低了后续索引器评估成本，同时保留初始的全范围扫描。

2.4. 高效架构扩展

2.4.1. Single-Pass mHC

在 DeepSeek-V4 中，我们提出 mHC（Xie et al., 2026），在相邻 Transformer block 之间维护 $n$ 条残差流。对每个 token，将这些残差流记为

$$X_l \in \mathbb{R}^{n\times d},$$

其中 $l$ 为 block 索引，$d$ 为隐藏维度。其更新方式为：

$$X_{l+1}=B_lX_l+C_lF_l(A_lX_l),\qquad (A_l,B_l,C_l)=H(X_l). \tag{2}$$

其中 $A_l\in\mathbb{R}^{1\times n}$、$C_l\in\mathbb{R}^{n\times 1}$、$B_l\in\mathbb{R}^{n\times n}$，这些均为由 $X_l$ 预测的逐 token 系数。系数预测器 $H$ 包括归一化与投影。

理想情况下，相邻两个 block 之间的残差变换应是从 $(X_{l-1},Y_{l-1})$ 到 $(X_l,\hat X_l)$ 的单一映射，其中

$$\hat X_l=A_lX_l,\qquad Y_{l-1}=F_{l-1}(\hat X_{l-1}).$$

这样的映射需要读取 $(n+1)d$ 个值、写入 $(n+1)d$ 个值，因此激活内存流量的下界为 $(2n+2)d$。

实际上，DeepSeek-V4 对式 (2) 采用多遍实现。由于数据依赖，三个内核必须顺序执行：

$$X_l=B_{l-1}X_{l-1}+C_{l-1}Y_{l-1} \qquad\text{残差更新，对 }n\text{ 维收缩} \tag{3}$$ $$(A_l,B_l,C_l)=H(X_l) \qquad\text{系数计算，对 }nd\text{ 维收缩} \tag{4}$$ $$\hat X_l=A_lX_l \qquad\text{输入混合，对 }n\text{ 维收缩}. \tag{5}$$

三个内核分别读取 $(n+1)d$、$nd$ 和 $nd$ 个值，总共写入 $(n+1)d$ 个值。若再计入 $F_l$ 内的 pre-norm，总激活内存流量为 $(4n+4)d$，是理论下界的两倍。

在 $H$ 中，归一化权重可离线折叠到投影权重中，而 RMS 除法可在投影完成后执行。因此，三个阶段中的两个可以共享一次残差遍历：残差更新不需要沿隐藏维进行归约，所以可以逐 tile 计算 $X_l$，并立即将其用于累计投影输出以及计算 RMS 所需的平方和。

输入混合无法融合进上述遍历，因为 $A_l$ 只有在全部隐藏维 tile 的归约完成后才可获得，因此需要再次读取 $X_l$。第二遍读取还可同时完成输入 pre-norm。该两遍实现总共需要 $(3n+2)d$ 的激活读写量，比理论下界多一次对 $X_l$ 的读取。

因此，我们提出 Single-Pass mHC：将输入混合系数向后错开一个 block，即每个 block 使用前一个 block 生成的混合系数，从而消除上述依赖：

$$X_{l+1}=B_lX_l+C_lF_l(A_{l-1}X_l),\qquad (A_l,B_l,C_l)=H(X_l). \tag{6}$$

此时输入混合使用 $A_{l-1}$ 而不是 $A_l$，因此不再依赖由 $X_l$ 计算出的系数。于是每个 $X_l$ 的 tile 都可以立即同时用于输入混合和系数预测，无需等待完整归约结束。实验上，这种错位只带来可忽略的性能下降。

预训练期间，我们保留现有的多内核实现，因为这一修改只改变每个 block 实际采用哪一组混合系数。部署时，我们将残差更新、输入混合和系数预测融合为单一内核 Mega-mHC。该内核实现普通 mHC 时需要 $(3n+2)d$ 的激活读写，而实现 Single-Pass mHC 时只需 $(2n+2)d$。Mega-mHC 沿隐藏维以 tile 方式处理 $X_l$；每个 tile 同时用于计算混合输入，并累计预测下一个 block 的 $(A_l,B_l,C_l)$ 所需量。该内核还整合输入 pre-norm 和 FP8 转换。这样，残差只需读取一次、写入一次，实现理想映射所需的 $(n+1)d$ 次读取与 $(n+1)d$ 次写入，并将原实现的激活内存流量减半。

2.4.2. Engram

我们为 DeepSeek-V4.1-Flash 加入 Engram（Cheng et al., 2026c），这是此前工作中提出的条件记忆模块，用于将记忆能力与计算解耦。我们沿用原始 Engram 的设计——tokenizer 压缩、多头哈希、上下文感知门控和多分支集成——并作两项修改。第一，我们移除短因果卷积，因为其性能收益不足以抵消在推理栈中增加的复杂度。第二，我们采用基于动量的更新并结合 Sinkhorn balancing 来优化 Engram 嵌入，详见第 2.5 节。

我们为 Engram 分配 196B 参数，平均分布于两个模块。每个模块使用 $N$-gram 阶数 $\{2,3,4\}$，包含 8 个哈希 head，并且每个阶数的总嵌入维度为 2048。每个 head 索引约 16M 条目的表，且不同表大小选择互不相同的素数。嵌入表以及 key/value 投影均使用 FP8 精度。两个模块分别放置在第 1 层和第 14 层（从 0 开始计数），以平衡训练流水线各阶段的内存使用。

推理时，确定性寻址允许通过后台 RDMA 传输从主机内存预取嵌入；第一个模块的预取可与第一个 Transformer block 的计算重叠。Engram 训练与推理的更多实现细节见第 3.1.3 节。

2.4.3. DSpark

我们为 DeepSeek-V4.1-Flash 配置 DSpark（Cheng et al., 2026a），这是一种将半自回归草稿生成与置信度调度验证结合的推测解码模块。

DSpark 的 drafter 由三个 Transformer block 组成，滑动注意力窗口为 128 token。一次前向传播即可并行计算 5 个草稿位置的基础 logits，同时使用轻量 Markov head 建模草稿 token 之间的依赖。置信度 head 预测各位置的条件接受概率，并据此估计前缀存活概率。调度器将这些估计值与预先测得的引擎吞吐曲线结合，动态选择每个请求的验证长度，目标是在当前系统负载下最大化预期的全系统 token 吞吐量。

与 DeepSeek-V3（DeepSeek-AI, 2024）中在整个预训练期间与主干联合训练的 MTP 模块不同，DSpark 在预训练完成后通过一个专门阶段引入。在这一阶段，我们冻结主干，只训练 DSpark。后训练阶段继续让 DSpark 与主干同时训练，但 DSpark 目标产生的梯度不会反向传播到主干。这样可以使 DSpark 始终与不断演化的策略保持一致，从而同时加速在线服务以及 RL、OPD 中的 rollout 生成。

2.4.4. FP4 主 KV 缓存

长上下文智能体工作负载要求每个请求维持大型 KV 缓存，从而提高服务成本。DeepSeek-V4 已经对 FP4 索引器 query 与 key 使用量化感知训练（QAT）（Jacob et al., 2018），以加速索引计算并减小索引器缓存。我们采用 OCP 标准的 MXFP4 格式（Rouhani et al., 2023），以兼容尽可能多的硬件平台；虽然实验中其他格式的精度更高，但兼容性更差。

现在，我们将 QAT 进一步扩展到主 KV 缓存。这里 FP4 的主要作用是减少存储，而不是加速矩阵乘法。注意力计算前先对缓存值反量化，使我们可以采用更精确的格式，而不要求硬件原生支持该格式的矩阵乘法，从而维持跨硬件平台兼容性。

在评估的约 4 bit 格式中，我们选择 E2M1 + 每 16 个通道一个 E4M3 scale 的方案，整体沿用 NVFP4（Alvarez et al., 2025），但省略其第二级全局 scale，以平衡精度和简洁性。省略全局 scale 后，主 KV 缓存仍有充足的动态范围：该格式可表示的最大幅值为

$$448\times 6 = 2688,$$

远高于缓存本身的幅值上界。在 DeepSeek-V4.1-Flash 中，训练得到的 RMSNorm 权重最大幅值约为 1。经过 RMS 归一化后，512 通道 KV latent 的 $L_2$ 范数最多约为

$$\sqrt{512}.$$

RoPE 保持该范数不变，因此旋转后所有通道的最大绝对值也被约束在

$$\sqrt{512}\approx 22.6.$$

训练过程中观测到的实际最大幅值约为 10。因此，移除全局 scale 不会带来可测量的精度下降，同时能简化缓存布局。

为在 DeepSeek-V4.1-Flash 中启用 FP4 主 KV 缓存存储，我们在后训练期间引入 QAT。非 RoPE 与 RoPE 分量使用同一种量化格式，并在 RoPE 之后对缓存进行量化。实验显示，在 RoPE 之前量化只会带来极小精度提升，却会增加解码开销。SWA KV 缓存对量化更敏感，因此仍保留 FP8。与 DeepSeek-V4 的 FP8 主 KV 缓存相比，该格式无论驻留 HBM 还是卸载到 SSD，存储占用都近似减半。

2.5. 优化

在 DeepSeek-V4 的优化配置基础上，我们做出若干新修改，使其更好地匹配当前架构设计。

第一，我们使用 head-wise Muon：应用 Muon 更新前，先按 head 拆分 Query 权重。

其动机如下。若将 Muon 视为预条件梯度下降，标准 Muon 为所有 head 共用一个预条件器，而 head-wise Muon 为不同 head 提供不同预条件器。该设计能更好地处理注意力 head 之间的异质性（Zhang et al., 2024; Zhang, 2026, Section 3）。实验显示，head-wise Muon 优于标准 Muon；GLM 5（Zeng et al., 2026）和 Kimi-K3（Team et al., 2026a）也验证了这一经验优势。

第二，若对新引入的 Engram 参数使用 Adam，会显著增加优化器状态的内存占用。为减少训练内存，我们改用动量更新 + Sinkhorn balancing 来优化 Engram 嵌入表、token embedding 和 prediction head。Sinkhorn balancing 此前已在 SinkGD（Scetbon et al., 2025）中用于线性层权重矩阵；这里我们将其扩展到这些大型参数矩阵。与 Muon 类似，该方法只需要一个动量缓冲区，同时实验性能优于 Adam。

基础配置。 对归一化层权重以及偏置、缩放因子等其他非矩阵参数，我们仍使用 AdamW（Loshchilov and Hutter, 2019）。语言模型主干中的线性变换权重矩阵、Engram 投影层和视觉-语言投影器使用 Muon（Jordan et al., 2024）。Query 与 Key 权重使用 head-wise Muon。Muon 使用解耦式 weight decay 和 Nesterov momentum（Nesterov, 1983; Liu et al., 2025）；归一化层权重同样施加 weight decay，而偏置和缩放因子不施加。Sinkhorn-balanced update 也使用 Nesterov momentum，但不使用 weight decay。

预训练期间，在学习率衰减阶段之前，视觉编码器保持冻结，其最终归一化层和视觉-语言投影器仍可训练。开始学习率衰减时，我们解冻视觉编码器，并以较小学习率与 LLM 联合优化。

用于 Engram / Embedding / Prediction Head 的 Sinkhorn-Balanced Update。 完整过程见算法 1。总体工作流与 Muon 相同，只是以 Sinkhorn balancing 替代 Newton-Schulz 正交化。记较大的矩阵维度为 $m$，对于嵌入表和 prediction head，该维度对应词表大小；隐藏维度记为 $n$。

给定 Nesterov 动量更新 $\hat G_t$，Sinkhorn balancing 寻找对角缩放矩阵 $D_r$ 与 $D_c$，使得

$$\Delta_t=\sqrt{n}\,U^{(K)}=\sqrt{n}\,D_r\hat G_tD_c,$$

并近似满足

$$\frac{1}{n}\sum_{j=1}^{n}(\Delta_t)_{ij}^{2}\approx 1, \qquad \frac{1}{m}\sum_{i=1}^{m}(\Delta_t)_{ij}^{2}\approx 1. \tag{7}$$

因此，该过程近似使更新矩阵的逐行 RMS 与逐列 RMS 都被均衡。这里，一行对应一个 token 索引或 n-gram 标识，一列对应一个隐藏特征。Sinkhorn balancing 利用这一 token-特征结构，沿行和列两个方向进行归一化。为保证数值稳定，对满足 $\rho_i\le \tau\bar\rho$ 的行进行 mask。因子 $\sqrt n$ 将单位行 $\ell_2$ 范数转换为单位逐行 RMS。此外，为匹配 Adam 的更新幅值，我们将有效学习率调整为

$$\tilde\eta_t=\gamma\eta_t.$$

取 $\gamma=0.18$，与 Moonlight（Liu et al., 2025）中使用的 0.2 因子接近。

从更广泛角度看，Sinkhorn balancing 与利用矩阵或张量轴结构的优化器密切相关（Shazeer and Stern, 2018; Zhang et al., 2025a; Wen et al., 2025; Glentis et al., 2025; Deng et al., 2026; Yuan et al., 2026; Xu et al., 2026a）。

算法 1｜带 Sinkhorn balancing 的动量更新

原始算法 1

输入： 权重 $W_t\in\mathbb{R}^{m\times n}$、梯度 $G_t$、上一时刻动量 $M_{t-1}$、动量系数 $\beta$、基础学习率 $\eta_t$、学习率修正系数 $\gamma$、数值常数 $\varepsilon$ 与 $\tau$，以及奇数次归一化步骤 $K$。

1. $M_t\leftarrow \beta M_{t-1}+(1-\beta)G_t$
2. $\hat G_t\leftarrow \beta M_t+(1-\beta)G_t$ 　// Nesterov momentum
3. $\rho_i\leftarrow \|\hat G_{t,i,:}\|_2$，且 $\bar\rho\leftarrow \frac{1}{m}\sum_{i=1}^{m}\rho_i$
4. $U^{(0)}\leftarrow \hat G_t$
5. 若 $\rho_i\le\tau\bar\rho$，则令 $U^{(0)}_{i,:}\leftarrow 0$ 　// mask 近零行
6. 对 $k=1,\ldots,K$：
   • 若 $k$ 为奇数，则对所有 $i=1,\ldots,m$：

$$U^{(k)}_{i,:}\leftarrow \frac{U^{(k-1)}_{i,:}} {\|U^{(k-1)}_{i,:}\|_2+\varepsilon}.$$

• 否则，对所有 $j=1,\ldots,n$：

$$U^{(k)}_{:,j}\leftarrow \frac{U^{(k-1)}_{:,j}} {\|U^{(k-1)}_{:,j}\|_2+\varepsilon}.$$

7. $\Delta_t\leftarrow \sqrt n\,U^{(K)}$ 　// 将单位行 $\ell_2$ 范数转换为单位行 RMS
8. $\tilde\eta_t\leftarrow \gamma\eta_t$ 　// 匹配 Adam 的更新幅值
9. $W_{t+1}\leftarrow W_t-\tilde\eta_t\Delta_t$

例如，Adafactor（Shazeer and Stern, 2018）以不同方式进行逐行和逐列归一化，而 Adam-mini（Zhang et al., 2025a）对嵌入表和 prediction head 使用另一种逐行归一化。不同归一化策略在优化性能与通信开销上可能存在差异，更深入的研究留待未来工作。