soft-moe-pytorch

Soft MoE - Pytorch

在Pytorch中实现<a href="https://arxiv.org/abs/2308.00951">Soft MoE（专家混合）</a>，由Brain的Vision团队提出。

这个MoE仅适用于非自回归编码器。然而，最近一些<a href="https://arxiv.org/abs/2305.18295">文本到图像模型</a>已经开始使用MoE并取得了很好的效果，因此可能适用于这些场景。

如果有人对如何使其适用于自回归有任何想法，请通过电子邮件或讨论告诉我。我思考了很久，但想不出好的方法。插槽方案的另一个问题是，随着序列长度的增加，路由会受到二次方的影响（很像注意力机制）。

致谢

感谢<a href="https://stability.ai/">StabilityAI</a>的慷慨赞助，以及其他所有赞助商
感谢<a href="https://github.com/arogozhnikov/einops">Einops</a>让我的工作变得轻松

安装

$ pip install soft-moe-pytorch

使用方法

import torch
from soft_moe_pytorch import SoftMoE

moe = SoftMoE(
    dim = 512,         # 模型维度
    seq_len = 1024,    # 最大序列长度（将自动计算插槽数量为seq_len // num_experts）- 你也可以直接设置num_slots
    num_experts = 4    # 专家数量 - （他们建议专家数量应该足够高，以使每个专家只获得1个插槽。不知道这是否是论文的弱点？）
)

x = torch.randn(1, 1024, 512)

out = moe(x) + x # (1, 1024, 512) - 在某一层的前馈网络位置添加transformer（这里同时展示了残差连接）

对于一个即兴变体，可以使用动态插槽，使插槽数量≈序列长度，只需导入DynamicSlotsSoftMoe即可

import torch
from soft_moe_pytorch import DynamicSlotsSoftMoE

# 无需指定序列长度或插槽数量

moe = DynamicSlotsSoftMoE(
    dim = 512,         # 模型维度
    num_experts = 4,   # 专家数量
    geglu = True
)

x = torch.randn(1, 1023, 512)

out = moe(x) + x # (1, 1023, 512)

待办事项

解决插槽数量固定的限制。考虑根据序列长度动态设置插槽数量的方法
一旦在分布式环境中处理了可变序列长度，将其添加到动态软MoE中
分发和组合张量也可以拆分并移入Experts类中，以更好地分配工作

引用

@misc{puigcerver2023sparse,
    title 	= {From Sparse to Soft Mixtures of Experts}, 
    author 	= {Joan Puigcerver and Carlos Riquelme and Basil Mustafa and Neil Houlsby},
    year 	= {2023},
    eprint 	= {2308.00951},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG}
}

@misc{shazeer2020glu,
    title   = {GLU Variants Improve Transformer},
    author  = {Noam Shazeer},
    year    = {2020},
    url     = {https://arxiv.org/abs/2002.05202}
}