Multilinear Mixture of Experts: Scalable Expert Specialization through Factorization

📄 arXiv: 2402.12550v4 📥 PDF

作者: James Oldfield, Markos Georgopoulos, Grigorios G. Chrysos, Christos Tzelepis, Yannis Panagakis, Mihalis A. Nicolaou, Jiankang Deng, Ioannis Patras

分类: cs.CV, cs.LG

发布日期: 2024-02-19 (更新: 2024-10-16)

备注: Accepted at NeurIPS 2024. Github: https://github.com/james-oldfield/muMoE. Project page: https://james-oldfield.github.io/muMoE

🔗 代码/项目: GITHUB

💡 一句话要点

提出多线性专家混合模型以解决专家特化的计算成本问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家 计算机视觉 模型特化 因式分解 深度学习 模型微调 CelebA分类

📋 核心要点

  1. 现有的混合专家方法在扩展专家数量以实现细粒度特化时面临高昂的计算成本。
  2. 本文提出的μMoE层通过因式分解形式的隐式计算，解决了专家特化的可扩展性问题。
  3. 实验结果显示，μMoE层在视觉任务中实现了更高的专家特化，并在CelebA分类中支持偏差修正。

📝 摘要（中文）

混合专家（MoE）范式提供了一种强大的方式，将密集层分解为更小的模块化计算，便于人类理解、调试和编辑。然而，随着专家数量的增加，计算成本显著上升，限制了细粒度特化的实现。本文提出了多线性混合专家（μMoE）层，专注于视觉模型，通过在因式分解形式下对大权重张量进行隐式计算，实现可扩展的专家特化。μMoE层避免了密集MoE在推理时的高成本，同时不继承稀疏MoE在训练中面临的离散专家路由问题。实验结果表明，μMoE层在视觉任务的基础模型微调中能够实现更具特化性的专家，并支持CelebA属性分类中的手动偏差修正。

🔬 方法详解

问题定义：现有的混合专家（MoE）方法在扩展专家数量时，推理时间成本过高，限制了其在细粒度特化中的应用。同时，稀疏MoE的离散专家路由在训练过程中存在不易优化的问题。

核心思路：本文提出的多线性混合专家（μMoE）层通过因式分解的方式，隐式计算大权重张量，从而实现专家的可扩展特化，避免了高推理成本和训练中的离散性问题。

技术框架：μMoE层的整体架构包括多个专家模块，每个模块通过因式分解的权重表示进行计算。该框架允许在不增加推理成本的情况下，动态选择和激活不同的专家。

关键创新：μMoE层的核心创新在于其因式分解的设计，使得在推理时能够有效地处理大规模权重张量，同时保持训练的可微性。这一设计与传统的稀疏MoE方法形成了鲜明对比。

关键设计：在μMoE层中，关键参数设置包括专家数量、因式分解的维度以及损失函数的设计。通过精细调节这些参数，能够在保持模型准确性的同时，实现专家的高效特化。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，μMoE层在视觉任务的基础模型微调中，专家特化程度显著提高，尤其在CelebA属性分类中实现了手动偏差修正。此外，使用μMoE块的GPT2和MLP-Mixer模型在保持相似准确率的同时，展现了更强的专家特化能力。

🎯 应用场景

该研究的潜在应用领域包括计算机视觉、自然语言处理等需要专家特化的任务。μMoE层的设计能够在不显著增加计算成本的情况下，提高模型的表现，具有重要的实际价值和广泛的应用前景。

📄 摘要（原文）

The Mixture of Experts (MoE) paradigm provides a powerful way to decompose dense layers into smaller, modular computations often more amenable to human interpretation, debugging, and editability. However, a major challenge lies in the computational cost of scaling the number of experts high enough to achieve fine-grained specialization. In this paper, we propose the Multilinear Mixture of Experts ($μ$MoE) layer to address this, focusing on vision models. $μ$MoE layers enable scalable expert specialization by performing an implicit computation on prohibitively large weight tensors entirely in factorized form. Consequently, $μ$MoEs (1) avoid the restrictively high inference-time costs of dense MoEs, yet (2) do not inherit the training issues of the popular sparse MoEs' discrete (non-differentiable) expert routing. We present both qualitative and quantitative evidence that scaling $μ$MoE layers when fine-tuning foundation models for vision tasks leads to more specialized experts at the class-level, further enabling manual bias correction in CelebA attribute classification. Finally, we show qualitative results demonstrating the expert specialism achieved when pre-training large GPT2 and MLP-Mixer models with parameter-matched $μ$MoE blocks at every layer, maintaining comparable accuracy. Our code is available at: https://github.com/james-oldfield/muMoE.