HedgeMamba：融合线性注意力与状态空间模型的高效序列建模

1. 项目概述：当注意力机制遇上线性复杂度

在自然语言处理领域，Transformer架构凭借其强大的注意力机制，几乎统治了所有前沿模型。但任何一个在真实业务场景里部署过大模型的人都知道，那个看似优雅的注意力计算，其时间复杂度是序列长度的平方（O(n²)）。这意味着，当你试图处理一篇长文档、一段长对话，或者仅仅是批量推理时，显存和算力的消耗会呈爆炸式增长。这不仅是成本问题，更是模型能否落地的关键瓶颈。

为了解决这个问题，社区一直在探索线性复杂度的序列建模方法。状态空间模型（State Space Models, SSM），特别是像Mamba这样的选择性扫描架构，提供了一种极具潜力的路径。它能以O(n)的复杂度处理序列，但早期的纯SSM模型在语言建模任务上的表现，尤其是在需要复杂上下文理解的任务上，往往难以匹敌成熟的Transformer。另一方面，线性注意力（Linear Attention）试图通过数学变换绕过Softmax的计算瓶颈，但简单的线性化往往会牺牲模型的表现力。

HedgeMamba项目的核心思路，正是将这两条技术路径的优势结合起来。它没有选择“替代”，而是选择了“融合”。具体来说，它利用Hedgehog特征映射技术，将标准的Softmax注意力机制优雅地线性化，然后将这个线性化的注意力计算，巧妙地嵌入到Mamba的SSM框架中。最终的目标是：得到一个在数学上等价于注意力机制，但在计算上享有线性复杂度的高效模块。这个模块可以直接用于构建新模型，但在这个项目中，它被用于一个更务实的场景——知识蒸馏。我们将一个庞大的、性能优异的Transformer教师模型（如Pythia）的知识，蒸馏到一个由我们新设计的、高效的HedgeMamba模块构成的学生模型中。

简单来说，我们不是在从头训练一个全新的线性模型，而是让一个高效的“学生”去模仿一个强大的“老师”的行为。这样，学生模型既能继承老师的大部分能力，又能在推理时跑得更快、吃得更少。这对于将大模型能力下沉到资源受限的边缘设备或追求极致吞吐的线上服务，有着直接且巨大的价值。

2. 核心架构设计：Hedgehog与Mamba的化学反应

要理解HedgeMamba，我们需要拆解它的两个核心技术支柱：Hedgehog特征映射和Mamba状态空间模型，并看它们是如何协同工作的。

2.1 Hedgehog特征映射：注意力线性化的钥匙

标准Transformer中的缩放点积注意力（Scaled Dot-Product Attention）公式为：Attention(Q, K, V) = softmax(QK^T / √d) V。这里的softmax操作是导致二次复杂度的元凶，因为它需要计算序列中所有位置对之间的相似度矩阵QK^T。

线性注意力的核心思想是找到一个特征映射函数φ(·)，使得softmax(QK^T)可以近似表示为φ(Q) φ(K)^T。这样，注意力计算就可以重排为(φ(Q) (φ(K)^T V))，利用矩阵乘法的结合律，先计算φ(K)^T V（一个[d, d]的矩阵），再与φ(Q)相乘，从而将复杂度从O(n²d)降为O(nd²)。当序列长度n远大于特征维度d时，节省的计算量是巨大的。

然而，找到一个既能保持注意力表达力，又计算高效的φ(·)并非易事。Hedgehog特征映射的提出，正是为了解决这个问题。它的设计非常巧妙：

1. 映射与拼接：对于查询Q或键K，先通过一个可学习的线性层φ进行投影。然后，将投影后的结果与其负值进行拼接：[φ(x), -φ(x)]。这个操作相当于同时考虑了特征的正向和负向激活。
2. Softmax归一化：接着，沿着特征维度（dim=-1）应用Softmax函数。注意，这里不是在序列长度维度上做Softmax，而是在这个扩展后的特征维度上做。这一步至关重要，它使得φ(Q) φ(K)^T的内积结果，在数学上能够模拟标准注意力中Softmax后的概率分布特性。项目附录C的代码明确使用了dim=-1的Softmax，并指出这是为了数值稳定性，替代了原论文中的指数操作。

注意：这里的Softmax沿着特征维，与标准注意力沿着序列维的Softmax有本质区别。它不计算位置间的依赖，而是对每个位置的特征表示进行归一化，为后续的线性内积模拟注意力权重创造条件。

