""" Mixture-of-Experts (MoE): Explainer + Interactive Routing Simulator -------------------------------------------------------------------- A Hugging Face Space (Gradio) that: 1. Explains how sparse Mixture-of-Experts layers work (gating, top-k routing, load balancing, capacity, sparse compute savings). 2. Lets the user run a live simulation: tokens are embedded, a gating network scores them against experts, top-k experts are chosen per token (optionally with noisy/capacity-limited routing like Switch Transformer / Mixtral / GShard), and the results are visualized (routing map, expert load, gate-weight histogram, FLOPs comparison). No trained model is required -- everything is a lightweight numpy simulation, so the Space is instant to load and free to run on CPU. """ import numpy as np import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import gradio as gr # -------------------------------------------------------------------- # Optional ZeroGPU compatibility shim # -------------------------------------------------------------------- # This app is pure CPU (numpy/matplotlib) and never needs a GPU. If the # Space's hardware is set to "CPU basic" (recommended), the `spaces` # package won't even be installed, and everything below is skipped. # # If someone instead deploys this on a "ZeroGPU" hardware tier, that # runtime requires at least one function decorated with @spaces.GPU to # be present at startup, or it throws a runtime error. This shim # satisfies that requirement with a trivial no-op, without ever # actually consuming GPU quota (the simulation itself never calls it). try: import spaces @spaces.GPU def _zerogpu_startup_check(): """No-op so ZeroGPU Spaces detect a valid @spaces.GPU function at startup. Never called by the app itself -- this Space does not use a GPU for anything.""" return True except ImportError: # `spaces` isn't installed -> we're on CPU basic (or running # locally). Nothing to do. pass # -------------------------------------------------------------------- # Core MoE routing simulation # -------------------------------------------------------------------- RNG_SEED = 0 def simulate_moe( num_tokens: int, embed_dim: int, num_experts: int, top_k: int, noise_std: float, capacity_factor: float, seed: int, ): """Simulate one forward pass of a sparse MoE gating layer. Returns a dict with routing assignments, gate weights, dropped tokens (due to capacity), and everything needed to plot. """ rng = np.random.default_rng(seed) # 1. Fake token embeddings (in a real model these come from the # previous transformer block's hidden states). # We give tokens soft "cluster" structure so routing isn't pure # noise -- similar to how semantically similar tokens tend to # prefer the same experts in trained MoEs. num_clusters = min(num_experts, max(2, num_experts // 2)) cluster_centers = rng.normal(0, 3.0, size=(num_clusters, embed_dim)) cluster_ids = rng.integers(0, num_clusters, size=num_tokens) tokens = cluster_centers[cluster_ids] + rng.normal(0, 1.0, size=(num_tokens, embed_dim)) # 2. Gating network: a single linear layer (token_dim -> num_experts) W_gate = rng.normal(0, 1.0 / np.sqrt(embed_dim), size=(embed_dim, num_experts)) logits = tokens @ W_gate # 3. Noisy gating (as in Shazeer et al. 2017 / Switch Transformer) # helps exploration & load balancing during training. if noise_std > 0: logits = logits + rng.normal(0, noise_std, size=logits.shape) gate_probs = _softmax(logits, axis=-1) # 4. Top-k expert selection per token topk_idx = np.argsort(-gate_probs, axis=-1)[:, :top_k] topk_val = np.take_along_axis(gate_probs, topk_idx, axis=-1) # re-normalize the chosen top-k weights so they sum to 1 (standard practice) topk_val = topk_val / topk_val.sum(axis=-1, keepdims=True) # 5. Expert capacity (Switch/GShard style): each expert can only # process `capacity` tokens per batch; overflow tokens are # dropped (their residual/skip connection carries them instead). capacity = max(1, int(np.ceil(capacity_factor * num_tokens * top_k / num_experts))) expert_fill = np.zeros(num_experts, dtype=int) kept_mask = np.ones_like(topk_idx, dtype=bool) # process tokens in order, first-come-first-served per expert for t in range(num_tokens): for k in range(top_k): e = topk_idx[t, k] if expert_fill[e] < capacity: expert_fill[e] += 1 else: kept_mask[t, k] = False # dropped: over capacity # 6. Load per expert (tokens actually routed, i.e. not dropped) load = np.zeros(num_experts, dtype=int) for t in range(num_tokens): for k in range(top_k): if kept_mask[t, k]: load[topk_idx[t, k]] += 1 dropped_tokens = int((~kept_mask).sum()) total_routes = num_tokens * top_k # 7. Load-balancing auxiliary loss (Switch Transformer formulation) # L_aux = num_experts * sum_e( f_e * P_e ) # f_e = fraction of tokens dispatched to expert e (top-1 proxy) top1 = topk_idx[:, 0] f = np.bincount(top1, minlength=num_experts) / num_tokens P = gate_probs.mean(axis=0) aux_loss = num_experts * np.sum(f * P) return dict( tokens=tokens, cluster_ids=cluster_ids, gate_probs=gate_probs, topk_idx=topk_idx, topk_val=topk_val, kept_mask=kept_mask, load=load, capacity=capacity, dropped_tokens=dropped_tokens, total_routes=total_routes, aux_loss=aux_loss, num_experts=num_experts, top_k=top_k, num_tokens=num_tokens, ) def _softmax(x, axis=-1): x = x - np.max(x, axis=axis, keepdims=True) e = np.exp(x) return e / np.sum(e, axis=axis, keepdims=True) # -------------------------------------------------------------------- # FLOPs / compute-savings estimate (dense vs sparse MoE) # -------------------------------------------------------------------- def compute_flops_comparison(num_experts, top_k, ffn_mult=4, embed_dim=1024): """Rough FFN FLOPs-per-token estimate: dense model uses ALL experts' worth of parameters for every token; MoE only activates top_k. Assumes each 'expert' is sized like one dense FFN block.""" per_expert_flops = 2 * embed_dim * (embed_dim * ffn_mult) * 2 # up+down proj, x2 for mul-add dense_equiv_flops = per_expert_flops * num_experts # a dense FFN with same total capacity moe_flops = per_expert_flops * top_k speedup = dense_equiv_flops / moe_flops return dense_equiv_flops, moe_flops, speedup # -------------------------------------------------------------------- # Plotting # -------------------------------------------------------------------- EXPERT_CMAP = plt.get_cmap("tab20") def plot_routing_map(sim): tokens, topk_idx, kept_mask = sim["tokens"], sim["topk_idx"], sim["kept_mask"] # 2D projection via first 2 dims (embed space already low-dim for viz clarity) if tokens.shape[1] > 2: # simple PCA-free projection: just take first two coordinates xy = tokens[:, :2] else: xy = tokens fig, ax = plt.subplots(figsize=(6, 5)) top1 = topk_idx[:, 0] top1_kept = kept_mask[:, 0] for e in range(sim["num_experts"]): mask = (top1 == e) & top1_kept ax.scatter(xy[mask, 0], xy[mask, 1], s=35, color=EXPERT_CMAP(e % 20), label=f"Expert {e}", alpha=0.85, edgecolors="white", linewidths=0.3) dropped = ~top1_kept if dropped.any(): ax.scatter(xy[dropped, 0], xy[dropped, 1], s=45, facecolors="none", edgecolors="red", linewidths=1.4, label="Dropped (over capacity)") ax.set_title("Token routing map (top-1 expert, colored)", fontsize=11) ax.set_xlabel("embedding dim 1") ax.set_ylabel("embedding dim 2") ax.legend(loc="upper right", fontsize=7, ncol=2, framealpha=0.9) fig.tight_layout() return fig def plot_expert_load(sim): load = sim["load"] capacity = sim["capacity"] fig, ax = plt.subplots(figsize=(6, 4)) colors = [EXPERT_CMAP(e % 20) for e in range(sim["num_experts"])] bars = ax.bar(range(sim["num_experts"]), load, color=colors, edgecolor="black", linewidth=0.4) ax.axhline(capacity, color="red", linestyle="--", linewidth=1.3, label=f"Capacity = {capacity}") ax.set_xlabel("Expert index") ax.set_ylabel("Tokens routed") ax.set_title("Expert load balance", fontsize=11) ax.set_xticks(range(sim["num_experts"])) ax.legend(fontsize=8) fig.tight_layout() return fig def plot_gate_weights(sim): vals = sim["topk_val"].flatten() fig, ax = plt.subplots(figsize=(6, 4)) ax.hist(vals, bins=25, color="#6C5CE7", edgecolor="white", alpha=0.9) ax.set_xlabel("Normalized gate weight (top-k, per token)") ax.set_ylabel("Count") ax.set_title("Gate-weight distribution", fontsize=11) fig.tight_layout() return fig # -------------------------------------------------------------------- # Gradio callback # -------------------------------------------------------------------- def run_simulation(num_tokens, embed_dim, num_experts, top_k, noise_std, capacity_factor, ffn_mult, seed): top_k = min(top_k, num_experts) sim = simulate_moe( num_tokens=int(num_tokens), embed_dim=int(embed_dim), num_experts=int(num_experts), top_k=int(top_k), noise_std=float(noise_std), capacity_factor=float(capacity_factor), seed=int(seed), ) dense_flops, moe_flops, speedup = compute_flops_comparison( int(num_experts), int(top_k), ffn_mult=int(ffn_mult), embed_dim=int(embed_dim) ) drop_pct = 100.0 * sim["dropped_tokens"] / sim["total_routes"] load = sim["load"] load_std = float(np.std(load)) load_mean = float(np.mean(load)) cv = load_std / load_mean if load_mean > 0 else 0.0 summary = f""" ### Routing summary | Metric | Value | |---|---| | Tokens | {sim['num_tokens']} | | Experts | {sim['num_experts']} | | Top-k | {sim['top_k']} | | Expert capacity (per expert) | {sim['capacity']} | | Dropped routes (over capacity) | {sim['dropped_tokens']} / {sim['total_routes']} ({drop_pct:.1f}%) | | Load balance (coefficient of variation, lower=better) | {cv:.3f} | | Load-balancing aux loss (Switch-style) | {sim['aux_loss']:.4f} | ### Compute comparison (FFN block, per token) | Model | FFN FLOPs/token | Relative | |---|---|---| | Dense (uses all {num_experts} experts) | {dense_flops:,.0f} | 1.0x | | Sparse MoE (top-{top_k} of {num_experts}) | {moe_flops:,.0f} | {1/speedup:.2f}x | **→ The sparse MoE layer does ~{speedup:.1f}x less FFN compute per token than an equivalently-sized dense layer**, while keeping (in a trained model) comparable total parameter capacity — this is MoE's core value proposition. """ return ( plot_routing_map(sim), plot_expert_load(sim), plot_gate_weights(sim), summary, ) # -------------------------------------------------------------------- # Explanatory markdown # -------------------------------------------------------------------- INTRO_MD = """ # 🧩 Mixture of Experts (MoE) — Explainer & Interactive Simulator **Mixture of Experts** is a neural network architecture where, instead of every input passing through one big dense feed-forward block, a layer is split into many smaller **expert** sub-networks, and a lightweight **gating (router) network** decides — *per token* — which one or few experts should process it. This is how models like **Mixtral 8x7B**, **DeepSeek-MoE**, **Switch Transformer**, and **GShard** scale to huge parameter counts while keeping inference compute roughly constant: a model might have hundreds of billions of *total* parameters, but only activate a few billion of them for any given token. ## How it works 1. **Experts.** Each MoE layer replaces one dense FFN with `N` independent FFN "experts" (`E_0 … E_{N-1}`), each with its own weights. 2. **Gating network.** A small linear layer scores each token against every expert: `logits = x @ W_gate`, then `softmax` turns scores into a probability distribution over experts. 3. **Top-k routing.** Only the top-`k` highest-scoring experts are actually run for that token (commonly k=1 for Switch Transformer, k=2 for Mixtral/ GShard). The token's output is the weighted sum of the chosen experts' outputs, weighted by their (renormalized) gate probabilities. 4. **Capacity limits.** To keep training/inference batched and efficient, each expert is only allowed to process a fixed number of tokens per batch (`capacity`). Tokens that overflow a full expert are **dropped** for that layer (they skip via the residual connection instead). 5. **Load balancing.** If left unconstrained, gating tends to collapse onto a few "favorite" experts. An **auxiliary load-balancing loss** penalizes uneven routing so all experts stay utilized and trained. ## Why it matters - **Sparse activation** → you get the representational capacity of a huge dense model, but the *compute cost per token* of a much smaller one. - **Specialization** → different experts can implicitly specialize (e.g. by language, syntax pattern, or topic), though in practice this specialization is soft and not perfectly interpretable. - **The catch** → routing imbalance, dropped tokens, communication overhead (in distributed training experts often live on different devices), and training instability are the main engineering challenges. --- ## 🎛️ Try it yourself below Tune the simulation parameters and watch: which expert each token gets routed to, how balanced the load is across experts, the gate-weight distribution, and the theoretical FLOPs savings vs. an equivalent dense layer. """ FOOTER_MD = """ --- ### Notes on this simulation - Token embeddings are **synthetic** (sampled from a few Gaussian clusters) purely to make routing patterns visually interpretable — no real language model is involved. - The gating network is a single random linear layer, matching the structure (not the trained weights) of a real MoE router. - Capacity-based dropping and the auxiliary load-balancing loss follow the formulation in the **Switch Transformer** paper (Fedus et al., 2021) and **GShard** (Lepikhin et al., 2020). - This is meant as an intuition-building tool, not a benchmark of any real model's routing behavior. """ # -------------------------------------------------------------------- # Gradio UI # -------------------------------------------------------------------- with gr.Blocks(title="Mixture of Experts — Explainer & Simulator") as demo: gr.Markdown(INTRO_MD) with gr.Row(): with gr.Column(scale=1): num_tokens = gr.Slider(8, 512, value=128, step=8, label="Number of tokens (batch)") embed_dim = gr.Slider(2, 64, value=16, step=1, label="Token embedding dim") num_experts = gr.Slider(2, 16, value=8, step=1, label="Number of experts") top_k = gr.Slider(1, 8, value=2, step=1, label="Top-k (experts activated per token)") noise_std = gr.Slider(0.0, 3.0, value=0.5, step=0.1, label="Gating noise (std)") capacity_factor = gr.Slider(0.5, 3.0, value=1.25, step=0.05, label="Capacity factor (>1 = slack, <1 = aggressive dropping)") ffn_mult = gr.Slider(1, 8, value=4, step=1, label="FFN expansion ratio (per expert)") seed = gr.Slider(0, 9999, value=0, step=1, label="Random seed") run_btn = gr.Button("▶ Run simulation", variant="primary") with gr.Column(scale=2): with gr.Row(): routing_plot = gr.Plot(label="Routing map") load_plot = gr.Plot(label="Expert load") gate_plot = gr.Plot(label="Gate weight distribution") summary_md = gr.Markdown() gr.Markdown(FOOTER_MD) inputs = [num_tokens, embed_dim, num_experts, top_k, noise_std, capacity_factor, ffn_mult, seed] outputs = [routing_plot, load_plot, gate_plot, summary_md] run_btn.click(fn=run_simulation, inputs=inputs, outputs=outputs) demo.load(fn=run_simulation, inputs=inputs, outputs=outputs) if __name__ == "__main__": try: demo.launch(theme=gr.themes.Soft()) except TypeError: # Gradio <6: theme belongs on Blocks(), not launch() -- fall back gracefully demo.launch()