SplitMoE：破解视频扩散模型中MoE的均匀性陷阱

读 SplitMoE 这篇文章的时候，我实验室里正好有一个视频生成项目在硬吃显存。跑一个 45 帧、720p 的扩散模型采样要将近 14GB 显存，其中一大部分没花在 attention 上，反而被逐帧共享的 FFN 给吞噬了。所以看到“视频扩散模型 + MoE”这种组合，我第一反应是：终于有人来收拾 FFN 了。但看完 NeurIPS 2026 Spotlight 版本的 SplitMoE，我发现它真正解决的问题比我预想的更微妙。它不是简单地把 DiT 里的 MLP 换成一堆专家，而是先把视觉 token 的“路由依据”重新定义了一遍——也就是标题里那个关键词：均匀性陷阱。这篇精读，我会把它的动机、架构设计和工程细节完整拆开讲，顺便把我自己复现时踩过的坑一并写出来。

1. 视频扩散模型的计算瓶颈：为什么大家盯上 MoE

1.1 视频扩散到底贵在哪里

视频扩散模型的去噪网络，这几年基本都收敛到 3D DiT 或基于 transformer 的 U-Net 结构。以 DiT 为例，输入是一段视频的若干帧，每帧被切成 16×16 的 patch，经过 patch embedding 后变成一串 token。假如输入是 16 帧 256×256 的视频，patch size 仍是 2×16×16，那么 token 数量约为 16×(256/16)×(256/16)=4096 个。这还只是短视频。如果拉到 64 帧、512×512，token 数直接到三万以上。

这个长度对 transformer 是致命的。虽然很多实现用 spatial-temporal 分离 attention 来缓解平方复杂度，但真正的问题不只是 attention——每一层 FFN 的 FLOPs 也是随 token 数线性增长的。attention 至少还有各种窗函数、flash attention 可以救，FFN 几乎没什么可救的，它就是每个 token 独立过一遍全连接网络，token 越多越贵。

所以视频扩散模型的训练和采样成本，本质上是一个“长序列 + 逐 token 大计算”的问题。一个个优化下来，FFN 逐渐成了大头。我实测过一个中等规模的 DiT 视频模型，FFN 部分的 FLOPs 占整体去噪网络的比例大概在 45%~55% 之间，取决于 attention 的窗口配置。也就是说，哪怕 attention 已经理想地逼近线性，你的显卡还是有一半时间浪费在这个看起来“没什么技术含量”的 MLP 上。

1.2 MoE 看起来很合适，但直接搬过去就出问题

MoE 的思路在 LLM 里已经被验证过很多次。把 FFN 替换成多个专家 FFN（通常 8 个或 16 个），每个 token 只激活 top-k 个专家。参数量上去了，单 token 的计算量却基本可控，因为稀疏激活的本质就是“用更多参数覆盖多样化的输入模式，同时只付出 k/E 的计算成本”。

这个思路放在视频扩散模型里太理想了：视频内容本身高度多样，有静态背景、有大动态运动、有不同光照，理论上天然适合让不同的专家去负责不同的“场景片段”。当时我第一版实验就是照着 LLM 的做法改的：把每个 DiT block 的 MLP 换成 8 个专家，top-2 激活，另外加一个标准的 load balancing loss。

结果比较难看。训练初期 FVD 还有一点改善，但步数一长，loss 就开始震荡，生成质量不升反降。最反常的是路由分布——我盯着 tensorboard 上的专家利用率看到，所有专家的被选概率都在 12% 附近抖动，极其均匀。这听起来不是好事吗？专家都忙起来了，负载均衡也完美。但生成质量告诉我：这种均匀不是健康的分工，而是一种“不知道该把 token 分给谁”的和稀泥。这就是 SplitMoE 里说的均匀性陷阱。

2. “均匀性陷阱”是一种更加隐性的失败：路由很均匀，模型却没变强

2.1 路由熵增的表象与机理

先说现象。普通 MoE 里我们通常监控三个指标：路由熵、专家利用率、平均 top-1 置信度。naive MoE 训练一段时间后，路由熵会迅速爬升到接近 log(E)，也就是每个专家的被选概率趋于 1/E。专家利用率曲线倒是很漂亮，几乎没有任何 expert 饿死，但如果你去看 router 输出的 logits，会发现它们几乎是同一个值，完全没有置信度可言。

问题出在视频 patch token 的特征分布上。语言模型里的 token 有语义离散性，“dog”这个词和“run”这个词在词嵌入空间里的分布是有明确边界的，router 很容易学会做“分诊”。但视频扩散模型里的 patch token 是连续视觉特征的切片，每个 patch 同时包含空间纹理、局部边缘、运动趋势、光照变化，特征之间是连续过渡的，没有一个天然的“类别标签”。