通过这套操作，Hedgehog得到了一个高效的、可导的φ(·)，使得线性注意力在多个基准测试上能够接近甚至匹配标准注意力的性能。

2.2 Mamba SSM：选择性扫描与硬件优化

Mamba是SSM家族中的佼佼者，它解决了传统SSM（如S4）在处理离散、信息密集的序列数据（如文本）时的两个关键短板：

1. 输入依赖性：传统SSM的参数（如状态转移矩阵）是静态的、与输入无关的。Mamba引入了输入依赖的机制，让SSM的关键参数（如离散化步长Δ、B和C矩阵）能够根据当前输入token动态变化。这使得模型能够像注意力一样，选择性地记住或忽略信息。
2. 硬件感知算法：Mamba设计了并行的选择性扫描算法，充分利用GPU的并行计算能力，避免了传统SSM递归计算的串行瓶颈，实现了高效的训练和推理。

Mamba块的核心是一个“SSM Mixer”，它接收输入序列，通过一个轻量的卷积层进行局部特征混合，然后通过上述输入依赖的机制参数化一个离散状态空间系统，最后通过扫描（Scan）操作计算出输出。此外，它还有一个独立的“门控”（Gate）分支，用于调制SSM的输出。

2.3 HedgeMamba混合模块：如何将两者组装

HedgeMamba的设计哲学是“站在巨人的肩膀上”。它没有完全重新发明轮子，而是对Mamba的SSM Mixer进行了针对性的改造，使其能够执行线性化的注意力计算。

观察Mamba的SSM方程（简化）：它的计算涉及输入V、动态参数B和C。如果我们把B和C分别类比为线性注意力中的φ(K)和φ(Q)，把SSM的扫描过程类比为线性注意力中φ(K)^T V的累积计算，那么Mamba的SSM天然就提供了一个计算线性注意力的框架。

因此，HedgeMamba模块（对应项目图4中间部分和附录C的Lst. 2）做了如下关键修改：

1. 引入值投影：在Mamba原有结构上，增加了一个v_proj线性层，用于从输入中投影出V（值）。
2. 应用Hedgehog映射：将Mamba中原本用于生成B和C的线性投影输出，分别送入两个独立的Hedgehog投影模块（hhog_k和hhog_q）。这样，我们就得到了线性注意力所需的φ(K)和φ(Q)。
3. 状态扩展与归一化：为了精确模拟注意力中Softmax的归一化效果，HedgeMamba对SSM的内部状态进行了“扩容”。它将参数A、dt和输入V都复制了一份，拼接在一起。在SSM扫描完成后，输出也被分为两部分：一部分是未归一化的“分子”（类似于φ(Q)(φ(K)^T V)），另一部分则是用于归一化的“分母”（类似于φ(Q)(φ(K)^T 1)）。最终的输出就是分子除以分母，完美实现了线性注意力中的归一化步骤。
4. 保留门控与卷积：Mamba原有的短卷积和门控分支被保留，它们提供了额外的局部建模和非线性能力。

通过这样的设计，HedgeMamba模块在前向传播的数学效果上近似等价于一个标准的注意力层，但其计算复杂度是线性的。它继承了Mamba的硬件高效性，同时通过Hedgehog映射注入了注意力机制强大的特征交互能力。

3. 两阶段蒸馏策略：从初始化到微调

拥有了高效的HedgeMamba学生模块，下一步就是如何让一个由该模块构建的学生网络，学会模仿一个庞大的Transformer教师网络。直接从头训练一个线性复杂度的模型达到SOTA性能非常困难，因此知识蒸馏成为了关键桥梁。本项目采用了一个精心设计的两阶段蒸馏策略，这是取得好结果的核心。

3.1 第一阶段：注意力行为模仿

第一阶段的目标，是让学生模型的HedgeMamba模块“初始化”出近似于教师模型注意力模块的行为模式。在这一阶段，我们冻结了HedgeMamba中从Mamba继承来的、用于实现“选择性”的动态参数（即输入依赖的Δ、B、C生成机制中的特定部分，以及门控和卷积分支）。

如何冻结？通过特殊的参数初始化，使这些模块在前向传播中变为“恒等映射”。

