近年来，大模型技术飞速发展，但其高昂的计算成本和内存需求也带来了部署难题。Google 近期发布的 Gemma 3n 模型，以其轻量级和高性能引起了广泛关注。其中一个关键创新点在于其核心架构 MatFormer (Matryoshka Transformer)。这篇文章将深入探讨 MatFormer 的原理、优势，以及它如何助力 Gemma 3n 实现高效的模型压缩与推理，并最终服务于各类资源受限的应用场景。MatFormer 的引入，为大模型小型化和部署提供了新的思路，也为未来的模型 scaling 和部署策略带来了变革性的影响。
MatFormer 的核心概念：嵌套Transformer
MatFormer，又名“俄罗斯套娃Transformer”，其核心思想是将多个不同尺寸的Transformer子模型嵌套在一个更大的Transformer模型中。与传统的模型蒸馏或混合专家模型 (MoE) 不同，MatFormer 通过共享权重和巧妙的矩阵切片技术，无需额外的训练或重构架构，即可提取出多个功能完备的子模型。这种嵌套架构避免了传统方法带来的精度损失和训练成本，极大地提升了模型部署的灵活性。 想象一下，一个完整的 Transformer 架构就像一个大型的乐高城堡，而 MatFormer 的目标是在这个城堡内部，构建多个较小的、功能完善的乐高小屋。这些小屋共享城堡的地基和部分砖块，但拥有各自独立的房间和用途。
FFN嵌套：降低参数量和推理延迟
MatFormer 主要针对 Transformer 架构中的前馈网络 (Feed-Forward Network, FFN) 进行了优化。FFN 通常由两个线性层组成，中间夹着一个激活函数（如 GELU 或 SwiGLU）。FFN 层是 Transformer 模型中参数量和推理延迟的主要贡献者。在 Llama 3 和 Gemma 3 这样的模型中，FFN 的中间层维度 (dff) 通常是输入 embedding 维度 (d) 的 4 到 6 倍，导致模型体积庞大。MatFormer 的创新之处在于，它没有为每一层 FFN 固定一个宽度，而是嵌套了多个不同宽度的 FFN。 例如，假设输入维度 d = 4096，最大的 FFN 维度 dff = 16384。我们可以定义一系列粒度：m1 = 4096, m2 = 8192, m3 = 12288, m4 = 16384。那么，宽度为 m_i 的子模型，只需使用完整权重矩阵的左上角子矩阵即可。 这种设计巧妙地实现了参数共享，避免了冗余计算，从而降低了模型大小和推理延迟。
联合训练：保证所有子模型的性能
如何确保所有嵌套的子模型都具有良好的性能？ MatFormer 采用了一种高效的训练方法：随机抽样。 在每个训练步骤中，MatFormer 不是同时计算所有子模型的损失，而是随机选择一个粒度的子模型进行评估。具体来说，MatFormer 使用以下多粒度损失函数：
L = Σ p(Mi) * L(Mi)
其中 Mi 是对应第 i 个粒度的子模型，p(Mi) 是抽样概率。 这种设计有两个关键优势：
- 避免了同时评估所有粒度的开销。
- 由于每个子模型都与更大的模型共享参数，因此，在每个步骤中训练不同的子模型，仍然可以更新共享的权重矩阵。
在 MatFormer 的论文中，作者最初采用均匀抽样，即每个子模型以相同的概率被选中。但后续实验表明，适当提高最大粒度（即完整模型）的损失权重，可以提升大型模型的性能，同时对小型子模型的性能影响不大。
Mix’n’Match：灵活定制模型大小
MatFormer 最令人兴奋的特性之一是 “Mix’n’Match”：在推理时，可以灵活地组合不同层级的 FFN 宽度，从而生成大量的自定义模型。 在训练过程中，MatFormer 优化了多个具有特定粒度的子模型。例如，一个子模型在所有层中使用宽度 m1，另一个子模型在所有层中使用宽度 m3。 然而，这并不意味着我们只能使用统一宽度的模型。在推理时，我们可以逐层改变 FFN 宽度。例如，第一层使用 m2，第二层使用 m3，第三层再次使用 m_2。由于所有 FFN 块都是嵌套的，因此这些混合模型无需重新训练、微调或重新定义架构。 这种逐层配置的可能性非常巨大。对于 L 层和 g 个粒度，存在 g^L 个可能的子模型。即使 L 和 g 的值很小，例如 32 层和 4 个粒度，也有数十亿种选择。 然而，并非所有组合都同样有效。MatFormer 的作者提出了一个简单的启发式方法来选择良好的配置：
- 平滑过渡： 倾向于选择 FFN 宽度在层之间变化缓慢且一致的子模型。避免在相邻层之间出现从小粒度到大粒度的剧烈跳跃。
- 单调递增： 倾向于使用单调或缓慢递增的序列，例如，在早期层中使用 m2，并在后续层中逐渐增加到 m3。
这种启发式方法与模型的训练方式相符：在训练过程中，抽样的子模型始终使用统一的宽度。因此，虽然混合的逐层组合从未经过显式训练，但与训练方式相似的配置具有更好的泛化能力。
MatLM 与 MatViT：跨模态应用
MatFormer 的作者将这种方法应用于两种基于 Transformer 的模型：LLM (MatLM) 和 ViT (MatViT)。 对于使用 MatFormer 的嵌套 FFN 训练的 LLM，所有粒度的子模型都优于单独或独立训练的类似大小的 vanilla Transformer LLM。 在 MatLM 的实验中，他们训练了 4 个嵌套粒度，其 d_ffn / d 比率为 {0.5, 1, 2, 4}，并将相应的模型命名为 S、M、L 和 XL。这里，XL 是完整模型。 即使 MatFormer 论文中训练的最大 Transformer 只有 850M 个参数，Gemma 3n 使用 MatFormer 构建 4B (有效 5B) 参数的模型，也表明了这种方法的可扩展性。 MatViT 在图像检索方面也展现出了独特的优势。典型的图像检索流程是，离线获取图库图像的 embedding 并存储，而实时创建查询图像的 embedding。这意味着图库 embedding 可以使用更大、更准确的模型，而查询 embedding 可能需要更小的模型以进行快速推理。由于 MatFormer 中的较小子模型是嵌套的并且联合训练，因此它们倾向于与较大模型共享相同的 embedding 空间，从而导致基于最近邻检索的有意义的距离。这使得可以使用 XL 模型进行图库编码，同时使用较小模型（例如 S 模型）进行查询。
Mix’n’Match 模型的性能优势
MatFormer 的一个重要优势体现在“Mix’n’Match”模型在精度与计算效率之间的权衡。这些模型往往位于 Pareto 最优曲线上，意味着在给定的计算资源下，可以创建尽可能精确的 Mix’n’Match 模型。 这对于实际部署具有重要意义。例如，在许多部署场景中，我们可能有足够的计算资源（GPU 内存）来加载一个 20B 参数的模型，但最接近的现有模型要么是 14B，要么是 33B。使用 14B 参数的模型会导致资源利用不足和性能下降，而 33B 参数的模型则无法放入内存。Mix’n’Match 方法允许我们免费创建接近完美尺寸的模型变体，并获得相应的精度提升。 Gemma 3n 基于 MatFormer 架构，使我们能够执行 Mix’n’Match 并创建更多介于 2B 和 4B 之间的模型尺寸。在性能方面，Gemma 3n E4B 的 MMLU 准确率为 62.3%，而 E2B 的准确率为 50.90%。Mix’n’Match 后的变体，例如 2.54B 的 MMLU 准确率为 55.40%，而 2.69B 的准确率为 57.70%，这明显优于 2B，略低于 4B。这使我们能够为设备端部署用例（智能手机、平板电脑、Mac 等）获得大小和精度都完美匹配的模型。 为了尝试创建 Gemma 3n 模型的 Mix’n’Match 变体，Gemma 3n 团队发布了 MatFormer Lab，这是一个 Colab notebook，您可以在其中对模型进行切片并将其推送到 Hugging Face。他们还发布了一些最佳切片配置（正如我们上面看到的，虽然我们可以生成 g^L 个模型，但并非所有 Mix’n’Match 都可能对于所需的尺寸而言是最佳的）在 hugging face 上 – google/gemma3n-slicing-configs。这些配置包括在层级别（更细粒度）或在块级别（包括 4 个局部层和 1 个全局层）更改 FFN 维度。
一致性与推测解码
MatFormer 作者关注的另一个方面是，将较小的模型作为较大模型的嵌套子模型进行训练，可以提高预测 token 的一致性。验证这种一致性的一种方法是查看推测解码的拒绝率。 在推测解码中，不是始终从大型目标模型中自回归地预测 token，而是存在一个较小的“草稿模型”和一个较大的“验证模型”（目标模型）。草稿模型执行 token 推测，即，它预测接下来的 n 个 token 及其概率，而目标模型并行地执行对预测 token 的验证。通过比较目标模型对草稿模型预测的相同 token 的概率，来执行拒绝抽样。当草稿模型与目标模型显着偏离时，草稿模型的预测将被拒绝。 可以看出，当草稿模型始终预测与目标模型相似的 token 分布时，可以获得最大的加速，从而减少拒绝。如果拒绝更多，那么我们将无法获得推测解码所期望的期望的推理加速。MatFormer 的作者表明，由于较小的子模型嵌套在较大的 (XL) 模型中并与其一起训练，因此它们显示出草稿和验证器模型所需的更一致的特性。 当使用独立训练的“S”大小的模型进行传统推测解码时，目标模型与 XL 模型相比，目标模型的自回归解码的推理速度提高了 10%，而使用“S”大小的 MatLM 模型则速度提高了 16%。 此外，由于两个 MatLM 模型（S 和 XL）是共同训练的，因此它们允许共享注意力缓存（这对于独立训练的模型来说是无法实现的，因为潜在表示会非常不同）。 我发现更有趣的是，在部署期间，通用 XL 模型可以存储在内存中，并根据计算约束用于提取自适应的较小模型。这种动态资源自适应模型提取和推理可能对边缘用例非常有用。
MatFormer 与 MoE 的区别
几乎所有主要的最新开放 LLM Transformer 都是 MoE 架构。混合专家 (MoE) 模型通过向每个 Transformer 层添加多个 FFN（称为专家）来增加容量。在推理时，路由器决定为每个 token 使用哪些专家，并且仅执行那些选定的路径。这降低了每次前向传递的计算成本，因为大多数专家都保持非活动状态。但是，所有专家权重仍然需要加载到内存中，因为路由器可以在运行时选择其中任何一个。这使得 MoE 模型难以部署在内存受限的设备上——虽然计算是稀疏的，但内存使用率仍然很高。 相比之下，MatFormer 不依赖于路由或多个专家。相反，每个 FFN 都设计为通过嵌套较小的子网络来支持多个粒度。在推理时，可以选择特定的子模型并独立运行。只需要将该子模型的参数加载到内存中。这使得 MatFormer 非常适合于设备端或低内存推理，在这些情况下，存储完整模型是不可行的。
总结
MatFormer 代表了我们在模型训练、推理和部署方面的一种重大转变。传统上，大型模型在大量的计算预算下进行训练，然后创建较小的变体——要么单独训练，要么从较大的变体中提取。MatFormer 通过共享权重训练将多个功能子模型嵌套在单个 Transformer 中来打破这种模式，从而消除了对昂贵的蒸馏、重新训练或像混合专家这样的复杂路由的需求。这种设计释放了一个平滑的资源自适应模型连续体——从轻量级的、可用于移动设备的部署到全容量推理引擎——所有这些都来自一个基本模型。 Gemma 3n 的 E2B 和 E4B 的成功、Mix’n’Match 变体的多功能性以及 MatViT 编码器展示的跨模态通用性表明，MatFormer 风格的架构可能成为研究实验室发布可扩展模型系列的新默认方式。 随着大模型技术的持续演进，MatFormer 无疑将成为推动模型小型化、高效推理和广泛应用的关键技术之一。它的出现，预示着大模型不再是高高在上的空中楼阁，而是可以灵活定制、触手可及的工具。