Understanding Mixture of Experts (MoE)
A Beginner’s Guide with PyTorchIf you’ve ever wondered how modern AI models handle complex tasks like answering questions or recommending movies, you might have heard of the Mixture of Experts (MoE). It’s a super cool way to make neural networks smarter and more efficient.Imagine a team of…
A Beginner’s Guide with PyTorchIf you’ve ever wondered how modern AI models handle complex tasks like answering questions or recommending movies, you might have heard of the Mixture of Experts (MoE). It’s a super cool way to make neural networks smarter and more efficient.Imagine a team of specialists working together, where each one is an expert in a specific area, and a manager picks the best ones for the job. That’s MoE in a nutshell!In this blog, we’ll dive into MoE by walking through a simple PyTorch example. We’ll build an MoE model, train it on a fake dataset, and compare it to a traditional neural networkWhat is Mixture of Experts (MoE)?Consider the use of classifying animals in photos. A single neural network might struggle because it has to learn everything — cats, dogs, birds — all at once. MoE splits the work among experts, smaller neural networks trained on specific parts of the problem (e.g., one expert for cats, another for dogs).A gating network acts like a manager, deciding which experts to call on for each input. Only the best experts contribute to the final answer, making MoE both accurate and efficient.In our code, we’ll create an MoE with three experts, pick the top two for each input, and compare it to a single large neural network. The dataset is a simple one with 5,000 samples, each with four features and three classes (think of it like classifying three types of fruit). Let’s break down the code and see how it all comes together.1. Building the Expert Modelclass Expert(nn.Module): def __init__(self, input_dim, hidden_dim, output_dim): super(Expert, self).__init__() self.layer1 = nn.Linear(input_dim, hidden_dim) self.layer2 = nn.Linear(hidden_dim, output_dim) def forward(self, x): x = torch.relu(self.layer1(x)) return torch.softmax(self.layer2(x), dim=1)Each expert is like a mini neural network with a specific job. In our case, we have three experts, and each one is a simple feed-forward network with two layers:Layer 1: Takes the input (4 features, like measurements of a fruit) and transforms it into a hidden layer of 32 units. We use ReLU (a function that keeps positive values and zeros out negatives) to add some non-linearity, helping the network learn complex patterns.Layer 2: Converts the hidden layer into 3 outputs (one for each class). The softmax function turns these into probabilities, so we know how likely each class is.Think of each expert as a specialist who’s really good at spotting one type of animal. With three experts, we cover all three classes, and each one has 259 parameters (weights and biases), making them lightweight.Parameter Breakdown for One ExpertLayer 1: 4 (input) → 32 (hidden)Weights: 4 × 32 = 128Biases: 32 (one per hidden unit)Total: 128 + 32 = 160 parametersLayer 2: 32 (hidden) → 3 (output)Weights: 32 × 3 = 96Biases: 3 (one per output unit)Total: 96 + 3 = 99 parametersTotal per expert: 160 + 99 = 259 parameters2. Creating the Gating Networkclass Gating(nn.Module): def __init__(self, input_dim, num_experts): super(Gating, self).__init__() self.layer1 = nn.Linear(input_dim, 16) # Single hidden layer self.layer2 = nn.Linear(16, num_experts) def forward(self, x): x = torch.relu(self.layer1(x)) return torch.softmax(self.layer2(x), dim=1)Layer 1: Takes the 4 input features and maps them to a small hidden layer of 16 units, again using ReLU to learn patterns.Layer 2: Outputs a score for each of the 3 experts. The softmax function turns these scores into probabilities (e.g., [0.5, 0.3, 0.2]), showing how much each expert should contribute.This gating network is super simple, with only 131 parameters, but it’s powerful enough to pick the right experts.3. Putting It Together: The MoE Modelclass MoE(nn.Module): def __init__(self, experts, top_k=2): super(MoE, self).__init__() self.experts = nn.ModuleList(experts) self.gating = Gating(input_dim=experts[0].layer1.in_features, num_experts=len(experts)) self.top_k = top_k def forward(self, x): weights = self.gating(x) top_k_weights, top_k_indices = torch.topk(weights, k=self.top_k, dim=1) top_k_weights = top_k_weights / top_k_weights.sum(dim=1, keepdim=True) batch_size = x.size(0) output = torch.zeros(batch_size, self.experts[0].layer2.out_features, device=x.device) for i in range(batch_size): for k in range(self.top_k): expert_idx = top_k_indices[i, k] weight = top_k_weights[i, k] expert_output = self.experts[expert_idx](x[i:i+1]) output[i] += weight * expert_output.squeeze(0) return outputOnly 649 parameters (2*259 + 131 from the gating network) are used during inference, making MoE super efficient compared to a traditional network.4. The Single Neural Network (for Comparison)class SingleNN(nn.Module): def __init__(self, input_dim, hidden_dim1, hidden_dim2, output_dim): super(SingleNN, self).__init__() self.layer1 = nn.Linear(input_dim, hidden_dim1) self.layer2 = nn.Linear(hidden_dim1, hidden_dim2) self.layer3 = nn.Linear(hidden_dim2, output_dim) def forward(self, x): x = torch.relu(self.layer1(x)) x = torch.relu(self.layer2(x)) return torch.softmax(self.layer3(x), dim=1)The neural network has layers 4→350→190→3. It’s a traditional network that tries to learn everything at once, with 68,827 parameters, way more than the MoE’s active 649.5. Counting Parametersdef count_parameters(model): return sum(p.numel() for p in model.parameters() if p.requires_grad)6. Creating a training datasetnum_samples = 5000input_dim = 4output_dim = 3x_data = torch.randn(num_samples, input_dim)y_data = torch.cat([torch.zeros(num_samples // 3), torch.ones(num_samples // 3), torch.full((num_samples - 2 * (num_samples // 3),), 2)]).long()for i in range(num_samples): if y_data[i] == 0: x_data[i, 0] += 1 elif y_data[i] == 1: x_data[i, 1] -= 1 elif y_data[i] == 2: x_data[i, 0] -= 1shuffled_indices = torch.randperm(num_samples)x_data, y_data = x_data[shuffled_indices], y_data[shuffled_indices]train_size = int(0.8 * num_samples)x_train, y_train = x_data[:train_size], y_data[:train_size]x_test, y_test = x_data[train_size:], y_data[train_size:]7. Training the Expertshidden_dim = 32epochs = 100learning_rate = 0.001experts = [Expert(input_dim, hidden_dim, output_dim) for _ in range(3)]optimizers = [optim.Adam(expert.parameters(), lr=learning_rate) for expert in experts]for i, expert in enumerate(experts): optimizer = optimizers[i] mask = y_train == i x_train_subset, y_train_subset = x_train[mask], y_train[mask] if len(x_train_subset) == 0: continue for epoch in range(epochs): optimizer.zero_grad() outputs = expert(x_train_subset) loss = nn.CrossEntropyLoss()(outputs, y_train_subset) loss.backward() optimizer.step()8. Training the MoE Modelmoe_model = MoE(experts, top_k=2)optimizer_moe = optim.Adam(moe_model.parameters(), lr=learning_rate)for epoch in range(epochs): optimizer_moe.zero_grad() outputs_moe = moe_model(x_train) loss_moe = nn.CrossEntropyLoss()(outputs_moe, y_train) loss_moe.backward() optimizer_moe.step() if (epoch + 1) % 20 == 0: print(f"MoE Epoch [{epoch+1}/{epochs}], Loss: {loss_moe.item():.4f}")Output:MoE Epoch [20/100], Loss: 1.0884MoE Epoch [40/100], Loss: 1.0494MoE Epoch [60/100], Loss: 1.0133MoE Epoch [80/100], Loss: 0.9794MoE Epoch [100/100], Loss: 0.94839. Training the Single Neural Networksingle_nn = SingleNN(input_dim=4, hidden_dim1=350, hidden_dim2=190, output_dim=3)optimizer_nn = optim.Adam(single_nn.parameters(), lr=learning_rate)for epoch in range(epochs): optimizer_nn.zero_grad() outputs_nn = single_nn(x_train) loss_nn = nn.CrossEntropyLoss()(outputs_nn, y_train) loss_nn.backward() optimizer_nn.step() if (epoch + 1) % 20 == 0: print(f"Single NN Epoch [{epoch+1}/{epochs}], Loss: {loss_nn.item():.4f}")Output:Single NN Epoch [20/100], Loss: 0.8549 Single NN Epoch [40/100], Loss: 0.8487 Single NN Epoch [60/100], Loss: 0.8468 Single NN Epoch [80/100], Loss: 0.8453Single NN Epoch [100/100], Loss: 0.843810. Evaluating the Modelsdef evaluate(model, x, y): with torch.no_grad(): outputs = model(x) _, predicted = torch.max(outputs, 1) correct = (predicted == y).sum().item() return correct / len(y)accuracy_moe = evaluate(moe_model, x_test, y_test)accuracy_nn = evaluate(single_nn, x_test, y_test)11. Showing the Resultsprint("\nResults Comparison:")print(f"MoE Total Parameters: {moe_total_params}")print(f"MoE Active Parameters (Inference): {moe_active_params}")print(f"Single NN Parameters: {nn_params}")print(f"MoE Model Accuracy: {accuracy_moe:.4f}")print(f"Single NN Accuracy: {accuracy_nn:.4f}")Output:Results Comparison: MoE Total Parameters: 908 MoE Active Parameters (Inference): 649 Single NN Parameters: 69013 MoE Model Accuracy: 0.6540 Single NN Accuracy: 0.655012. SummaryThe Mixture of Experts (MoE) model used far fewer parameters than the single neural network, only 649 active ones during inference compared to 69,013 in the single model.Even with this huge drop in size, the MoE still achieved almost the same accuracy (0.6540 vs. 0.6550). This shows that the MoE model is much more efficient, using fewer resources while delivering similar performance.This story is published under the Generative AI publication. Connect with us on LinkedIn and follow Zeniteq to stay in the loop with the latest AI stories. Let’s shape the future of AI together!Understanding Mixture of Experts (MoE) was originally published in Generative AI on Medium, where people are continuing the conversation by highlighting and responding to this story.Source: Generative AI Pub — Published — Category: Image AI