• 状态矩阵Λ：通过将衰减率λ初始化为0，并使生成Δ的权重Wu为0、偏置bu为一个特定常数，使得Λ矩阵变为单位矩阵，SSM的动态性被暂时取消。
• 卷积层：将卷积核权重初始化为[0, 0, ..., 1]（最后一个元素为1），偏置初始化为0，使其成为一个简单的延迟操作，近似于恒等映射。
• 门控层：将权重初始化为0，偏置初始化为SiLU⁻¹(1)，使得门控输出恒为1，不起调制作用。

这样做的目的是，在第一阶段，HedgeMamba模块退化为一个纯粹的、静态的线性注意力层。训练的目标是最小化学生和教师模型在相同输入下，对应层输出的均方误差（MSE）或KL散度。这个阶段使用的学习率相对较高（如0.01），因为它需要学习一个相对宏观的映射关系。

实操心得：第一阶段的本质是“功能对齐”。它不要求学生模型学会选择性地关注信息（那是Mamba动态部分的工作），而是先学会如何像注意力一样，对信息进行加权聚合。这为第二阶段引入动态选择性打下了坚实的基础。如果跳过此阶段，直接进行第二阶段，模型很容易陷入局部最优，无法有效学习到注意力机制的精髓。

3.2 第二阶段：全参数联合微调

在第二阶段，我们解除所有参数的冻结，让HedgeMamba模块中Mamba的动态部分（选择性扫描机制）、卷积和门控全部参与训练。此时，学生模型不再仅仅是模仿静态的注意力输出，而是要学习如何动态地、选择性地运用从第一阶段学到的“注意力式”的信息融合能力。

这一阶段的学习率要设置得非常低（例如1e-5量级），因为它是在一个已经相对良好的初始化上进行精细调整。训练目标通常是标准的语言建模损失（如交叉熵），数据则来自教师模型在大量文本上预测的分布，或者直接使用原始文本数据。

项目中的敏感性分析（图3和表5）清晰地证明了这两个阶段缺一不可。当所有训练token都用于第二阶段（0/100 split）时，验证集困惑度（PPL）最终停滞在一个较高的水平。而即使只将10%的token预算用于第一阶段（10/90 split），也能获得最佳的最终性能。这说明，第一阶段提供的良好初始化，是后续动态机制能够有效学习的前提。

3.3 训练配置与超参选择

根据附录A.2的实现细节，我们可以还原出关键的训练配置：

• 框架：PyTorch，使用分布式数据并行（DDP）和混合精度训练（bfloat16）以加速并节省显存。
• 优化器：AdamW，这是当前训练LLM的事实标准。β1=0.9， β2=0.95，权重衰减（weight decay）为0.1，梯度裁剪（gradient clipping）为1.0。这些是经过大量实践检验的稳定配置。
• 学习率调度：采用线性预热（warm-up）接余弦衰减（cosine decay）到峰值学习率的10%。这是Transformer家族模型常用的调度策略，能稳定训练初期并让模型在后期更好地收敛。
• 硬件与规模：实验在8张NVIDIA A100 GPU的节点上进行。蒸馏的总token预算约为100亿（10B），对应20万训练步，每步batch size对应约4.9万token。对于1B参数的学生模型，完成一次10B token的蒸馏需要大约12天9小时。作者指出，由于Mamba选择性扫描实现中对模型维度（d_model）有256的硬限制，当他们的模型维度达到2048时，需要序列化处理，导致了额外的计算开销。

4. 实验结果深度剖析与方案选型启示

项目的实验结果表格（表5、表6）包含了丰富的信息，我们需要从中解读出HedgeMamba的有效性、优势边界以及对实际应用的指导意义。

4.1 核心性能对比：HedgeMamba vs. 基线

在1B参数规模下（表5），对比以下几组有代表性：

1. 教师模型 (Pythia)：作为性能上限参考。
2. Hedgehog基线：仅使用线性注意力（Hedgehog特征映射）替换Transformer中的注意力。这是最直接的线性化方案。
3. HedgeMamba (Ours)：我们提出的混合方案。

从Perplexity（PPL，越低越好）看，HedgeMamba（14.11）显著优于纯Hedgehog基线（14.89），并且更接近教师模型（13.86）。这说明引入Mamba的动态SSM机制带来了增益。