在这种情况下，router 要最大化训练似然，最安全的策略反而就是把概率抹平。因为视觉特征高度混合，任何一个专家单独处理某类 token 都很难比“平均处理”更有效，于是 router 学到的“最优解”就是均匀分配。更麻烦的是，一旦 router 输出均匀，所有 token 都被所有专家用几乎相同权重处理，各专家接收到的输入分布同样均匀，也就没有动力去分化自己的权重。这是一个负反馈循环。

2.2 为什么均匀分布对视频扩散是坏消息

如果我们把 8 个专家训练后的权重拿出来看就会发现，它们之间的余弦相似度非常高，几乎退化成同一个 MLP。那跟用一个宽 FFN 没什么区别，但 MoE 还得额外付出一大堆代价：

• 计算开销更大了。专家网络虽然只激活 top-2，但每个专家的 hidden size 通常和 dense FFN 一样，所以单 token 的实际 MACs 反而比 dense 翻倍。
• 显存更大了。8 个完整专家 FFN 的参数是参数量的大头，即使不激活也得常驻显存或者频繁换入换出。
• 训练不稳定。为了维持“表面均匀”，load balancing loss 一直在和重建损失打架，loss 权重稍高一点，去噪质量就掉，稍低一点，专家利用率又开始倾斜，总有一个动态平衡要调。
• 长视频上尤其吃亏。序列越长，均匀路由带来的是大量无效计算——每个 token 都被平均地“咨询”了所有专家，却没有得到针对性的特征变换。

一句话总结：路由均匀是这个模型的“死法”，但它死得比显式崩溃更隐蔽——loss 没有爆，指标没有归零，只是用了 8 个专家的钱，买了 1 个专家的效果。

2.3 真正的问题不是“负载不均”，而是“无处可专”

现在可以回到 SplitMoE 的核心论断了。它指出，视频扩散模型里 MoE 失败的根源不是负载不均衡，而是“任务分配缺乏结构依据”。负载均衡只是结果，不是目标。你不可能靠一个辅助 loss 让 router 产生专业化分工，因为 router 的输入本身就缺乏一个线性可分的分界。

你看 LLM 的 MoE 为什么能work，本质是因为 token 所属的语义空间存在大量直线可分的子区域。视频 patch 也有空间和时间两个维度，但它们在 embedding 里是纠缠在一起的。SplitMoE 的动机非常直接：与其指望 router 在纠缠的特征里自己画出“空间专属区”和“时间专属区”，不如直接动手把 embedding 切成两半，把空间主导的成分交给空间专家，把时间主导的成分交给时间专家。这样路由不再是自由发挥，而是被约束在了一个更合理的结构里。

3. SplitMoE 的核心设计拆解：先 Split，再路由

3.1 按通道拆分：把“空间纹理”和“时间运动”分给不同专家

SplitMoE 不是在 token 维度上做文章的，也不是在时间轴上的动态对齐，而是在每个 token 的 hidden embedding 的通道（channel）维度上做拆分。

具体来说，它会设定一个拆分比例 r（比如 0.5）。对于每个 token 的 hidden state h，拆成 h_spatial = h[:, :, : int(dr)] 和 h_temporal = h[:, :, int(dr):]。前者被送入空间专家组，后者被送入时间专家组。两组专家完全独立，各自有自己的 router，每一组内部再走 top-k 稀疏激活。

这个设计第一次看会有点反直觉：通道维而已，真能区分空间和时间吗？实际上视频 DiT 的通道在训练过程中确实会逐步出现频率分层。靠前的通道偏向低频结构（物体的轮廓、背景布局），靠后的通道偏向高频细节与短期运动残差。虽然这种分层不是绝对干净的，但 SplitMoE 只需要它提供一个大致的结构性先验，不需要完美的因果解耦。

我理解 SplitMoE 能work的关键在于：它把“专家选择”从 token level 的不可区分任务，降维成了“子空间内特征变换”的可区分任务。空间专家只需要在 h_spatial 这个子空间里做映射，时间专家只需要在 h_temporal 里做映射，它们的输入分布天然不同，分化压力自热就会出现。

3.2 Split-Aware Gating 与组内稀疏激活

路由部分，SplitMoE 用了一个我称之为“split-aware gating”的结构。每组内部的 router 只看自己那一段子向量，不看全量 hidden state。这个细节很重要——如果 router 输入全量 embedding，它很可能在某个训练阶段重新学会把不同 sub-space 的信息混合，退化成原来的均匀路由。只拿子向量做路由，本质上是用信息隔离来防止“和稀泥”。

一个简化的 PyTorch 风格伪代码长这样：

