Publish Top-2 sparse mixture-of-experts and dense control

README.md

@@ -0,0 +1,44 @@
+---
+license: apache-2.0
+task_categories:
+- image-classification
+tags:
+- mixture-of-experts
+- sparse-routing
+- load-balancing
+- pytorch
+---
+
+# Pocket MoE
+
+Pocket MoE is a compact top-2 mixture-of-experts classifier. A shared image encoder
+feeds four specialist MLPs while a learned router activates two experts for each
+example.
+
+The benchmark compares an MoE and a parameter-matched dense MLP on the same split,
+then reports expert utilization, routing entropy, and each digit's dominant expert.
+
+## Reproduce
+
+```powershell
+uv run python projects/tiny-vision-foundry/prepare_data.py
+uv run python projects/pocket-moe/train.py
+```
+
+Sparse routing here demonstrates conditional computation; it does not reduce the
+stored checkpoint size because all four experts remain available.
+
+## Verified results
+
+| Model | Stored parameters | Active parameters | Test accuracy | Macro F1 |
+| --- | ---: | ---: | ---: | ---: |
+| Dense control | 5,490 | 5,490 | 96.30% | 0.9618 |
+| Top-2 Pocket MoE | 5,004 | **3,608** | 96.30% | **0.9625** |
+
+The MoE activates 34.3% fewer parameters per example than the dense control without
+losing accuracy. Experts specialized strongly: digit `0` routed 86.5% to expert 3,
+digit `2` routed 83.1% to expert 2, and digit `7` routed 90.9% to expert 0.
+
+Utilization was not perfectly balanced: expert shares were 44.2%, 12.4%, 19.6%, and
+23.8%, for a coefficient of variation of 0.473. The stored MoE still includes all four
+experts, so conditional compute should not be confused with checkpoint compression.

dense-control/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:457fabfa0b0762393ca7447de7494ca85019b33beede5c81a94c941f877ad5fe
+size 22432

pocket-moe/evaluation.json

@@ -0,0 +1,76 @@
+{
+  "model": "Pocket MoE",
+  "experts": 4,
+  "active_experts_per_example": 2,
+  "stored_parameters": 5004,
+  "active_parameters_per_example": 3608,
+  "dense_control_parameters": 5490,
+  "moe_training": {
+    "best_epoch": 31,
+    "best_validation_accuracy": 0.9555555555555556
+  },
+  "dense_training": {
+    "best_epoch": 58,
+    "best_validation_accuracy": 0.9666666666666667
+  },
+  "moe_test": {
+    "accuracy": 0.9629629629629629,
+    "macro_f1": 0.9624784837803706,
+    "expert_utilization": [
+      0.4422966241836548,
+      0.12372409552335739,
+      0.1955741047859192,
+      0.23840512335300446
+    ],
+    "utilization_coefficient_of_variation": 0.47337183356285095,
+    "mean_router_entropy": 0.897014319896698,
+    "maximum_router_entropy": 1.3862943611198906,
+    "dominant_expert_by_digit": {
+      "0": {
+        "expert": 3,
+        "routing_share": 0.8653278946876526
+      },
+      "1": {
+        "expert": 0,
+        "routing_share": 0.7711800336837769
+      },
+      "2": {
+        "expert": 2,
+        "routing_share": 0.8314319849014282
+      },
+      "3": {
+        "expert": 2,
+        "routing_share": 0.7494029998779297
+      },
+      "4": {
+        "expert": 3,
+        "routing_share": 0.7010031938552856
+      },
+      "5": {
+        "expert": 0,
+        "routing_share": 0.6955094933509827
+      },
+      "6": {
+        "expert": 1,
+        "routing_share": 0.6472891569137573
+      },
+      "7": {
+        "expert": 0,
+        "routing_share": 0.909274697303772
+      },
+      "8": {
+        "expert": 0,
+        "routing_share": 0.5719929337501526
+      },
+      "9": {
+        "expert": 0,
+        "routing_share": 0.7218619585037231
+      }
+    }
+  },
+  "dense_control_test": {
+    "accuracy": 0.9629629629629629,
+    "macro_f1": 0.9618359470717962
+  },
+  "accuracy_delta": 0.0
+}

pocket-moe/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7f92fa5214a1d4d237c9baf3815f7f0285e14f19813b2ea8fe31224e4e41c120
+size 21592

