Research Papers
EN ZH

专家混合 · 高效 AI

Switch Transformer:每个词只走一个专家,参数冲上万亿

Switch Transformer 把混合专家简化为每个词只路由到一个专家,在同等算力下让 T5 预训练最高快 7 倍,并用 bfloat16 训练把稀疏模型扩到 1.6 万亿参数。

Google Research · 论文发表 2021-01-11 · 3 分钟读完

Switch Transformer:每个词只走一个专家,参数冲上万亿

快速答案

Switch Transformer 是把稀疏混合专家(MoE)真正简化到能大规模训练的那篇论文:它不再把每个词发给好几个专家,而是只路由给一个专家。Google Research 靠这招在同等算力、同样数据下把 T5 预训练加速最高 7 倍,并训出了高达 1.6 万亿参数的稀疏模型——而每个词的计算量(FLOPs)几乎不变,因为每个词只触发一个专家。

每个词只走一个专家

普通 Transformer 对所有词复用同一套前馈权重。MoE 则把这一层前馈换成许多并行的“专家”前馈网络,再加一个路由器决定每个词交给哪些专家处理。2021 年之前的共识——可追溯到 Shazeer 2017 年的 MoE——是必须用 top-k 路由、且 k ≥ 2,即每个词至少发给两个专家,这样路由器才有可用的梯度,选择也不会塌缩。

Switch Transformer 的核心赌注是:k = 1 就够了。路由器对所有专家做 softmax,把词只发给得分最高的那一个,并用路由概率给该专家输出加权,梯度照样能传。这个“Switch 层”就是全部诀窍,而它一次带来三重收益:路由计算减半,设备间搬运的数据大致减半,实现也大幅简化——每个词只有一个去处,而不是数量不定的一小撮。

让它稳下来的几个技巧

稀疏模型素来以难训著称,Switch Transformer 真正的贡献是把这一点修好的工程手艺:

  • 选择性精度。 不稳定主要来自路由器在 bfloat16 下的指数运算。办法是只把路由器这一段局部转成 float32,其余仍留在 bfloat16。这让大型稀疏模型首次能在低精度下训练,又不必付出全程 float32 的通信代价。
  • 容量因子与负载均衡损失。 每个专家有固定缓冲区,超出的词会被“丢弃”(经残差直通)。容量因子调节这个缓冲,辅助的负载均衡损失则逼路由器把词摊匀,而不是把几个偏爱的专家挤爆。
  • 更小的初始化与专家 dropout。 调小权重初始化、在微调时提高专家内部的 dropout,才让大得多的模型不至于发散,也不会在小的下游任务上过拟合。

关键结果

  • 预训练最高快 7 倍。 在同等算力下,以达到相同质量的实际耗时计,Switch 比 T5-Base 与 T5-Large 快最多 7 倍。
  • 1.6 万亿参数。 Switch-C 把稀疏参数扩到 1.6 万亿,而每个词的 FLOPs 大致维持在 T5-Base 水平。
  • 比 T5-XXL 快 4 倍。 在 Colossal Clean Crawled Corpus(C4)上预训练的万亿参数 Switch 模型,达到 T5-XXL 质量约快 4 倍。
  • 101 种语言全面提升。 在多语种预训练上对比 mT5-Base,Switch 版本在全部 101 种语言上都有提升,而不只是高资源语言。
  • 可蒸馏回稠密模型。 大型稀疏模型能蒸馏成稠密模型,保留约 30% 的稀疏对稠密的质量增益,同时甩掉庞大的参数量——在难以部署众多专家时很有用。

为什么现在重要

Switch Transformer 是今天 MoE 大模型一脉相承的前身。Mixtral、DeepSeek-V3、Qwen 的 MoE 模型背后那套“每个词路由给少数专家、均衡负载、把参数量扩得远超算力预算”的配方,正是 Switch 思路的精炼版。多数现代系统其实又退回了 top-2 路由以换质量,所以 Switch 最激进的主张——k = 1 就够——恰是被业界部分否定的那部分。但真正胜出的框架是 Switch 的:把稀疏当作在不增加推理成本的前提下扩容量的实用手段。要想明白为什么前沿实验室能拿出万亿参数的 MoE 模型、每词成本却和小得多的稠密模型相当,这就是让该论点变得可信的那篇论文。

局限与存疑

万亿参数的标题是参数量,不是能力主张——这些模型只是更快地追平了小得多的稠密模型的质量,并没有解锁新能力,FLOPs 对齐的对比才诚实。稀疏 MoE 还是拿算力换内存和带宽:你得存下并切分每一个专家,所以 1.6 万亿的模型即便每个词只用到其中一小块,也得用一大簇加速器才装得下权重。容量上限处的丢词意味着部分词拿不到专家;在小数据集上微调这些模型,明显容易过拟合和不稳定。而给论文命名的 k = 1,恰是后续工作最常推翻的选择——这说明单专家路由简单又快,却未必是质量最优点。

常见问题

Switch Transformer 和普通混合专家有什么不同?

早期 MoE 层把每个词路由给至少两个专家(top-k,k ≥ 2)。Switch Transformer 只把每个词路由给一个专家(k = 1),从而把路由和通信成本减半,也让这一层更易实现和训练。

Switch Transformer 有多少参数?

最大的 Switch-C 有 1.6 万亿参数。由于每个词只激活一个专家,尽管参数量惊人,每个词的计算量仍接近 T5-Base 这样的稠密模型。

Switch Transformer 比 T5 快多少?

在同等算力下,预训练比 T5-Base 与 T5-Large 最高快 7 倍;万亿参数版本达到与 T5-XXL 相当的质量约快 4 倍。

Switch Transformer 影响了 Mixtral、DeepSeek 这类模型吗?

是的。Switch Transformer 把稀疏 MoE 在语言模型上推广开来,今天的 MoE 大模型——Mixtral、DeepSeek-V3、Qwen MoE——都建立在它的负载均衡与路由机制之上,尽管多数又回到了 top-2 路由,而非 Switch 的单专家选择。

一句话:每个词只走一个专家、再均衡负载,你就能把参数扩到万亿,而在推理时无需为之买单。阅读 arXiv 原文。