从下游任务看，在Arc-E（53.66 vs 52.74）、BoolQ（55.20 vs 54.80）、Winogrande（52.17 vs 50.91）、HellaSwag（41.87 vs 40.79）等多个常识推理和语言理解任务上，HedgeMamba都一致地超越了Hedgehog基线。特别是在需要较长上下文推理的Lambada（预测句子最后一个词）任务上，HedgeMamba（32.31）比基线（30.60）有近2个点的提升，这很可能得益于Mamba的选择性状态机制对长程依赖的更好建模。

4.2 模型规模缩放分析

表6展示了从160M、410M到1B参数规模下的蒸馏效果。一个清晰的趋势是：随着模型规模增大，所有方法的性能都在提升，但HedgeMamba相对于Hedgehog基线的优势始终存在。

• 在160M小模型上，HedgeMamba的PPL（26.84）远低于基线（35.95）和教师（39.38），甚至在多项任务上超过教师。这可能是蒸馏过程起到了“提纯”和正则化的作用。
• 在1B模型上，HedgeMamba在绝大多数任务上超越了基线，并且其PPL与教师的差距（14.11 vs 13.86）远小于基线与教师的差距（14.89 vs 13.86）。

这表明，HedgeMamba的架构优势在不同规模下都具有鲁棒性，并且模型越大，其吸收教师知识、逼近教师性能的潜力越大。

4.3 两阶段分词分配敏感性分析

表5中间部分研究了第一阶段和第二阶段训练token分配比例的影响。结论非常明确：

• 纯第一阶段（100/0）效果最差：模型只学会了静态模仿，无法进行动态选择，在需要推理的任务上（如Lambada, 12.12）表现崩盘。
• 纯第二阶段（0/100）次优：如前所述，缺乏良好初始化，性能上限受限。
• 10/90或25/75分配通常最佳：仅用少量资源进行第一阶段初始化，就能为第二阶段的动态学习奠定坚实基础，在PPL和多项下游任务上取得最佳平衡。这为资源分配提供了直接指导：不需要对第一阶段投入过多预算，但其不可或缺。

4.4 方案选型启示：何时考虑HedgeMamba？

基于以上分析，在技术选型时可以考虑以下几点：

1. 目标场景：如果你的核心痛点是在长序列推理（如长文档处理、对话历史建模）且对延迟和吞吐要求苛刻，那么线性复杂度模型是必然方向。HedgeMamba提供了一个性能接近Transformer的强基线。
2. 资源与阶段：如果你有一个训练好的、性能优秀的Transformer教师模型（例如业务场景中的精调大模型），并且希望得到一个高效的部署版本，那么这种两阶段蒸馏方案非常合适。第一阶段可以复用教师模型的中间层特征进行快速对齐。
3. 对比其他方案：需要与其他高效架构对比，如纯Mamba、其他线性注意力变体（如FlashAttention-2虽然优化了计算，但复杂度仍是O(n²)，只是常数更优）、以及模型压缩技术（如量化、剪枝）。HedgeMamba的优势在于其性能潜力，劣势在于架构相对复杂，训练和实现需要更精细的调优。
4. 从头训练：本项目聚焦蒸馏。若想从头训练一个HedgeMamba模型，需要谨慎。两阶段策略可以借鉴：先以较高学习率训练“静态”部分，再以低学习率微调动态部分。但其数据效率和最终性能上限仍需更多探索。

5. 实操复现指南与核心代码解读

对于希望复现或在其基础上进行开发的工程师而言，理解代码实现细节至关重要。项目附录C提供的伪代码清晰地勾勒出了核心模块。

5.1 整体层结构替换

如Lst. 1所示，整体框架基于Pythia（一种GPT-NeoX架构的变体）。该架构的特点是前馈网络（MLP）和注意力层是并行（Parallel）的，而非Transformer中常见的串行（Sequential）结构。我们的改造非常直接：将原有的attention模块替换为自定义的HedgeMambaMixer。前向传播时，输入同时经过LayerNorm后送入HedgeMamba模块和MLP模块，两者的输出再与残差连接相加。

# 伪代码示意，展示并行结构 def forward(self, x): # 注意力流 attn_path = self.attention(self.norm1(x)) # MLP流 mlp_path = self.mlp(self.norm2(x)) # 残差连接 return x + attn_path + mlp_path

5.2 HedgeMambaMixer 核心实现

Lst. 2是核心所在。我们来拆解关键步骤：