def split_moe_block(x, r=0.5, top_k=1, num_experts=8): d = x.shape[-1] split_idx = int(d * r) x_s, x_t = x[..., :split_idx], x[..., split_idx:] # 每组独立的 router 与 expert FFN y_s = route_and_activate(x_s, spatial_router, spatial_experts, top_k) y_t = route_and_activate(x_t, temporal_router, temporal_experts, top_k) # 共享专家接收完整输入，捕捉公共基础模式 y_shared = shared_expert(x) return y_s + y_t + y_shared

这里 top_k 通常取 1 或 2。取 1 的时候极稀疏，每组只激活一个专家，加上共享专家，总激活专家数远小于 naive MoE，计算收益最明显。取 2 的时候表达能力更强，适合视频内容复杂度很高的主实验。默认配置里我看到的做法是空间组和时间组都是 4 expert、每组 top-1，另带一个共享专家，等于每层每 token 只激活 3 个专家 FFN，其中两个的输入维度还减半了，总计算量非常可控。

共享专家的存在也值得一提。纯空间专家和纯时间专家毕竟各自只看到了一半通道，总有些通用模式（比如基础边缘、整体亮度变化的公共处理）不适合被硬塞进某一边。共享专家负责兜底，它接收完整 hidden state，输出直接加到最终结果上。

3.3 与 Naive MoE 的架构对比：一张表格

我把这两者放在一起看会更直观。下表是我根据论文配置整理的一个对比，数字按 hidden size d=1152、FFN hidden 4608 估算，具体超参与原始论文配置为准：

这里有一点要特别注意：naive MoE 的 MACs 不一定真的低于 dense，因为它激活两个完整尺寸的专家 FFN，每个专家的 hidden size 和 dense FFN 一样大，总计算量翻倍。SplitMoE 把输入拆半之后，两个专家加起来约等于原来的一个 dense FFN，共享专家再补一点，计算量才真正控制住。这也是这篇文章能在视频扩散模型里落地的直接原因。

4. 效果评测与实现细节：训练收敛、显存足迹、推理吞吐

4.1 评测指标与主要结果

论文的主实验我简单梳理一下，方便读者对着自己的场景判断。他们主要是在 UCF101 和两个文本到视频生成基准上做的评测，指标是 FVD、IS、CLIPSIM 这一类。和同 FLOPs 的 dense DiT 相比，SplitMoE 在 FVD 上的下降幅度我记得比较明显，大约接近 10%~15% 的量级，具体数字还是要看论文原始表格。

更有意义的是训练曲线。SplitMoE 的收敛曲线在最开始几乎和 dense 重合，训练到中期开始分开，且没有出现 naive MoE 那种明显的 loss 震荡。这个现象我认为是设计成功的信号：additive 结构的共享专家保证了训练的稳定性，而空间/时间两个独立专家组的传入特征不同，自然会让各自专家的权重走向分化。没有震荡，意味着 load balancing loss 和重建 loss 的矛盾被结构性地缓解了。

在推理侧，SplitMoE 主要收益来自激活专家数减少和每个专家的输入维度减半。同一块 A100 上，跑 60 帧 512×512 的采样，显存峰值比 dense DiT 大概能省下接近 30%~40%，采样吞吐提升明显。这个提升在长序列扩散模型里非常关键，因为很多现场 demo 卡脖子就卡在采样时显存爆炸。

4.2 训练侧的 load balancing 与稳定性

训练稳定性是视频 MoE 最容易翻车的地方，SplitMoE 也没有完全放弃辅助 loss。它的做法是改成了 split-aware 的约束：空间组和时间组分别计算各自的 expert-level load balance loss，而不是在整个专家池上算一个全局均匀项。

比如空间组有 4 个专家，最后的负载比例可以是不均匀的，因为视频里静态场景和动态场景的比例本来就不同；时间组同理。因此它不再要求“所有专家使用率一致”，而是约束“每个专家不要饿死或过载”到一个合理的区间。这个改动让 aux loss 的权重可以开得比较低，对重建损失的干扰明显减少。

我复现的时候还额外加了两个技巧，强烈建议你也加上。一个是 router z-loss，目的是让 router 的 logits 幅度不要无限涨，数值稳定很多；另一个是 warmup：训练的前几千步固定 router 输出为均匀分布，让专家 FFN 先适应基本特征，之后再放开 router 学习。直接从头训练 router 的话，前期梯度很容易把 router 推到某种局部死区里。

4.3 推理侧的专家并行与内存权衡

推理性价比是 SplitMoE 真正值钱的地方，但也需要工程配合。专家参数全量常驻显存的话，8 个专家的参数还是会占不少显存。SplitMoE 的解法在我看来也很务实：因为空间组和时间组互不依赖，可以天然拆到不同的设备上；空间专家放在一张卡，时间专家放在另一张卡，token 的前半段和后半段分别走两条设备，最后通过 all-reduce 拼回来。这比在 token 维度上做 expert parallelism 的通信量小得多。

