soft-moe-pytorch

Install the package using pip:

$ pip install soft-moe-pytorch

The SoftMoE class implements the Soft Mixture of Experts as proposed by Brain's Vision team. It is designed for non-autoregressive encoders.

When using SoftMoE, you can specify seq_len (which automatically calculates the number of slots as seq_len // num_experts) or set num_slots directly.

import torch
from soft_moe_pytorch import SoftMoE

moe = SoftMoE(
    dim = 512,         # model dimensions
    seq_len = 1024,    # max sequence length (will automatically calculate number of slots as seq_len // num_experts) - you can also set num_slots directly
    num_experts = 4    # number of experts
)

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

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

If you need to handle variable sequence lengths where the number of slots is not fixed, use DynamicSlotsSoftMoE. In this variant, the number of slots is dynamic (approximately equal to the sequence length), so you do not need to specify seq_len or num_slots.

import torch
from soft_moe_pytorch import DynamicSlotsSoftMoE

# sequence length or number of slots need not be specified

moe = DynamicSlotsSoftMoE(
    dim = 512,         # model dimensions
    num_experts = 4,   # number of experts
    geglu = True
)

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

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

SoftMoE

Used for fixed-slot MoE implementations.

Arguments:

• dim (int): Model dimensions.
• seq_len (int, optional): Max sequence length. If provided, num_slots is calculated as seq_len // num_experts.
• num_experts (int): Number of experts.
• num_slots (int, optional): Direct specification of the number of slots.

DynamicSlotsSoftMoE

An improvised variant that supports dynamic slots so that num_slots $\approx$ sequence length.

Arguments:

• dim (int): Model dimensions.
• num_experts (int): Number of experts.
• geglu (bool): Whether to use GEGLU activation.