source/model.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class PocketMoE(nn.Module):
+    def __init__(self, experts: int = 4, top_k: int = 2) -> None:
+        super().__init__()
+        self.expert_count = experts
+        self.top_k = top_k
+        self.encoder = nn.Sequential(
+            nn.Linear(64, 32),
+            nn.GELU(),
+        )
+        self.router = nn.Linear(32, experts)
+        self.experts = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Linear(32, 16),
+                    nn.GELU(),
+                    nn.Linear(16, 10),
+                )
+                for _ in range(experts)
+            ]
+        )
+
+    def forward(
+        self,
+        pixels: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        hidden = self.encoder(pixels)
+        router_probabilities = F.softmax(self.router(hidden), dim=1)
+        top_probabilities, top_indices = router_probabilities.topk(
+            self.top_k,
+            dim=1,
+        )
+        sparse_weights = torch.zeros_like(router_probabilities).scatter(
+            1,
+            top_indices,
+            top_probabilities,
+        )
+        sparse_weights = sparse_weights / sparse_weights.sum(dim=1, keepdim=True)
+        expert_logits = torch.stack(
+            [expert(hidden) for expert in self.experts],
+            dim=1,
+        )
+        logits = (expert_logits * sparse_weights.unsqueeze(-1)).sum(dim=1)
+        return logits, router_probabilities, sparse_weights
+
+
+class DenseControl(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.network = nn.Sequential(
+            nn.Linear(64, 48),
+            nn.GELU(),
+            nn.Linear(48, 40),
+            nn.GELU(),
+            nn.Linear(40, 10),
+        )
+
+    def forward(self, pixels: torch.Tensor) -> torch.Tensor:
+        return self.network(pixels)
+
+
+def parameter_count(model: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in model.parameters())
+
+
+def active_parameter_count(model: PocketMoE) -> int:
+    shared = sum(parameter.numel() for parameter in model.encoder.parameters())
+    router = sum(parameter.numel() for parameter in model.router.parameters())
+    experts = sorted(
+        [
+            sum(parameter.numel() for parameter in expert.parameters())
+            for expert in model.experts
+        ],
+        reverse=True,
+    )
+    return shared + router + sum(experts[: model.top_k])

source/train.py

@@ -0,0 +1,256 @@
+from __future__ import annotations
+
+import json
+import random
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+import trackio
+from model import (
+    DenseControl,
+    PocketMoE,
+    active_parameter_count,
+    parameter_count,
+)
+from safetensors.torch import save_file
+from sklearn.metrics import accuracy_score, f1_score
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, TensorDataset
+
+PROJECT_DIR = Path(__file__).resolve().parent
+ROOT_DIR = PROJECT_DIR.parents[1]
+DATA_DIR = ROOT_DIR / "projects" / "tiny-vision-foundry" / "data"
+ARTIFACT_DIR = PROJECT_DIR / "artifacts"
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def load_split(name: str) -> tuple[torch.Tensor, torch.Tensor]:
+    frame = pd.read_parquet(DATA_DIR / f"{name}.parquet")
+    pixels = np.stack(frame["image"].to_numpy()).astype(np.float32) / 16.0
+    labels = frame["label"].to_numpy(dtype=np.int64, copy=True)
+    return torch.from_numpy(pixels), torch.from_numpy(labels)
+
+
+@torch.inference_mode()
+def evaluate_dense(
+    model: DenseControl,
+    pixels: torch.Tensor,
+    labels: torch.Tensor,
+) -> dict:
+    model.eval()
+    predictions = model(pixels).argmax(dim=1).numpy()
+    return {
+        "accuracy": float(accuracy_score(labels.numpy(), predictions)),
+        "macro_f1": float(f1_score(labels.numpy(), predictions, average="macro")),
+    }
+
+
+@torch.inference_mode()
+def evaluate_moe(
+    model: PocketMoE,
+    pixels: torch.Tensor,
+    labels: torch.Tensor,
+) -> dict:
+    model.eval()
+    logits, router_probabilities, sparse_weights = model(pixels)
+    predictions = logits.argmax(dim=1).numpy()
+    utilization = sparse_weights.mean(dim=0).numpy()
+    dominant_by_class = {}
+    for label in range(10):
+        selected = labels == label
+        class_utilization = sparse_weights[selected].mean(dim=0)
+        dominant_by_class[str(label)] = {
+            "expert": int(class_utilization.argmax()),
+            "routing_share": float(class_utilization.max()),
+        }
+    entropy = -(
+        router_probabilities * torch.log(torch.clamp(router_probabilities, min=1e-9))
+    ).sum(dim=1)
+    return {
+        "accuracy": float(accuracy_score(labels.numpy(), predictions)),
+        "macro_f1": float(f1_score(labels.numpy(), predictions, average="macro")),
+        "expert_utilization": utilization.tolist(),
+        "utilization_coefficient_of_variation": float(
+            utilization.std() / utilization.mean()
+        ),
+        "mean_router_entropy": float(entropy.mean()),
+        "maximum_router_entropy": float(np.log(model.expert_count)),
+        "dominant_expert_by_digit": dominant_by_class,
+    }
+
+
+def train_dense(
+    train_pixels: torch.Tensor,
+    train_labels: torch.Tensor,
+    validation_pixels: torch.Tensor,
+    validation_labels: torch.Tensor,
+) -> tuple[DenseControl, dict]:
+    seed_everything(2042)
+    model = DenseControl()
+    loader = DataLoader(
+        TensorDataset(train_pixels, train_labels),
+        batch_size=64,
+        shuffle=True,
+        generator=torch.Generator().manual_seed(2042),
+    )
+    optimizer = torch.optim.AdamW(model.parameters(), lr=0.0025, weight_decay=0.002)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=90)
+    best_accuracy = -1.0
+    best_epoch = 0
+    best_state = None
+    for epoch in range(1, 91):
+        model.train()
+        for pixels, labels in loader:
+            loss = F.cross_entropy(model(pixels), labels)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+        scheduler.step()
+        validation = evaluate_dense(model, validation_pixels, validation_labels)
+        if validation["accuracy"] > best_accuracy:
+            best_accuracy = validation["accuracy"]
+            best_epoch = epoch
+            best_state = {
+                key: value.detach().cpu().clone()
+                for key, value in model.state_dict().items()
+            }
+    assert best_state is not None
+    model.load_state_dict(best_state)
+    return model, {
+        "best_epoch": best_epoch,
+        "best_validation_accuracy": best_accuracy,
+    }
+
+
+def train_moe(
+    train_pixels: torch.Tensor,
+    train_labels: torch.Tensor,
+    validation_pixels: torch.Tensor,
+    validation_labels: torch.Tensor,
+) -> tuple[PocketMoE, dict]:
+    seed_everything(2042)
+    model = PocketMoE()
+    loader = DataLoader(
+        TensorDataset(train_pixels, train_labels),
+        batch_size=64,
+        shuffle=True,
+        generator=torch.Generator().manual_seed(2042),
+    )
+    optimizer = torch.optim.AdamW(model.parameters(), lr=0.0025, weight_decay=0.002)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=90)
+    best_accuracy = -1.0
+    best_epoch = 0
+    best_state = None
+    for epoch in range(1, 91):
+        model.train()
+        losses = []
+        for pixels, labels in loader:
+            logits, router_probabilities, _ = model(pixels)
+            classification = F.cross_entropy(logits, labels)
+            importance = router_probabilities.mean(dim=0)
+            balance = ((importance * model.expert_count - 1) ** 2).mean()
+            loss = classification + 0.025 * balance
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+            losses.append(loss.item())
+        scheduler.step()
+        validation = evaluate_moe(model, validation_pixels, validation_labels)
+        if validation["accuracy"] > best_accuracy:
+            best_accuracy = validation["accuracy"]
+            best_epoch = epoch
+            best_state = {
+                key: value.detach().cpu().clone()
+                for key, value in model.state_dict().items()
+            }
+        if epoch == 1 or epoch % 10 == 0:
+            trackio.log(
+                {
+                    "epoch": epoch,
+                    "moe_train_loss": float(np.mean(losses)),
+                    "moe_validation_accuracy": validation["accuracy"],
+                    "moe_utilization_cv": validation[
+                        "utilization_coefficient_of_variation"
+                    ],
+                    "learning_rate": scheduler.get_last_lr()[0],
+                }
+            )
+    assert best_state is not None
+    model.load_state_dict(best_state)
+    return model, {
+        "best_epoch": best_epoch,
+        "best_validation_accuracy": best_accuracy,
+    }
+
+
+def main() -> None:
+    train_pixels, train_labels = load_split("train")
+    validation_pixels, validation_labels = load_split("validation")
+    test_pixels, test_labels = load_split("test")
+    trackio.init(
+        project="pocket-moe",
+        name="top2-versus-dense-v1",
+        config={
+            "experts": 4,
+            "active_experts": 2,
+            "moe_parameters": parameter_count(PocketMoE()),
+            "dense_parameters": parameter_count(DenseControl()),
+        },
+    )
+    dense, dense_training = train_dense(
+        train_pixels,
+        train_labels,
+        validation_pixels,
+        validation_labels,
+    )
+    moe, moe_training = train_moe(
+        train_pixels,
+        train_labels,
+        validation_pixels,
+        validation_labels,
+    )
+    dense_test = evaluate_dense(dense, test_pixels, test_labels)
+    moe_test = evaluate_moe(moe, test_pixels, test_labels)
+    results = {
+        "model": "Pocket MoE",
+        "experts": moe.expert_count,
+        "active_experts_per_example": moe.top_k,
+        "stored_parameters": parameter_count(moe),
+        "active_parameters_per_example": active_parameter_count(moe),
+        "dense_control_parameters": parameter_count(dense),
+        "moe_training": moe_training,
+        "dense_training": dense_training,
+        "moe_test": moe_test,
+        "dense_control_test": dense_test,
+        "accuracy_delta": moe_test["accuracy"] - dense_test["accuracy"],
+    }
+    trackio.log(
+        {
+            "moe_test_accuracy": moe_test["accuracy"],
+            "dense_test_accuracy": dense_test["accuracy"],
+            "moe_test_utilization_cv": moe_test["utilization_coefficient_of_variation"],
+        }
+    )
+    trackio.finish()
+    moe_dir = ARTIFACT_DIR / "pocket-moe"
+    dense_dir = ARTIFACT_DIR / "dense-control"
+    moe_dir.mkdir(parents=True, exist_ok=True)
+    dense_dir.mkdir(parents=True, exist_ok=True)
+    save_file(moe.state_dict(), moe_dir / "model.safetensors")
+    save_file(dense.state_dict(), dense_dir / "model.safetensors")
+    (moe_dir / "evaluation.json").write_text(
+        json.dumps(results, indent=2),
+        encoding="utf-8",
+    )
+    print(json.dumps(results, indent=2))
+
+
+if __name__ == "__main__":
+    main()