苏剑林解析K3大模型架构优化细节

2026-09-09 17:20:53

文章摘要

苏剑林解析迭代开源大模型K3的架构优化细节，K3是K2的迭代版本，整合团队过往研究成果。本文重点介绍Stable LatentMoE混合专家模块的稳定性优化方案，以及保留MLA注意力架构、改用NoPE位置编码的设计取舍，指出K3的优化均为平衡效果、效率与稳定性的渐进调整。

近期推出的开源大模型K3，是K2模型的迭代版本，并非从零搭建的全新架构，而是基于团队过往的系列研究成果逐步演进而来，整合了在模型效果、推理效率与训练稳定性上的最新探索与优化。相较于从零开始的全新设计，K3的研发路径更像是对已有技术的集大成式整合，每一处调整都有明确的实践与理论支撑。接下来我们将详细拆解K3在架构层面的核心设计思路。

K3的整体架构可以概括为KDA、MLA、Stable LatentMoE与AttnRes的组合，训练阶段使用的优化器为Moonlight版Muon，同时针对注意力模块的权重做了Per-Head形式的调整，进一步提升了训练的正确性与稳定性。其中AttnRes与KDA的相关设计已经在过往的技术报告中做过详细介绍，本文将重点围绕Stable LatentMoE与MLA两个模块展开分析，讲解团队在优化这两个模块时的思考与实践。

混合专家系统的稳定性优化

K3所采用的混合专家系统被命名为Stable LatentMoE，该设计源自LatentMoE的相关研究，其核心优势在于可以在相近的训练与推理成本下实现更优的模型效果，但原版LatentMoE在实际应用中存在训练稳定性不足的问题，团队因此针对该模块做了一系列针对性优化。

当前主流的MoE架构中，单个专家模块通常采用SwiGLU作为激活函数，其主要非线性来源为SiLU函数。但在实际训练中，SwiGLU容易出现数值异常的问题：当W₁的某一行向量与输入向量同向时，会导致点积结果过大，极端情况下甚至会出现O(‖x‖⁴)级别的异常值，影响训练的稳定性。

为了解决这个问题，团队首先将SiLU替换为SiTU（Sigmoid Tanh Unit），将门控部分的输出限制在(-β, β)的区间内，其中β取值为4。进一步的测试发现，仅通过该调整仍无法完全避免数值膨胀的问题，因此团队又在线性变换环节加入了softcap运算，最终形成了SiTU-GLU的激活结构，其中β₁=4，β₂=25。相较于此前部分模型使用的Hard Clip操作，softcap在相同的数值限制下可以取得更好的效果，因此成为了团队的选择。

LatentMoE与传统MoE的核心区别在于，其会先对输入进行降维，再执行2n选2k的MoE路由，最后通过升维恢复原始维度，这样可以在保持训推成本相近的前提下获得更优的效果。但原版LatentMoE在降维和升维之间没有额外的归一化操作，加上中间的MoE模块，会形成四个矩阵连乘的结构，训练过程极易出现不稳定的情况。

团队最初尝试在降维后的输出与升维前的输入两个位置都添加RMS Norm，但通过消融实验发现，仅在升维前的位置添加RMS Norm就可以达到最佳的稳定效果，因此最终采用了最小改动的原则，仅保留了该位置的归一化操作。后续的对比实验显示，这一调整不仅可以稳定训练过程，还能在部分基准测试中带来效果提升，其原因可能是该归一化操作平衡了路由专家与共享专家的比例，同时也为模块增加了微弱的非线性变换，提升了模型的等效深度。

K3最初的MoE设计为448选8，引入LatentMoE后调整为896选16，虽然稀疏度保持不变，但总专家数量的增加会加剧负载不均衡的问题。与前代模型K2一样，K3采用了Loss-Free的负载均衡方案，但此前使用的SignSGD式更新规则在总专家数量较大时表现不稳定，因此团队引入了QB（Quantile Balancing）算法。该算法的核心优势在于数学逻辑更严谨，且无需额外的超参数。

QB的核心运算为求解全局分位数，但分位数属于非线性运算，直接朴素计算会带来极大的通信开销。团队最初尝试通过局部求解分位数再进行全局平均的方式优化，但在模型规模进一步扩大后发现该方法的精度不足，因此最终采用了分箱近似（直方图估计）的方案：将需要计算分位数的分数压缩到0~1的区间后，通过分箱估计分数的分布，再从分布中读取对应的分位数。测试显示，使用1000个分箱就可以达到足够的精度，相较于10000个分箱并未带来更优的负载均衡效果。基于分布的可加性，该方案可以以极低的通信开销实现跨机器、跨梯度累积的分布信息聚合，从而获得全局的近似分位数。

