[论文解读] Spatial Mixture-of-Experts

Nikoli Dryden, Torsten Hoefler|arXiv (Cornell University)|Nov 24, 2022

Human Mobility and Location-Based Analysis被引用 5

一句话总结

本文提出空间混合专家（SMoE）层，一种稀疏门控神经网络组件，通过使用张量路由函数在像素级别对专家进行路由，从而学习数据中的细粒度空间结构。该方法通过引入空间归纳偏置并结合自监督路由损失和专家误差抑制，在中等范围天气预测与集合后处理任务中实现了最先进性能，优于卷积神经网络和标准MoE模型，在密集回归任务中表现更优。

ABSTRACT

Many data have an underlying dependence on spatial location; it may be weather on the Earth, a simulation on a mesh, or a registered image. Yet this feature is rarely taken advantage of, and violates common assumptions made by many neural network layers, such as translation equivariance. Further, many works that do incorporate locality fail to capture fine-grained structure. To address this, we introduce the Spatial Mixture-of-Experts (SMoE) layer, a sparsely-gated layer that learns spatial structure in the input domain and routes experts at a fine-grained level to utilize it. We also develop new techniques to train SMoEs, including a self-supervised routing loss and damping expert errors. Finally, we show strong results for SMoEs on numerous tasks, and set new state-of-the-art results for medium-range weather prediction and post-processing ensemble weather forecasts.

研究动机与目标

• 为解决标准神经网络在捕捉具有位置结构数据中的细粒度空间依赖关系方面的局限性。
• 克服卷积神经网络与标准混合专家模型在密集回归任务中难以学习位置依赖模式的问题。
• 开发一种路由机制，使专家能够针对不同空间位置实现专业化，同时保持计算效率。
• 通过新颖的训练技术引入空间归纳偏置，提升天气预报任务的模型性能。

提出的方法

• 提出空间混合专家（SMoE）层，利用可学习的门控函数在像素级别稀疏地将输入路由至共享专家组。
• 引入张量路由，一种轻量且高效的路由函数，显式建模输入域中的空间结构。
• 采用自监督路由分类损失，通过预测路由错误为门控提供更强的学习信号。
• 应用专家误差抑制技术，抑制错误路由样本的梯度更新，提升模型鲁棒性与性能。
• 在具有相似空间维度的层之间共享路由门控，以减少参数量并提升泛化能力。
• 将SMoE层嵌入U-Net与ResNet架构中，替换标准卷积层，实现空间专业化。

实验结果

研究问题

• RQ1 混合专家架构能否有效学习具有位置结构数据中的细粒度空间依赖关系？
• RQ2 与粗粒度路由或标准卷积相比，MoE中的细粒度像素级路由是否能在密集回归任务中带来性能提升？
• RQ3 当端到端路由无法捕捉空间结构时，自监督路由损失是否能改善门控训练？
• RQ4 在存在路由错误的情况下，专家误差抑制如何影响模型收敛与性能？
• RQ5 SMoE层是否能在多样任务（如天气预测与图像分类）中实现良好泛化，同时相比先前方法使用更少参数？

主要发现

• SMoE在WeatherBench基准上实现了中等范围天气预测的最先进结果，优于先前的SOTA方法。
• 在ENS-10集合天气预报后处理数据集上，SMoE在所有预报时效与集合规模下均创下SOTA记录，CRPS与EECRPS评分更低。
• 在具有位置依赖扩散系数的受控热传导任务中，SMoE成功学习了空间结构，而标准CNN与MoE模型则收敛至平均扩散系数。
• 在ImageNet-1k数据集上，当SMoE替换所有3×3卷积时，仅使用56%的参数量，但保持了具有竞争力的准确率，优于LRLCN模型。
• 张量路由、自监督路由损失与专家误差抑制的结合显著提升了性能，尤其在需要细粒度空间建模的任务中表现突出。
• SMoE能有效扩展至大空间区域（如全球天气数据），且性能与效率不下降。