1. 投影层：x_proj层将输入映射，用于生成SSM的动态参数dt（离散化步长）、B和C。注意，这里的B和C的维度与注意力头的设计相关，被重塑为[batch, seq_len, num_heads, head_dim]。
2. Hedgehog映射：将B和C分别送入hhog_k和hhog_q（两个独立的HedgehogProjection实例），得到φ(K)和φ(Q)。
3. 位置编码：对φ(Q)和φ(K)应用旋转位置编码（RoPE），这是保持模型感知序列位置能力的关键，与原始Transformer和Pythia保持一致。
4. 值投影与状态扩容：通过v_proj得到V。然后，为了进行归一化，需要复制一份“分母”计算所需的量。这里巧妙地将SSM的状态矩阵A、步长dt和输入V都沿着特定维度复制并拼接（torch.cat）。复制的V部分被替换为全1的向量，因为在注意力归一化中，分母相当于φ(Q)(φ(K)^T * 1)。
5. 选择性扫描：调用Mamba原生的selective_scan_fn函数，传入扩容后的参数进行计算。这个函数是Mamba高效实现的核心，以并行的方式完成了理论上串行的状态扫描。
6. 归一化与门控：扫描输出被拆分为分子和分母两部分，执行除法完成归一化。最后，结果与门控分支（gate）的输出进行元素级相乘，引入非线性。
7. 输出投影：通过最后的out_proj线性层将输出映射回模型维度。

5.3 HedgehogProjection 模块

Lst. 3的实现简洁明了。重点在于前向传播：

1. 通过一个线性层self.phi进行投影。
2. 将结果与其负值拼接：torch.cat([x, -x], dim=-1)。
3. 在拼接后的特征维度（最后一个维度，dim=-1）上应用Softmax。

注意事项：这里的Softmax是沿着特征维（dim=-1），这与标准注意力在序列维（dim=1或dim=2）做Softmax有本质区别。它的作用是让φ(x)的每个特征维度成为一个概率分布，从而使得后续的φ(Q)φ(K)^T内积能够模拟出类似注意力权重的行为。使用Softmax而非原论文的指数函数，是为了增强数值稳定性，防止在深度网络中产生梯度爆炸或溢出。

5.4 复现环境搭建与调试建议

1. 环境依赖：需要PyTorch（>=1.12以支持torch.compile可能带来的优化）、CUDA工具链。Mamba的选择性扫描内核需要特定的CUDA编译，可参考Mamba官方仓库的安装指南。
2. 代码集成：将HedgehogProjection、HedgeMambaMixer和修改后的层结构集成到现有的Transformer架构（如HuggingFace的Pythia或GPT-NeoX）中。注意对齐维度，特别是多头注意力的头数（num_heads）和头维度（head_dim）。
3. 蒸馏流程：
   • 阶段一：加载教师模型（如EleutherAI/pythia-1b）。实例化学生模型，并按照3.1节所述初始化Mamba动态参数、卷积和门控权重。使用均方误差（MSE）损失，在教师模型的中层特征或输出logits上进行蒸馏。
   • 阶段二：加载阶段一训练好的学生模型检查点。使用语言建模损失（交叉熵），在大量文本数据上继续训练。学习率应调至很低（如1e-5）。
4. 调试技巧：
   • 初始化检查：在阶段一开始前，可以手动验证冻结模块的输出是否接近恒等映射。例如，给门控分支输入随机张量，看输出是否全为1（在SiLU激活下）。
   • 梯度流：使用torch.autograd.detect_anomaly()或在训练初期检查各模块的梯度范数，确保第一阶段动态参数梯度为零，第二阶段梯度正常。
   • 内存监控：线性复杂度模型在长序列下内存增长应近似线性。如果出现异常增长，检查是否有无意中产生的[seq_len, seq_len]中间矩阵。

6. 常见问题、潜在挑战与未来方向

在实际操作中，你可能会遇到以下问题或产生相关思考：

6.1 性能与效率的权衡

• Q：HedgeMamba真的比标准Attention快吗？

  • A：在理论复杂度上，对于长序列（n >> d），O(n)远优于O(n²)。但在实际硬件运行中，常数项和内存访问模式至关重要。Mamba的选择性扫描内核经过高度优化，在长序列上已显示出显著的速度和内存优势。然而，HedgeMamba增加了额外的线性投影和Hedgehog映射，可能引入一些开销。最终收益需要在目标序列长度和硬件上进行实测。对于短序列（如<512），标准Attention的优化实现（如FlashAttention）可能仍然更快。