注意力机制的设计取舍

K3的注意力模块采用了KDA与MLA的混合设计，本文将重点讨论MLA模块的设计思路。或许有读者会疑惑，部分最新模型已经放弃了MLA的设计，为何K3仍选择该方案？实际上团队在选择注意力架构时经过了慎重的考量，MLA依然是当前场景下的最优选择之一。

早在一年前，团队就通过多篇技术博客从实验与理论两个层面探究了MLA的优势，当时得出的结论是：在相同的训练成本与推理成本下，MLA可能是效果最好的全注意力变体。这一结论在今天依然基本成立，但也存在一些细微的变化。

在固定训练成本与KV Cache大小的前提下，MLA依然是近乎最优的注意力架构，但随着推测解码（MTP）技术的出现，推理阶段的速度不再仅由KV Cache决定。MLA在推理阶段表现为较大维度的多查询注意力，会消耗大量的计算资源，因此“MLA+MTP”的组合在推理速度上并不占优。

但注意力架构的选择需要综合多方面的因素，MTP只是其中一个维度。如果更换其他注意力设计，虽然可能在MTP上表现更优，但在其他方面未必能超过MLA。例如，如果将MLA的维度调小，其效果很难超过原版MLA，且会大幅增加KV Cache的占用量，在实际应用中并不现实。而如果将架构放大为其他形式，虽然效果可以追平MLA，但训练成本会显著提升，在缩放定律的视角下并不划算，同时Prefill阶段的计算开销也会增加，而当前主流的Agent/Coding场景中，每一轮的Prefill长度通常并不短。

因此，理想的优于MLA的注意力架构需要同时满足多个条件：效果不弱于MLA、训练与Prefill成本不超过MLA、KV Cache占用量小于MLA、推理阶段的计算量小于MLA。但就目前的技术进展来看，尚未有简单优雅的注意力设计可以同时满足所有这些条件，因此团队只能根据自身的场景做出取舍。在与KDA模块混合的背景下，MLA的部分问题得到了缓解，因此团队最终依然选择了MLA作为核心的注意力架构。

这里可以补充讨论一下相关模型的注意力设计。部分最新模型表面上放弃了MLA，采用了全新的注意力架构，但仔细分析后可以发现，其依然带有MLA的设计思路，并且符合此前提到的几个优化方向。团队过往的研究曾指出，无形增大的head_dims是MLA效果的关键，同时给定KV Cache大小的前提下，效果最好的注意力架构是head_dims等于KV Cache大小、K与V共享的多查询注意力。相关模型正是采用了高维度的共享K/V的MQA，配合特定的位置编码，这其实就是MLA在推理阶段的形式。但这样的设计会导致训练与Prefill阶段的计算量大幅增加，且由于没有线性注意力，每一层都需要保留KV Cache，会带来不小的资源开销，因此这类模型引入了稀疏化与压缩操作来节省计算量与KV Cache占用，虽然这些优化可以有效缓解问题，但也带来了基础设施复杂度提升的代价，其最优性仍有待进一步验证。总体而言，从MLA到这类新架构更像是一次传承与升级，而非完全的抛弃与重造。

K3的MLA模块还有一个细节值得讨论：团队在保持标准MLA结构的前提下，移除了RoPE位置编码，改为使用NoPE。这一改动最早出现在Kimi Linear的设计中，但在K3发布后引发了不少讨论。

实际上，K3依然可以添加RoPE位置编码，且添加后模型效果并未出现明显变化，因此团队本着简洁性的原则选择了移除RoPE。但需要注意的是，对于纯MLA架构的模型，RoPE是不可或缺的，移除后会导致效果明显下降。K3可以使用NoPE的原因在于其是KDA与MLA的混合架构。

根据团队过往的旋转位置编码完备性分析，任意正交矩阵的幂都可以用于构建广义的RoPE，常规的RoPE采用的是简单的旋转矩阵，而其他变体则采用了不同的正交变换，同样取得了不错的效果。在相关推导中，正交情形下的位置编码变换可以等价为给Q、K添加特定的线性变换，而KDA正是更一般化的这类变换，因此KDA与MLA的混合架构本身就自带了类似RoPE的位置编码效果，可以说K3并非完全没有位置编码，而是通过KDA模块隐含提供了广义的位置编码。

至于为何保留了特定的维度拼接操作，团队给出了多个原因：首先是为了更好地适配现有的基础设施，无需重写大量代码；其次，如果调整为更简洁的投影方式，虽然代码更优雅，但会增加计算量，且并未带来效果提升，综合来看反而得不偿失；最后，K3已经引入了多个新的设计变量，团队不想在MLA模块中再引入过多的未知变量，毕竟架构优化需要循序渐进。