如果只有单卡，那就继续走专家缓存策略：把不活跃的专家参数换到 CPU 内存里，激活前再 prefetch 回来。SplitMoE 每组只有 4 个专家，且组内 top-1，每次只需要 prefetch 一个专家，比 naive 的 top-2 少了一半的换入压力。实际跑下来，单卡推理时 CPU-GPU 的 H2D 传输占比明显变小。

不过我建议你先别急着双卡部署。如果你的视频长度比较短、token 序列在 1000 以下，单卡把全部专家都常驻显存跑反而更省心，因为 expert prefetch 的管理开销可能抵消掉稀疏激活的收益。长序列、大视频场景下再上专家并行，收益才会非常明显。

5. 迁移到自己项目之前的复盘：适用边界与复现避坑清单

5.1 什么情况下 SplitMoE 才值得上

SplitMoE 不是什么场景都能救。我总结下来，至少满足以下条件时，收益才会明显大于成本：

• 你的视频 token 序列足够长（一般至少 2000 token 以上），否则 attention 和 FFN 的比例决定了 FFN 并不是瓶颈。
• 视频内容本身的时空复杂度差异大。如果全是固定机位的静态场景，空间专家和时间专家几乎没有分化空间。
• 你有充足的显存或者能接受多卡专家并行，否则参数量的膨胀会很尴尬。
• 你的基准任务是生成质量导向，而不是分类或轻量推理导向。

反过来，如果只是做短视频分类或者很小的视频 patch 扩散模型，SplitMoE 大概率帮倒忙。参数量上去了，训练也变慢，最后的收益完全无法覆盖成本。

5.2 复现时的关键超参与通行做法

我复现时把关键的几个配置整理出来了，按优先级排：

1. split ratio r：默认 0.5 比较稳，但建议先用 0.3 跑通再调，高 r 意味着空间专家承担更多通道，需要更强的空间特征表达。
2. 每组专家数：4 个够用，8 个时分化会更细，但对数据量要求也更高。
3. top_k：默认 1，表达能力不够再加到 2，但每个 token 激活的专家数翻倍，计算收益就少了。
4. 共享专家占比：共享专家的 hidden 一般设为专家 FFN 的 1/4 到 1/2 即可，太大反而会吸收过多信息，专家组失去意义。
5. aux loss 权重：SplitMoE 通常可以比 naive MoE 小 5~10 倍，我建议从 0.01 起步，观察专家利用率再过调。

一个特别容易忽略的地方是 patch embedding 的方式。如果你用的是 3D patch，空间和时间信息在 patch embedding 阶段就已经混合了一部分，通道拆分的效果会打折扣。我更建议在现有实现里先试试帧级 patch（即 2D patch + time embedding），跑通后再尝试 3D patch 版本对比。这不是论文里的原话，是我的个人实战经验，但确实会影响最终效果。

5.3 几个容易踩的坑

最后说说坑，这些我都是真金白银喂出来的。

第一个坑是路由熵监控缺失。如果你只盯着 FVD 看，完全不知道 router 是否又退化到均匀分布，那你可能跑到一半才发现白练了。训练时务必把每个 group 的路由熵、每个专家的被选概率记下来。一旦熵过快逼近 log(E)，先检查是不是 split 比 r 设置不合理。

第二个坑是梯度裁剪。视频扩散模型训练本身梯度噪声就大，加了 MoE 之后，router 和 expert 的梯度尺度不一致，一个简单的全局 grad clip 往往不够。我最后是把 grad clip max norm 从 1.0 降到了 0.5 左右，再配合 freeze router 的 warmup，才算稳下来。

第三个坑是 auxiliary loss 对不同视频长度的行为差异。我跑 16 帧短视频和 64 帧长视频时，同样的 aux loss 权重，短视频下一切正常，长视频下第一轮评估时 FVD 突然恶化。原因大概是长视频里时间专家的负载天然更低，全局均衡约束在长序列上过于苛刻。这也印证了论文中 split-aware 分组的必要性——别对整个专家池做统一均衡，一定要分组约束。

第四个容易被低估的是 activation memory。MoE 专家虽然不全被激活，但每个专家都可能有中间 buffer，反向传播时这些 buffer 会全部进入计算图。我一开始只关注参数显存，结果发现 activation 反而成了瓶颈。解决办法是尽量使用显存高效的 expert linear 实现，或者在前向时手动 detach 掉非激活专家分支的计算图。

这类结构性的问题，往往比调参本身更能决定项目生死。我把这篇论文反复看了几遍之后最大的感触是：MoE 用在视频扩散模型上，厉害的不是“稀疏激活”这个外壳，而是你能不能给路边唯一一个可靠的“分工依据”。SplitMoE 用通道级拆分回答了一个比模型架构更基本的问题——视频 token 凭什么被分到不同专家。带着这个视角再去看你自己的视频生成项目，应该会有完全不一样的判断。