• Q：两阶段蒸馏的训练总成本是否很高？

  • A：是的，这是一个需要权衡的点。虽然学生模型推理高效，但蒸馏过程本身需要大量计算。10B token的蒸馏在8xA100上耗时约12天。这需要对比：1）直接训练一个同等大小的标准Transformer；2）训练一个纯Mamba模型。如果最终部署的推理请求量极大，那么前期的一次性高成本蒸馏投资是值得的。反之，如果推理需求不大，则需谨慎评估。

6.2 实现与调优难点

• Q：Mamba选择性扫描的维度限制（256）对模型设计影响大吗？

  • A：影响显著。原版Mamba的CUDA内核为状态维度（在HedgeMamba中对应d_state或head_dim的某种映射）设置了256的上限。当项目中使用2048的模型维度时，被迫进行序列化处理，导致了8倍以上的训练时间增长。这是目前将Mamba用于超大宽度模型的一个主要工程瓶颈。解决方案包括：1）等待社区优化内核；2）调整模型结构，例如使用更多的头数（num_heads）但更小的头维度（head_dim），使相关维度不超过256。

• Q：Hedgehog映射中的Softmax维度容易搞错，如何验证其正确性？

  • A：一个简单的验证方法是构造一个微型测试。创建随机张量Q和K，形状为[batch, seq_len, dim]。分别用标准Softmax注意力（在seq_len维）和你的Hedgehog投影（在dim维，拼接后）计算QK^T的近似值。虽然结果不会完全一致，但应观察到相似的结构（如较高的值在对角线附近）。更重要的验证是在完整的蒸馏任务中，如果第一阶段损失能顺利下降，则说明映射基本有效。

6.3 理论理解与扩展

• Q：HedgeMamba如何保证其线性计算与注意力机制的等价性？

  • A：这是一种近似等价，而非严格数学等价。Hedgehog特征映射的设计目标就是找到一个φ(·)，使得<φ(q), φ(k)>近似于exp(q·k / √d)。Softmax沿着特征维的应用，是为了让φ(x)具有某种“能量”归一化特性，从而让内积模拟出概率分布的点积。Mamba的SSM框架则提供了一个高效计算φ(Q)(φ(K)^T V)的机制。两者的结合，是在表达能力和计算效率之间寻求一个工程上最优的平衡点。

• Q：这个方法可以扩展到视觉或多模态任务吗？

  • A：非常有潜力。视觉Transformer同样受注意力二次复杂度的困扰。线性注意力在视觉任务上已有探索。HedgeMamba的思路可以迁移：将图像序列视为patch序列，用Hedgehog进行线性化，并用SSM进行高效混合。关键在于如何适应图像的二维结构（可能需要对SSM进行修改，如使用二维扫描顺序）。对于多模态任务，不同的模态可以共享或拥有独立的HedgeMamba模块，来处理跨模态的线性注意力交互。

6.4 未来优化方向

从我个人的实验经验来看，HedgeMamba是一个将前沿研究进行工程化落地的优秀范例。要将其用于生产，还可以从以下几个方向深入：

1. 内核融合与极致优化：当前实现可能包含多个独立的核函数调用（线性层、Hedgehog Softmax、扫描）。可以尝试编写融合CUDA内核，将Hedgehog投影与SSM参数生成甚至扫描过程融合，减少内存读写和启动开销。
2. 量化与部署友好性：研究HedgeMamba模块的量化特性。由于SSM涉及递归和复杂激活，其量化可能比标准Attention更敏感。探索INT8甚至FP4量化方案，对于端侧部署至关重要。
3. 探索更优的特征映射：Hedgehog是一种成功的映射，但未必是终点。可以探索其他能产生正定核函数的特征映射，或者可学习的映射网络，看是否能以更低的计算成本达到更好的近似效果。
4. 与MoE结合：将HedgeMamba作为专家混合（Mixture of Experts）模型中的专家，每个专家负责处理序列的不同子空间或模式，可以在保持线性复杂度的同时，大幅增加模型容量。

这个项目清晰地展示了一条路径：通过蒸馏将Transformer的知识迁移到高效的线性架构中。它不仅仅是提出了一个混合模块，更提供了一套完整的、可复现的训练方法论。对于受限于推理成本的团队来说，这类工作提供了将大模型能力“瘦身”并加速的切实可行的技术选项。在实际应用中，你需要根据自身的数据特点、序列长度分布和硬件条件，对架构细节和训练配方进行微调，才能榨取出最大的性能效率比。