LLMs and generative AI / LLM architecture and theory
Comparing Dense, Sparse, and Mixture of Experts for LLMs
Contents
- Introduction to Dense, Sparse, and Mixture of Experts
- Dense Neural Networks
- Advantages of Dense Networks
- Limitations of Dense Networks
- Sparse Neural Networks
- Advantages of Sparse Networks
- Limitations of Sparse Networks
- Mixture of Experts (MoE)
- Advantages of Mixture of Experts
- Limitations of Mixture of Experts
- Comparing Dense, Sparse, and MoE Architectures
- Real-Life Example: Language Translation
- Real-Life Example: Image Classification
- Choosing the Right Architecture
- Additional Resources
Introduction to Dense, Sparse, and Mixture of Experts
Dense, sparse, and mixture of experts are different architectures used in Large Language Models (LLMs). Each has its own strengths and use cases. This presentation will explore these architectures, their differences, and when to use each one.
import matplotlib.pyplot as plt
def plot_architecture(arch_type):
fig, ax = plt.subplots(figsize=(8, 6))
if arch_type == 'dense':
ax.add_patch(plt.Rectangle((0, 0), 1, 1, fill=False))
ax.text(0.5, 0.5, 'Dense', ha='center', va='center', fontsize=20)
elif arch_type == 'sparse':
for i in range(10):
for j in range(10):
if np.random.rand() > 0.7:
ax.add_patch(plt.Circle((i/10, j/10), 0.03))
ax.text(0.5, 1.05, 'Sparse', ha='center', va='center', fontsize=20)
else: # mixture of experts
for i in range(3):
ax.add_patch(plt.Rectangle((i/3, 0), 0.3, 0.3, fill=False))
ax.text((i+0.5)/3, 0.15, f'Expert {i+1}', ha='center', va='center')
ax.add_patch(plt.Rectangle((0, 0.5), 1, 0.3, fill=False))
ax.text(0.5, 0.65, 'Router', ha='center', va='center')
ax.axis('off')
plt.show()
plot_architecture('dense')
plot_architecture('sparse')
plot_architecture('mixture')
Dense Neural Networks
Dense neural networks, also known as fully connected networks, are the traditional architecture where each neuron is connected to every neuron in the adjacent layers. They are simple to implement and understand but can be computationally expensive for large models.
import torch.nn as nn
class DenseNetwork(nn.Module):
def __init__(self, input_size, hidden_size, output_size):
super(DenseNetwork, self).__init__()
self.layer1 = nn.Linear(input_size, hidden_size)
self.layer2 = nn.Linear(hidden_size, hidden_size)
self.layer3 = nn.Linear(hidden_size, output_size)
self.relu = nn.ReLU()
def forward(self, x):
x = self.relu(self.layer1(x))
x = self.relu(self.layer2(x))
x = self.layer3(x)
return x
# Example usage
model = DenseNetwork(input_size=100, hidden_size=50, output_size=10)
input_tensor = torch.randn(1, 100)
output = model(input_tensor)
print(f"Input shape: {input_tensor.shape}")
print(f"Output shape: {output.shape}")
Advantages of Dense Networks
Dense networks excel in tasks where all input features are equally important. They can capture complex relationships between features and are suitable for smaller datasets. Dense networks are often used in classification tasks and as building blocks in more complex architectures.
import matplotlib.pyplot as plt
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.neural_network import MLPClassifier
# Generate a synthetic dataset
X, y = make_classification(n_samples=1000, n_features=2, n_informative=2,
n_redundant=0, n_clusters_per_class=1, random_state=42)
# Split the data
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# Train a dense neural network
mlp = MLPClassifier(hidden_layer_sizes=(100, 50), max_iter=300, random_state=42)
mlp.fit(X_train, y_train)
# Plot decision boundary
xx, yy = np.meshgrid(np.linspace(X[:, 0].min()-1, X[:, 0].max()+1, 100),
np.linspace(X[:, 1].min()-1, X[:, 1].max()+1, 100))
Z = mlp.predict(np.c_[xx.ravel(), yy.ravel()]).reshape(xx.shape)
plt.figure(figsize=(10, 8))
plt.contourf(xx, yy, Z, alpha=0.8, cmap=plt.cm.RdYlBu)
plt.scatter(X[:, 0], X[:, 1], c=y, cmap=plt.cm.RdYlBu, edgecolor='black')
plt.title("Decision Boundary of Dense Neural Network")
plt.xlabel("Feature 1")
plt.ylabel("Feature 2")
plt.show()
print(f"Accuracy on test set: {mlp.score(X_test, y_test):.2f}")
Limitations of Dense Networks
While dense networks are powerful, they have limitations. As the network size increases, the number of parameters grows quadratically, leading to high computational costs and potential overfitting. This makes them less suitable for very large models like state-of-the-art LLMs.
import torch.nn as nn
import matplotlib.pyplot as plt
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
input_sizes = [100, 1000, 10000, 100000]
hidden_sizes = [50, 500, 5000, 50000]
param_counts = []
for input_size, hidden_size in zip(input_sizes, hidden_sizes):
model = nn.Sequential(
nn.Linear(input_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, 10)
)
param_counts.append(count_parameters(model))
plt.figure(figsize=(10, 6))
plt.plot(input_sizes, param_counts, marker='o')
plt.title("Parameter Count vs Input Size in Dense Networks")
plt.xlabel("Input Size")
plt.ylabel("Number of Parameters")
plt.xscale('log')
plt.yscale('log')
plt.grid(True)
plt.show()
for size, count in zip(input_sizes, param_counts):
print(f"Input size: {size}, Parameters: {count:,}")
Sparse Neural Networks
Sparse neural networks have fewer connections between neurons compared to dense networks. In these networks, many weights are zero, leading to a more efficient representation. Sparsity can be achieved through pruning, where less important connections are removed, or through sparse initialization.
import torch.nn as nn
class SparseLinear(nn.Module):
def __init__(self, in_features, out_features, sparsity=0.9):
super(SparseLinear, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.weight = nn.Parameter(torch.Tensor(out_features, in_features))
self.bias = nn.Parameter(torch.Tensor(out_features))
self.reset_parameters()
self.apply_sparsity(sparsity)
def reset_parameters(self):
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
nn.init.uniform_(self.bias, -bound, bound)
def apply_sparsity(self, sparsity):
mask = torch.rand(self.weight.shape) > sparsity
self.weight.data *= mask.float()
def forward(self, input):
return nn.functional.linear(input, self.weight, self.bias)
# Example usage
sparse_layer = SparseLinear(100, 50, sparsity=0.8)
input_tensor = torch.randn(1, 100)
output = sparse_layer(input_tensor)
print(f"Input shape: {input_tensor.shape}")
print(f"Output shape: {output.shape}")
print(f"Sparsity: {(sparse_layer.weight == 0).float().mean().item():.2f}")
Advantages of Sparse Networks
Sparse networks offer several advantages over dense networks. They require less memory and computational resources, making them more efficient for large-scale models. Sparsity can also act as a form of regularization, potentially improving generalization. In LLMs, sparsity can help capture task-specific information more effectively.
import torch.nn as nn
import time
import matplotlib.pyplot as plt
class SparseLinear(nn.Module):
def __init__(self, in_features, out_features, sparsity=0.9):
super(SparseLinear, self).__init__()
self.linear = nn.Linear(in_features, out_features)
mask = torch.rand(self.linear.weight.shape) > sparsity
self.linear.weight.data *= mask.float()
def forward(self, x):
return self.linear(x)
def compare_performance(input_size, hidden_size, output_size, batch_size, sparsity):
dense_model = nn.Sequential(nn.Linear(input_size, hidden_size), nn.ReLU(), nn.Linear(hidden_size, output_size))
sparse_model = nn.Sequential(SparseLinear(input_size, hidden_size, sparsity), nn.ReLU(), SparseLinear(hidden_size, output_size, sparsity))
input_data = torch.randn(batch_size, input_size)
start_time = time.time()
dense_output = dense_model(input_data)
dense_time = time.time() - start_time
start_time = time.time()
sparse_output = sparse_model(input_data)
sparse_time = time.time() - start_time
dense_params = sum(p.numel() for p in dense_model.parameters())
sparse_params = sum((p != 0).sum().item() for p in sparse_model.parameters())
return dense_time, sparse_time, dense_params, sparse_params
input_sizes = [1000, 5000, 10000, 50000]
results = [compare_performance(size, size//2, 100, 64, 0.9) for size in input_samples]
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(input_sizes, [r[0] for r in results], label='Dense')
plt.plot(input_sizes, [r[1] for r in results], label='Sparse')
plt.title('Inference Time')
plt.xlabel('Input Size')
plt.ylabel('Time (s)')
plt.legend()
plt.subplot(1, 2, 2)
plt.plot(input_sizes, [r[2] for r in results], label='Dense')
plt.plot(input_sizes, [r[3] for r in results], label='Sparse')
plt.title('Parameter Count')
plt.xlabel('Input Size')
plt.ylabel('Number of Parameters')
plt.legend()
plt.tight_layout()
plt.show()
Limitations of Sparse Networks
Despite their advantages, sparse networks have some drawbacks. Training sparse networks from scratch can be challenging, and they may struggle to capture complex patterns that dense networks can learn. Finding the right sparsity level and structure is crucial for performance.
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
class SparseLinear(nn.Module):
def __init__(self, in_features, out_features, sparsity=0.9):
super(SparseLinear, self).__init__()
self.linear = nn.Linear(in_features, out_features)
self.sparsity = sparsity
self.apply_sparsity()
def apply_sparsity(self):
mask = torch.rand(self.linear.weight.shape) > self.sparsity
self.linear.weight.data *= mask.float()
def forward(self, x):
return self.linear(x)
def train_model(model, X, y, epochs=100):
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
losses = []
for epoch in range(epochs):
optimizer.zero_grad()
outputs = model(X)
loss = criterion(outputs, y)
loss.backward()
optimizer.step()
if isinstance(model[0], SparseLinear):
model[0].apply_sparsity()
losses.append(loss.item())
return losses
# Generate synthetic data
X = torch.randn(100, 10)
y = torch.randn(100, 1)
dense_model = nn.Sequential(nn.Linear(10, 1))
sparse_model = nn.Sequential(SparseLinear(10, 1, sparsity=0.5))
dense_losses = train_model(dense_model, X, y)
sparse_losses = train_model(sparse_model, X, y)
plt.figure(figsize=(10, 6))
plt.plot(dense_losses, label='Dense')
plt.plot(sparse_losses, label='Sparse')
plt.title('Training Loss Comparison')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.show()
print(f"Final dense loss: {dense_losses[-1]:.4f}")
print(f"Final sparse loss: {sparse_losses[-1]:.4f}")
Mixture of Experts (MoE)
Mixture of Experts (MoE) is an architecture that combines multiple "expert" networks with a router network. The router decides which expert(s) to use for each input. This allows the model to specialize different experts for different types of inputs or tasks, potentially improving overall performance.
import torch.nn as nn
import torch.nn.functional as F
class MixtureOfExperts(nn.Module):
def __init__(self, input_size, output_size, num_experts):
super(MixtureOfExperts, self).__init__()
self.router = nn.Linear(input_size, num_experts)
self.experts = nn.ModuleList([nn.Linear(input_size, output_size) for _ in range(num_experts)])
def forward(self, x):
router_probs = F.softmax(self.router(x), dim=-1)
expert_outputs = torch.stack([expert(x) for expert in self.experts])
output = torch.sum(router_probs.unsqueeze(-1) * expert_outputs, dim=0)
return output, router_probs
# Example usage
moe = MixtureOfExperts(input_size=10, output_size=5, num_experts=3)
input_tensor = torch.randn(2, 10)
output, router_probs = moe(input_tensor)
print(f"Input shape: {input_tensor.shape}")
print(f"Output shape: {output.shape}")
print(f"Router probabilities shape: {router_probs.shape}")
Advantages of Mixture of Experts
MoE architectures offer several benefits, especially for large language models. They allow for conditional computation, where only a subset of the model is active for each input. This can lead to more efficient use of model capacity and improved performance on diverse tasks.
import matplotlib.pyplot as plt
def simulate_moe_performance(num_tasks, num_experts):
task_performance = np.random.rand(num_tasks, num_experts)
moe_performance = np.max(task_performance, axis=1)
average_model_performance = np.mean(task_performance, axis=1)
return moe_performance, average_model_performance
num_tasks = 100
num_experts = 5
moe_perf, avg_perf = simulate_moe_performance(num_tasks, num_experts)
plt.figure(figsize=(10, 6))
plt.scatter(range(num_tasks), moe_perf, label='MoE', alpha=0.7)
plt.scatter(range(num_tasks), avg_perf, label='Average Model', alpha=0.7)
plt.title('MoE vs Average Model Performance Across Tasks')
plt.xlabel('Task')
plt.ylabel('Performance')
plt.legend()
plt.show()
print(f"Average MoE performance: {np.mean(moe_perf):.3f}")
print(f"Average single model performance: {np.mean(avg_perf):.3f}")
Limitations of Mixture of Experts
While MoE models can be powerful, they also come with challenges. The routing mechanism adds complexity to the model, which can make training more difficult. There's also a risk of load balancing issues, where some experts may be underutilized.
import matplotlib.pyplot as plt
def simulate_expert_utilization(num_samples, num_experts, skewness=2):
# Simulate skewed expert utilization
utilization = np.random.pareto(skewness, num_experts)
utilization = utilization / np.sum(utilization)
expert_counts = np.random.choice(num_experts, num_samples, p=utilization)
return np.bincount(expert_counts, minlength=num_experts) / num_samples
num_samples = 10000
num_experts = 8
utilization = simulate_expert_utilization(num_samples, num_experts)
plt.figure(figsize=(10, 6))
plt.bar(range(num_experts), utilization)
plt.title('Expert Utilization in MoE')
plt.xlabel('Expert')
plt.ylabel('Utilization')
plt.show()
gini_coefficient = 1 - 2 * np.sum((np.cumsum(np.sort(utilization)) / np.sum(utilization)) * (1 / num_experts))
print(f"Gini coefficient (inequality measure): {gini_coefficient:.3f}")
Comparing Dense, Sparse, and MoE Architectures
Each architecture has its strengths and weaknesses. Dense networks are simple and effective for smaller models. Sparse networks can be more efficient and scalable. MoE can offer a balance between model capacity and computational efficiency.
import matplotlib.pyplot as plt
def compare_architectures(model_sizes):
dense_params = model_sizes**2
sparse_params = model_sizes * np.log(model_sizes)
moe_params = model_sizes * np.sqrt(model_sizes)
return dense_params, sparse_params, moe_params
model_sizes = np.logspace(2, 5, num=20)
dense, sparse, moe = compare_architectures(model_sizes)
plt.figure(figsize=(12, 6))
plt.loglog(model_sizes, dense, label='Dense')
plt.loglog(model_sizes, sparse, label='Sparse')
plt.loglog(model_sizes, moe, label='MoE')
plt.title('Parameter Count vs Model Size')
plt.xlabel('Model Size')
plt.ylabel('Number of Parameters')
plt.legend()
plt.grid(True)
plt.show()
print("Relative parameter counts for largest model size:")
print(f"Dense: {dense[-1]:.2e}")
print(f"Sparse: {sparse[-1]:.2e}")
print(f"MoE: {moe[-1]:.2e}")
Real-Life Example: Language Translation
In language translation, different architectures can be employed based on the specific requirements. Dense networks might be suitable for small-scale, general-purpose translation. Sparse networks could be used for efficient, large-scale models. MoE could excel in multi-lingual translation by specializing experts for different language pairs.
import matplotlib.pyplot as plt
def simulate_translation_quality(num_language_pairs, architecture):
if architecture == 'dense':
return np.random.normal(0.7, 0.1, num_language_pairs)
elif architecture == 'sparse':
return np.random.normal(0.75, 0.15, num_language_pairs)
elif architecture == 'moe':
return np.random.normal(0.8, 0.2, num_language_pairs)
num_pairs = 50
dense_quality = simulate_translation_quality(num_pairs, 'dense')
sparse_quality = simulate_translation_quality(num_pairs, 'sparse')
moe_quality = simulate_translation_quality(num_pairs, 'moe')
plt.figure(figsize=(12, 6))
plt.violinplot([dense_quality, sparse_quality, moe_quality])
plt.xticks([1, 2, 3], ['Dense', 'Sparse', 'MoE'])
plt.title('Translation Quality Distribution')
plt.ylabel('BLEU Score')
plt.show()
print(f"Average BLEU scores:")
print(f"Dense: {np.mean(dense_quality):.3f}")
print(f"Sparse: {np.mean(sparse_quality):.3f}")
print(f"MoE: {np.mean(moe_quality):.3f}")
Real-Life Example: Image Classification
In image classification tasks, the choice of architecture can significantly impact performance and efficiency. Dense networks are often used for smaller datasets or when interpretability is important. Sparse networks can be effective for large-scale image classification, reducing computational requirements. MoE architectures could potentially specialize in different types of images or objects.
import matplotlib.pyplot as plt
def simulate_classification_performance(num_classes, architecture):
if architecture == 'dense':
return np.random.normal(0.8, 0.1, num_classes)
elif architecture == 'sparse':
return np.random.normal(0.85, 0.08, num_classes)
elif architecture == 'moe':
return np.clip(np.random.normal(0.9, 0.15, num_classes), 0, 1)
num_classes = 100
dense_perf = simulate_classification_performance(num_classes, 'dense')
sparse_perf = simulate_classification_performance(num_classes, 'sparse')
moe_perf = simulate_classification_performance(num_classes, 'moe')
plt.figure(figsize=(12, 6))
plt.scatter(range(num_classes), dense_perf, alpha=0.6, label='Dense')
plt.scatter(range(num_classes), sparse_perf, alpha=0.6, label='Sparse')
plt.scatter(range(num_classes), moe_perf, alpha=0.6, label='MoE')
plt.title('Classification Performance Across Classes')
plt.xlabel('Class')
plt.ylabel('Accuracy')
plt.legend()
plt.show()
print(f"Average accuracies:")
print(f"Dense: {np.mean(dense_perf):.3f}")
print(f"Sparse: {np.mean(sparse_perf):.3f}")
print(f"MoE: {np.mean(moe_perf):.3f}")
Choosing the Right Architecture
Selecting the optimal architecture depends on various factors including dataset size, computational resources, and specific task requirements. Dense networks are suitable for smaller models and when all input features are equally important. Sparse networks excel in large-scale applications with limited resources. MoE architectures can be beneficial for multi-task learning or when dealing with diverse input distributions.
import matplotlib.pyplot as plt
def architecture_score(data_size, compute_resources, task_diversity):
dense_score = 10 - (0.5 * data_size + 0.3 * compute_resources + 0.2 * task_diversity)
sparse_score = 7 + (0.4 * data_size + 0.4 * compute_resources + 0.2 * task_diversity)
moe_score = 5 + (0.3 * data_size + 0.3 * compute_resources + 0.4 * task_diversity)
return max(dense_score, 0), max(sparse_score, 0), max(moe_score, 0)
data_sizes = np.linspace(0, 10, 50)
compute_resources = 5 # Fixed value for simplicity
task_diversity = 5 # Fixed value for simplicity
dense_scores = []
sparse_scores = []
moe_scores = []
for size in data_sizes:
d, s, m = architecture_score(size, compute_resources, task_diversity)
dense_scores.append(d)
sparse_scores.append(s)
moe_scores.append(m)
plt.figure(figsize=(12, 6))
plt.plot(data_sizes, dense_scores, label='Dense')
plt.plot(data_sizes, sparse_scores, label='Sparse')
plt.plot(data_sizes, moe_scores, label='MoE')
plt.title('Architecture Suitability vs Data Size')
plt.xlabel('Data Size')
plt.ylabel('Suitability Score')
plt.legend()
plt.show()
Additional Resources
For more in-depth information on these architectures and their applications in LLMs, consider exploring the following resources:
- "Sparse is Enough in Scaling Transformers" (arXiv:2111.12763)
- "GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding" (arXiv:2006.16668)
- "Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity" (arXiv:2101.03961)
These papers provide detailed insights into the implementation and performance of sparse and mixture of experts models in large-scale language tasks.