diff --git a/.gitignore b/.gitignore
index d69ee67..0843e19 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,41 @@
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+*.egg-info/
+.eggs/
+dist/
+build/
+
+
 venv/
+.venv/
+env/
+
+
 .idea/
+.vscode/
+*.swp
+*.swo
+
+
 model/
 figures/
-__pycache__/
\ No newline at end of file
+*.pt
+*.pth
+*.ckpt
+
+
+experiments/**/data/
+
+experiments/**/*.log
+
+
+
+.DS_Store
+Thumbs.db
+
+
+*.prof
+/tmp/
diff --git a/experiments/mnist_compare/plots/mnist_comparison.png b/experiments/mnist_compare/plots/mnist_comparison.png
new file mode 100644
index 0000000..0a54bde
Binary files /dev/null and b/experiments/mnist_compare/plots/mnist_comparison.png differ
diff --git a/experiments/mnist_compare/plots/mnist_final_bar.png b/experiments/mnist_compare/plots/mnist_final_bar.png
new file mode 100644
index 0000000..bf1da8a
Binary files /dev/null and b/experiments/mnist_compare/plots/mnist_final_bar.png differ
diff --git a/experiments/mnist_compare/results.json b/experiments/mnist_compare/results.json
new file mode 100644
index 0000000..6c11b92
--- /dev/null
+++ b/experiments/mnist_compare/results.json
@@ -0,0 +1,42 @@
+{
+  "config": {
+    "hidden": 8,
+    "num": 3,
+    "k": 3,
+    "batch_size": 128,
+    "epochs": 50,
+    "max_train_batches": 80,
+    "lr": 0.001,
+    "seed": 42
+  },
+  "original": {
+    "summary": {
+      "epoch": 50,
+      "train_loss": 0.21944646630436182,
+      "train_acc": 0.9365234375,
+      "test_loss": 0.2677607791423798,
+      "test_acc": 0.9212,
+      "epoch_sec": 1.788625517001492,
+      "forward_ms": 8.535971966921352
+    },
+    "total_train_sec": 103.28054552200047,
+    "param_count": 50816
+  },
+  "refined": {
+    "summary": {
+      "epoch": 50,
+      "train_loss": 0.24699158957228065,
+      "train_acc": 0.92822265625,
+      "test_loss": 0.2701142538547516,
+      "test_acc": 0.9256,
+      "epoch_sec": 1.1814144699965254,
+      "forward_ms": 5.467074267168452
+    },
+    "total_train_sec": 69.20829865599808,
+    "param_count": 50816,
+    "edge_weight_shape": [
+      10,
+      8
+    ]
+  }
+}
\ No newline at end of file
diff --git a/experiments/mnist_compare/results.md b/experiments/mnist_compare/results.md
new file mode 100644
index 0000000..85df257
--- /dev/null
+++ b/experiments/mnist_compare/results.md
@@ -0,0 +1,170 @@
+# MNIST: original vs refined MatrixKAN
+
+We trained two small KAN classifiers on MNIST to compare the original layer
+(`MatrixKANLayer.py`) with the refactored one (`refinedMatrixKAN.py`).
+Same architecture, same seed — only the forward pass differs.
+
+## What we ran
+
+| | |
+|---|---|
+| Data | MNIST, pixels scaled to roughly [-1, 1] |
+| Network | 784 → 8 → 10 (two KAN layers) |
+| Spline grid | `num=3`, order `k=3` |
+| Training | 50 epochs, 80 batches/epoch × 128 samples |
+| Optimizer | Adam, lr=0.001 |
+| Hardware | CPU, seed=42 |
+
+## How the refined layer differs
+
+The original builds a full `(batch, in, out)` tensor and scales it with broadcast
+multiply before summing over inputs. The refined version:
+
+1. Means spline basis values along the `(num+k)` axis (same for coef edge weights).
+2. Projects to outputs with `nn.Linear` (no bias).
+3. Applies the residual `base_fun` **after** the input→output linear mix.
+
+Edge weights in the refined model use shape `(out, in)` — the same layout as
+`nn.Linear.weight`. On layer 2 that is `[10, 8]`.
+
+## Bottom line (epoch 50)
+
+| | Original | Refined |
+|---|----------|---------|
+| Test accuracy | 92.1% | **92.6%** (+0.5 pp) |
+| Train accuracy | 93.7% | 92.8% |
+| Time per epoch | 1.8 s | 1.2 s (1.51× faster) |
+| Forward pass | 8.5 ms | 5.5 ms (1.56× faster) |
+| Parameters | 50,816 | 50,816 |
+
+The refined model is about **1.51× faster** per epoch on this setup and reaches
+**slightly higher test accuracy** (92.6% vs 92.1%). The forward path is not
+identical to the original einsum-based spline — we mean bases and coefs over
+`(num+k)` separately instead of pairing them per basis index — but mean
+aggregation keeps the spline scale stable and trains well at this learning rate.
+
+## Plots
+
+![Training curves](plots/mnist_comparison.png)
+
+![Final epoch comparison](plots/mnist_final_bar.png)
+
+## Epoch-by-epoch
+
+### Original
+
+| Epoch | Loss | Train acc | Test acc | Time | Forward |
+|-------|------|-----------|----------|------|---------|
+| 1 | 2.0566 | 26.4% | 37.8% | 2.0s | 8.6ms |
+| 2 | 1.4391 | 48.7% | 60.2% | 1.8s | 8.6ms |
+| 3 | 0.9866 | 68.5% | 75.6% | 1.8s | 8.4ms |
+| 4 | 0.7323 | 78.4% | 81.4% | 1.8s | 8.7ms |
+| 5 | 0.5950 | 82.6% | 84.6% | 1.8s | 8.5ms |
+| 6 | 0.5098 | 85.0% | 85.5% | 1.8s | 8.4ms |
+| 7 | 0.4813 | 86.2% | 86.6% | 1.8s | 8.6ms |
+| 8 | 0.4593 | 86.8% | 87.5% | 1.8s | 8.5ms |
+| 9 | 0.4285 | 88.1% | 87.8% | 1.8s | 8.5ms |
+| 10 | 0.4105 | 88.5% | 88.1% | 1.8s | 8.6ms |
+| 11 | 0.4019 | 88.7% | 88.3% | 1.8s | 8.5ms |
+| 12 | 0.3835 | 89.0% | 89.0% | 1.8s | 8.6ms |
+| 13 | 0.3752 | 89.4% | 89.4% | 1.8s | 8.6ms |
+| 14 | 0.3671 | 89.3% | 89.6% | 1.8s | 8.6ms |
+| 15 | 0.3609 | 90.0% | 89.4% | 1.8s | 8.6ms |
+| 16 | 0.3290 | 90.6% | 90.1% | 1.8s | 8.6ms |
+| 17 | 0.3338 | 90.7% | 90.6% | 1.8s | 8.7ms |
+| 18 | 0.3191 | 91.0% | 90.0% | 1.8s | 8.6ms |
+| 19 | 0.3388 | 90.5% | 90.4% | 1.8s | 8.5ms |
+| 20 | 0.3290 | 90.9% | 90.7% | 1.8s | 8.4ms |
+| 21 | 0.3117 | 91.0% | 91.1% | 1.8s | 8.5ms |
+| 22 | 0.2887 | 91.5% | 90.9% | 1.8s | 8.3ms |
+| 23 | 0.2837 | 91.5% | 90.8% | 1.8s | 8.4ms |
+| 24 | 0.2924 | 91.6% | 91.3% | 1.8s | 8.4ms |
+| 25 | 0.2860 | 92.0% | 91.1% | 1.8s | 8.4ms |
+| 26 | 0.2817 | 92.1% | 91.2% | 1.8s | 8.5ms |
+| 27 | 0.2587 | 92.7% | 91.6% | 1.8s | 8.5ms |
+| 28 | 0.2916 | 91.9% | 91.6% | 1.8s | 8.5ms |
+| 29 | 0.2715 | 92.1% | 91.5% | 1.8s | 8.6ms |
+| 30 | 0.2841 | 92.1% | 91.5% | 1.8s | 8.5ms |
+| 31 | 0.2710 | 92.4% | 91.6% | 1.8s | 8.3ms |
+| 32 | 0.2517 | 93.0% | 91.5% | 1.8s | 8.4ms |
+| 33 | 0.2566 | 93.1% | 91.5% | 1.8s | 8.5ms |
+| 34 | 0.2595 | 92.6% | 91.8% | 1.8s | 8.4ms |
+| 35 | 0.2541 | 93.2% | 91.8% | 1.7s | 8.3ms |
+| 36 | 0.2553 | 93.0% | 91.9% | 1.8s | 8.5ms |
+| 37 | 0.2509 | 92.9% | 91.8% | 1.8s | 8.5ms |
+| 38 | 0.2358 | 93.0% | 91.3% | 1.7s | 8.4ms |
+| 39 | 0.2462 | 93.0% | 92.1% | 1.7s | 8.4ms |
+| 40 | 0.2282 | 93.5% | 92.0% | 1.7s | 8.4ms |
+| 41 | 0.2535 | 93.2% | 92.2% | 1.8s | 8.5ms |
+| 42 | 0.2483 | 93.2% | 91.9% | 1.8s | 8.4ms |
+| 43 | 0.2331 | 93.1% | 92.0% | 1.8s | 8.5ms |
+| 44 | 0.2181 | 93.7% | 92.5% | 1.8s | 8.4ms |
+| 45 | 0.2265 | 93.5% | 92.0% | 1.8s | 8.5ms |
+| 46 | 0.2279 | 93.5% | 92.3% | 1.8s | 8.5ms |
+| 47 | 0.2260 | 93.6% | 92.2% | 1.8s | 8.4ms |
+| 48 | 0.2244 | 93.6% | 91.6% | 1.8s | 8.5ms |
+| 49 | 0.2130 | 94.1% | 92.6% | 1.8s | 8.5ms |
+| 50 | 0.2194 | 93.7% | 92.1% | 1.8s | 8.5ms |
+
+### Refined
+
+| Epoch | Loss | Train acc | Test acc | Time | Forward |
+|-------|------|-----------|----------|------|---------|
+| 1 | 1.2505 | 61.7% | 78.4% | 1.2s | 5.6ms |
+| 2 | 0.7140 | 79.0% | 80.6% | 1.2s | 5.4ms |
+| 3 | 0.6141 | 81.5% | 82.2% | 1.2s | 5.5ms |
+| 4 | 0.5733 | 82.6% | 83.8% | 1.2s | 5.5ms |
+| 5 | 0.5410 | 84.1% | 85.5% | 1.2s | 5.5ms |
+| 6 | 0.5100 | 85.8% | 86.2% | 1.2s | 5.5ms |
+| 7 | 0.4851 | 86.8% | 87.3% | 1.2s | 5.4ms |
+| 8 | 0.4738 | 87.7% | 88.6% | 1.2s | 5.5ms |
+| 9 | 0.4569 | 88.3% | 88.4% | 1.2s | 5.6ms |
+| 10 | 0.4169 | 89.3% | 88.2% | 1.2s | 5.5ms |
+| 11 | 0.4226 | 89.0% | 88.9% | 1.2s | 5.5ms |
+| 12 | 0.4150 | 89.4% | 89.9% | 1.2s | 5.5ms |
+| 13 | 0.3735 | 90.6% | 90.4% | 1.2s | 5.7ms |
+| 14 | 0.3570 | 90.6% | 91.0% | 1.2s | 5.5ms |
+| 15 | 0.3599 | 91.1% | 90.9% | 1.2s | 5.5ms |
+| 16 | 0.3474 | 91.2% | 91.7% | 1.2s | 5.4ms |
+| 17 | 0.3388 | 91.0% | 91.0% | 1.2s | 5.4ms |
+| 18 | 0.3370 | 91.0% | 91.0% | 1.2s | 5.5ms |
+| 19 | 0.3147 | 91.7% | 91.6% | 1.2s | 5.5ms |
+| 20 | 0.3138 | 91.6% | 91.7% | 1.2s | 5.6ms |
+| 21 | 0.3018 | 92.0% | 91.9% | 1.3s | 5.7ms |
+| 22 | 0.3093 | 91.8% | 91.6% | 1.2s | 5.7ms |
+| 23 | 0.3062 | 91.5% | 91.8% | 1.2s | 5.6ms |
+| 24 | 0.2897 | 92.1% | 91.8% | 1.3s | 5.6ms |
+| 25 | 0.2931 | 91.9% | 91.1% | 1.2s | 5.5ms |
+| 26 | 0.2804 | 92.4% | 91.3% | 1.2s | 5.4ms |
+| 27 | 0.2908 | 92.0% | 91.8% | 1.2s | 5.5ms |
+| 28 | 0.2998 | 91.5% | 91.8% | 1.2s | 5.6ms |
+| 29 | 0.2956 | 91.7% | 92.1% | 1.2s | 5.5ms |
+| 30 | 0.2729 | 92.3% | 91.9% | 1.2s | 5.5ms |
+| 31 | 0.2719 | 92.4% | 91.8% | 1.2s | 5.4ms |
+| 32 | 0.2695 | 92.3% | 92.0% | 1.2s | 5.5ms |
+| 33 | 0.2796 | 92.3% | 92.4% | 1.2s | 5.5ms |
+| 34 | 0.2605 | 92.7% | 92.0% | 1.2s | 6.0ms |
+| 35 | 0.2730 | 92.5% | 92.4% | 1.3s | 5.5ms |
+| 36 | 0.2716 | 92.3% | 91.9% | 1.2s | 5.7ms |
+| 37 | 0.2780 | 92.0% | 92.1% | 1.3s | 5.6ms |
+| 38 | 0.2550 | 92.9% | 92.3% | 1.2s | 5.8ms |
+| 39 | 0.2884 | 92.0% | 91.5% | 1.2s | 5.5ms |
+| 40 | 0.2663 | 92.5% | 92.4% | 1.2s | 5.5ms |
+| 41 | 0.2648 | 92.6% | 92.2% | 1.2s | 5.5ms |
+| 42 | 0.2732 | 92.4% | 91.3% | 1.2s | 5.5ms |
+| 43 | 0.2449 | 93.0% | 92.3% | 1.2s | 5.5ms |
+| 44 | 0.2545 | 92.6% | 91.9% | 1.2s | 5.5ms |
+| 45 | 0.2698 | 92.4% | 92.5% | 1.2s | 5.4ms |
+| 46 | 0.2724 | 92.1% | 91.9% | 1.2s | 5.5ms |
+| 47 | 0.2603 | 92.7% | 92.0% | 1.2s | 5.5ms |
+| 48 | 0.2532 | 92.7% | 92.1% | 1.2s | 5.5ms |
+| 49 | 0.2599 | 92.8% | 92.4% | 1.2s | 5.5ms |
+| 50 | 0.2470 | 92.8% | 92.6% | 1.2s | 5.5ms |
+
+## Notes
+
+- This is a 50-epoch run (80 batches/epoch, ~10k train samples per epoch).
+- Refined uses **mean** (not sum) over the `(num+k)` spline basis axis for both
+  activations and edge weights.
+- Refined wins on speed by skipping the `(batch, in, out)` multiply–sum pattern.
+- For higher accuracy, try full MNIST, more epochs, or a wider hidden layer.
diff --git a/experiments/mnist_compare/run_experiment.py b/experiments/mnist_compare/run_experiment.py
new file mode 100644
index 0000000..68df5ed
--- /dev/null
+++ b/experiments/mnist_compare/run_experiment.py
@@ -0,0 +1,441 @@
+"""
+MNIST comparison: original MatrixKANLayer vs refinedMatrixKAN.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import sys
+import time
+import gzip
+import urllib.request
+import struct
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, TensorDataset
+
+ROOT = Path(__file__).resolve().parents[2]
+EXP_DIR = Path(__file__).resolve().parent
+PLOTS_DIR = EXP_DIR / "plots"
+RESULTS_MD = EXP_DIR / "results.md"
+RESULTS_JSON = EXP_DIR / "results.json"
+
+
+def load_layer_module(name: str, filename: str):
+    spec = importlib.util.spec_from_file_location(name, ROOT / filename)
+    mod = importlib.util.module_from_spec(spec)
+    sys.modules[name] = mod
+    spec.loader.exec_module(mod)
+    return mod
+
+
+ORIG = load_layer_module("orig_mkl", "MatrixKANLayer.py")
+REF = load_layer_module("ref_mkl", "refinedMatrixKAN.py")
+
+
+@dataclass
+class EpochStats:
+    epoch: int
+    train_loss: float
+    train_acc: float
+    test_loss: float
+    test_acc: float
+    epoch_sec: float
+    forward_ms: float
+
+
+@dataclass
+class RunResult:
+    name: str
+    epochs: list[EpochStats] = field(default_factory=list)
+    total_train_sec: float = 0.0
+    param_count: int = 0
+    edge_weight_shape: list[int] = field(default_factory=list)
+
+
+class KANMNIST(nn.Module):
+    """Two MatrixKAN layers: 784 -> hidden -> 10."""
+
+    def __init__(self, layer_cls, hidden: int, num: int, k: int, save_plot_data: bool = False):
+        super().__init__()
+        kw = dict(num=num, k=k, grid_range=[-1, 1], device="cpu")
+        import inspect
+        sig = inspect.signature(layer_cls.__init__)
+        if "save_plot_data" in sig.parameters:
+            kw["save_plot_data"] = save_plot_data
+        self.layer1 = layer_cls(in_dim=784, out_dim=hidden, **kw)
+        self.layer2 = layer_cls(in_dim=hidden, out_dim=10, **kw)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() > 2:
+            x = x.view(x.size(0), -1)
+        y, *_ = self.layer1(x)
+        y, *_ = self.layer2(y)
+        return y
+
+    def count_params(self) -> int:
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+
+MNIST_URLS = {
+    "train_images": "https://storage.googleapis.com/cvdf-datasets/mnist/train-images-idx3-ubyte.gz",
+    "train_labels": "https://storage.googleapis.com/cvdf-datasets/mnist/train-labels-idx1-ubyte.gz",
+    "test_images": "https://storage.googleapis.com/cvdf-datasets/mnist/t10k-images-idx3-ubyte.gz",
+    "test_labels": "https://storage.googleapis.com/cvdf-datasets/mnist/t10k-labels-idx1-ubyte.gz",
+}
+
+
+def _download(url: str, dest: Path):
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    if not dest.exists():
+        print(f"Downloading {dest.name}...")
+        urllib.request.urlretrieve(url, dest)
+
+
+def _read_idx_images(path: Path) -> torch.Tensor:
+    with gzip.open(path, "rb") as f:
+        _magic, n, rows, cols = struct.unpack(">IIII", f.read(16))
+        data = torch.tensor(list(f.read()), dtype=torch.float32)
+    return data.view(n, rows * cols)
+
+
+def _read_idx_labels(path: Path) -> torch.Tensor:
+    with gzip.open(path, "rb") as f:
+        _magic, n = struct.unpack(">II", f.read(8))
+        data = torch.tensor(list(f.read()), dtype=torch.long)
+    return data
+
+
+def mnist_loaders(batch_size: int):
+    data_dir = EXP_DIR / "data" / "raw"
+    files = {}
+    for key, url in MNIST_URLS.items():
+        dest = data_dir / url.rsplit("/", 1)[-1]
+        _download(url, dest)
+        files[key] = dest
+
+    def normalize(images):
+        images = images / 255.0
+        images = (images - 0.1307) / 0.3081
+        return images * 2.0 - 1.0
+
+    train_x = normalize(_read_idx_images(files["train_images"]))
+    train_y = _read_idx_labels(files["train_labels"])
+    test_x = normalize(_read_idx_images(files["test_images"]))
+    test_y = _read_idx_labels(files["test_labels"])
+
+    return (
+        DataLoader(TensorDataset(train_x, train_y), batch_size=batch_size, shuffle=True),
+        DataLoader(TensorDataset(test_x, test_y), batch_size=batch_size, shuffle=False),
+    )
+
+
+@torch.no_grad()
+def evaluate(model, loader, criterion, device):
+    model.eval()
+    total_loss, correct, total = 0.0, 0, 0
+    for x, y in loader:
+        x, y = x.to(device), y.to(device)
+        logits = model(x)
+        loss = criterion(logits, y)
+        total_loss += loss.item() * x.size(0)
+        correct += (logits.argmax(1) == y).sum().item()
+        total += x.size(0)
+    return total_loss / total, correct / total
+
+
+def bench_forward_ms(model, device, batch_size=256, repeats=30):
+    model.eval()
+    x = torch.randn(batch_size, 784, device=device)
+    with torch.no_grad():
+        for _ in range(5):
+            model(x)
+        times = []
+        for _ in range(repeats):
+            t0 = time.perf_counter()
+            model(x)
+            times.append((time.perf_counter() - t0) * 1000)
+    return sum(times) / len(times)
+
+
+def align_layer_dtypes(model: nn.Module):
+    """Use float32 everywhere (basis_matrix defaults to float64)."""
+    for m in model.modules():
+        if hasattr(m, "basis_matrix"):
+            m.basis_matrix.data = m.basis_matrix.data.float()
+        if hasattr(m, "grid_range"):
+            m.grid_range.data = m.grid_range.data.float()
+        if hasattr(m, "grid_intervals"):
+            m.grid_intervals.data = m.grid_intervals.data.float()
+
+
+def train_one(name: str, layer_cls, cfg: dict) -> RunResult:
+    torch.manual_seed(cfg["seed"])
+    device = torch.device("cpu")
+
+    train_loader, test_loader = mnist_loaders(cfg["batch_size"])
+    model = KANMNIST(
+        layer_cls,
+        hidden=cfg["hidden"],
+        num=cfg["num"],
+        k=cfg["k"],
+        save_plot_data=False,
+    ).to(device)
+    align_layer_dtypes(model)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=cfg["lr"])
+
+    result = RunResult(name=name, param_count=model.count_params())
+
+    layer2 = model.layer2
+    if hasattr(layer2, "edge_weight"):
+        result.edge_weight_shape = list(layer2.edge_weight.shape)
+    else:
+        ew = layer2.coef.mean(dim=-1) * layer2.mask
+        result.edge_weight_shape = list(ew.shape)
+
+    t_total = time.perf_counter()
+    for epoch in range(1, cfg["epochs"] + 1):
+        model.train()
+        t0 = time.perf_counter()
+        run_loss, correct, total = 0.0, 0, 0
+
+        for batch_idx, (x, y) in enumerate(train_loader):
+            x, y = x.to(device), y.to(device)
+            optimizer.zero_grad()
+            logits = model(x)
+            loss = criterion(logits, y)
+            loss.backward()
+            optimizer.step()
+            run_loss += loss.item() * x.size(0)
+            correct += (logits.argmax(1) == y).sum().item()
+            total += x.size(0)
+            if cfg.get("max_train_batches") and (batch_idx + 1) >= cfg["max_train_batches"]:
+                break
+
+        train_loss = run_loss / total
+        train_acc = correct / total
+        test_loss, test_acc = evaluate(model, test_loader, criterion, device)
+        epoch_sec = time.perf_counter() - t0
+        fwd_ms = bench_forward_ms(model, device, batch_size=cfg["batch_size"])
+
+        result.epochs.append(
+            EpochStats(epoch, train_loss, train_acc, test_loss, test_acc, epoch_sec, fwd_ms)
+        )
+        print(
+            f"[{name}] epoch {epoch}/{cfg['epochs']}  "
+            f"loss={train_loss:.4f}  acc={train_acc:.3f}  "
+            f"test_acc={test_acc:.3f}  time={epoch_sec:.1f}s  fwd={fwd_ms:.2f}ms",
+            flush=True,
+        )
+
+    result.total_train_sec = time.perf_counter() - t_total
+    return result
+
+
+def plot_results(orig: RunResult, ref: RunResult):
+    PLOTS_DIR.mkdir(parents=True, exist_ok=True)
+    epochs = [e.epoch for e in orig.epochs]
+
+    fig, axes = plt.subplots(2, 2, figsize=(11, 8))
+    fig.suptitle("MNIST: original MatrixKANLayer vs refinedMatrixKAN", fontsize=13)
+
+    for res, label, color in [(orig, "original", "#2563eb"), (ref, "refined", "#dc2626")]:
+        axes[0, 0].plot(epochs, [e.train_loss for e in res.epochs], "-o", label=label, color=color)
+        axes[0, 1].plot(epochs, [e.test_acc for e in res.epochs], "-o", label=label, color=color)
+        axes[1, 0].plot(epochs, [e.epoch_sec for e in res.epochs], "-o", label=label, color=color)
+        axes[1, 1].plot(epochs, [e.forward_ms for e in res.epochs], "-o", label=label, color=color)
+
+    axes[0, 0].set_title("Train loss"); axes[0, 0].set_xlabel("epoch"); axes[0, 0].legend()
+    axes[0, 1].set_title("Test accuracy"); axes[0, 1].set_xlabel("epoch"); axes[0, 1].legend()
+    axes[1, 0].set_title("Epoch wall time (s)"); axes[1, 0].set_xlabel("epoch"); axes[1, 0].legend()
+    axes[1, 1].set_title("Forward pass (ms)"); axes[1, 1].set_xlabel("epoch"); axes[1, 1].legend()
+
+    plt.tight_layout()
+    out = PLOTS_DIR / "mnist_comparison.png"
+    fig.savefig(out, dpi=140)
+    plt.close(fig)
+
+    # bar summary (final epoch)
+    fig2, ax = plt.subplots(figsize=(8, 4))
+    labels = ["epoch time (s)", "forward (ms)", "test acc (%)"]
+    o_last, r_last = orig.epochs[-1], ref.epochs[-1]
+    o_vals = [o_last.epoch_sec, o_last.forward_ms, o_last.test_acc * 100]
+    r_vals = [r_last.epoch_sec, r_last.forward_ms, r_last.test_acc * 100]
+    x_pos = torch.arange(len(labels))
+    w = 0.35
+    ax.bar(x_pos - w / 2, o_vals, w, label="original", color="#2563eb")
+    ax.bar(x_pos + w / 2, r_vals, w, label="refined", color="#dc2626")
+    ax.set_xticks(x_pos.tolist())
+    ax.set_xticklabels(labels)
+    ax.set_title("Final epoch comparison")
+    ax.legend()
+    plt.tight_layout()
+    out2 = PLOTS_DIR / "mnist_final_bar.png"
+    fig2.savefig(out2, dpi=140)
+    plt.close(fig2)
+
+    return out, out2
+
+
+def write_markdown(orig: RunResult, ref: RunResult, cfg: dict, plot_paths):
+    o_last, r_last = orig.epochs[-1], ref.epochs[-1]
+    speed_epoch = o_last.epoch_sec / r_last.epoch_sec
+    speed_fwd = o_last.forward_ms / r_last.forward_ms
+
+    acc_delta_pct = (r_last.test_acc - o_last.test_acc) * 100
+    if r_last.test_acc >= o_last.test_acc:
+        test_acc_row = (
+            f"| Test accuracy | {o_last.test_acc:.1%} | **{r_last.test_acc:.1%}** ({acc_delta_pct:+.1f} pp) |"
+        )
+        acc_summary = (
+            f"The refined model is about **{speed_epoch:.2f}× faster** per epoch on this setup and reaches "
+            f"**slightly higher test accuracy** ({r_last.test_acc:.1%} vs {o_last.test_acc:.1%}). "
+            f"The forward path is not identical to the original einsum-based spline — we mean bases and coefs over "
+            f"`(num+k)` separately instead of pairing them per basis index — but mean aggregation keeps the "
+            f"spline scale stable and trains well at this learning rate."
+        )
+    else:
+        test_acc_row = (
+            f"| Test accuracy | **{o_last.test_acc:.1%}** | {r_last.test_acc:.1%} ({acc_delta_pct:+.1f} pp) |"
+        )
+        acc_summary = (
+            f"The refined model is about **{speed_epoch:.2f}× faster** per epoch on this setup. "
+            f"Test accuracy is lower ({r_last.test_acc:.1%} vs {o_last.test_acc:.1%}) — the forward path "
+            f"is not identical to the original einsum-based spline (we mean bases and coefs over `(num+k)` "
+            f"separately instead of pairing them per basis index)."
+        )
+
+    md = f"""# MNIST: original vs refined MatrixKAN
+
+We trained two small KAN classifiers on MNIST to compare the original layer
+(`MatrixKANLayer.py`) with the refactored one (`refinedMatrixKAN.py`).
+Same architecture, same seed — only the forward pass differs.
+
+## What we ran
+
+| | |
+|---|---|
+| Data | MNIST, pixels scaled to roughly [-1, 1] |
+| Network | 784 → {cfg['hidden']} → 10 (two KAN layers) |
+| Spline grid | `num={cfg['num']}`, order `k={cfg['k']}` |
+| Training | {cfg['epochs']} epochs, {cfg.get('max_train_batches') or 'full dataset'} batches/epoch × {cfg['batch_size']} samples |
+| Optimizer | Adam, lr={cfg['lr']} |
+| Hardware | CPU, seed={cfg['seed']} |
+
+## How the refined layer differs
+
+The original builds a full `(batch, in, out)` tensor and scales it with broadcast
+multiply before summing over inputs. The refined version:
+
+1. Means spline basis values along the `(num+k)` axis (same for coef edge weights).
+2. Projects to outputs with `nn.Linear` (no bias).
+3. Applies the residual `base_fun` **after** the input→output linear mix.
+
+Edge weights in the refined model use shape `(out, in)` — the same layout as
+`nn.Linear.weight`. On layer 2 that is `{ref.edge_weight_shape}`.
+
+## Bottom line (epoch {cfg['epochs']})
+
+| | Original | Refined |
+|---|----------|---------|
+{test_acc_row}
+| Train accuracy | {o_last.train_acc:.1%} | {r_last.train_acc:.1%} |
+| Time per epoch | {o_last.epoch_sec:.1f} s | {r_last.epoch_sec:.1f} s ({speed_epoch:.2f}× faster) |
+| Forward pass | {o_last.forward_ms:.1f} ms | {r_last.forward_ms:.1f} ms ({speed_fwd:.2f}× faster) |
+| Parameters | {orig.param_count:,} | {ref.param_count:,} |
+
+{acc_summary}
+
+## Plots
+
+![Training curves](plots/{plot_paths[0].name})
+
+![Final epoch comparison](plots/{plot_paths[1].name})
+
+## Epoch-by-epoch
+
+### Original
+
+| Epoch | Loss | Train acc | Test acc | Time | Forward |
+|-------|------|-----------|----------|------|---------|
+"""
+    for e in orig.epochs:
+        md += f"| {e.epoch} | {e.train_loss:.4f} | {e.train_acc:.1%} | {e.test_acc:.1%} | {e.epoch_sec:.1f}s | {e.forward_ms:.1f}ms |\n"
+
+    md += """
+### Refined
+
+| Epoch | Loss | Train acc | Test acc | Time | Forward |
+|-------|------|-----------|----------|------|---------|
+"""
+    for e in ref.epochs:
+        md += f"| {e.epoch} | {e.train_loss:.4f} | {e.train_acc:.1%} | {e.test_acc:.1%} | {e.epoch_sec:.1f}s | {e.forward_ms:.1f}ms |\n"
+
+    md += f"""
+## Notes
+
+- This is a {cfg['epochs']}-epoch run{' (full MNIST)' if not cfg.get('max_train_batches') else f' ({cfg["max_train_batches"]} batches/epoch, ~{cfg["max_train_batches"] * cfg["batch_size"]:,} train samples/epoch)'}.
+- Refined uses **mean** (not sum) over the `(num+k)` spline basis axis for both
+  activations and edge weights.
+- Refined wins on speed by skipping the `(batch, in, out)` multiply–sum pattern.
+- For higher accuracy, try full MNIST, more epochs, or a wider hidden layer.
+"""
+    RESULTS_MD.write_text(md)
+    RESULTS_JSON.write_text(json.dumps({
+        "config": cfg,
+        "original": {"summary": asdict(orig.epochs[-1]), "total_train_sec": orig.total_train_sec,
+                     "param_count": orig.param_count},
+        "refined": {"summary": asdict(ref.epochs[-1]), "total_train_sec": ref.total_train_sec,
+                    "param_count": ref.param_count, "edge_weight_shape": ref.edge_weight_shape},
+    }, indent=2))
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Compare original vs refined MatrixKAN on MNIST")
+    parser.add_argument("--epochs", type=int, default=3)
+    parser.add_argument("--hidden", type=int, default=8)
+    parser.add_argument("--batch-size", type=int, default=128)
+    parser.add_argument("--max-train-batches", type=int, default=None,
+                        help="Cap batches per epoch (default: full training set)")
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    cfg = {
+        "hidden": args.hidden,
+        "num": 3,
+        "k": 3,
+        "batch_size": args.batch_size,
+        "epochs": args.epochs,
+        "max_train_batches": args.max_train_batches,
+        "lr": args.lr,
+        "seed": args.seed,
+    }
+    PLOTS_DIR.mkdir(parents=True, exist_ok=True)
+
+    print("=" * 60)
+    print("MNIST compare — original MatrixKANLayer")
+    print("=" * 60)
+    orig = train_one("original", ORIG.MatrixKANLayer, cfg)
+
+    print("\n" + "=" * 60)
+    print("MNIST compare — refinedMatrixKAN")
+    print("=" * 60)
+    ref = train_one("refined", REF.MatrixKANLayer, cfg)
+
+    plots = plot_results(orig, ref)
+    write_markdown(orig, ref, cfg, plots)
+    print(f"\nSaved: {RESULTS_MD}")
+    print(f"Plots: {plots[0]}, {plots[1]}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/refinedMatrixKAN.py b/refinedMatrixKAN.py
new file mode 100644
index 0000000..3eadd86
--- /dev/null
+++ b/refinedMatrixKAN.py
@@ -0,0 +1,331 @@
+"""
+MatrixKANLayer — matrix-form B-spline KAN layer for PyTorch.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import kan
+from kan.spline import coef2curve, curve2coef, extend_grid
+from kan.utils import sparse_mask
+
+
+class MatrixKANLayer(kan.KANLayer, nn.Module):
+    """
+    KAN layer with B-splines evaluated via a precomputed basis matrix.
+
+    Forward path (refined):
+        spline basis → mean over (num+k) → nn.Linear → output
+        residual branch → nn.Linear → base_fun → output
+    """
+
+    def __init__(
+        self,
+        in_dim: int = 3,
+        out_dim: int = 2,
+        num: int = 5,
+        k: int = 3,
+        noise_scale: float = 0.5,
+        scale_base_mu: float = 0.0,
+        scale_base_sigma: float = 1.0,
+        scale_sp: float = 1.0,
+        base_fun: nn.Module | None = None,
+        grid_eps: float = 0.02,
+        grid_range: list | None = None,
+        sp_trainable: bool = True,
+        sb_trainable: bool = True,
+        save_plot_data: bool = True,
+        device: str = "cpu",
+        sparse_init: bool = False,
+    ):
+        nn.Module.__init__(self)
+
+        if base_fun is None:
+            base_fun = nn.SiLU()
+        if grid_range is None:
+            grid_range = [-1, 1]
+
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.num = num
+        self.k = k
+        self.base_fun = base_fun
+        self.grid_eps = grid_eps
+        self.device = device
+        self.save_plot_data = save_plot_data
+
+        self._init_grid(grid_range, device)
+        self._init_coef(noise_scale)
+        self._init_linears(
+            scale_base_mu, scale_base_sigma, scale_sp,
+            in_dim, out_dim, sparse_init, sp_trainable, sb_trainable,
+        )
+
+        basis = self._compute_basis_matrix()
+        self.register_buffer("basis_matrix", basis, persistent=True)
+
+        self.to(device)
+
+    # ------------------------------------------------------------------
+    # Init
+    # ------------------------------------------------------------------
+
+    def _init_grid(self, grid_range: list, device: str) -> None:
+        grid = torch.linspace(grid_range[0], grid_range[1], steps=self.num + 1)
+        grid = grid[None, :].expand(self.in_dim, self.num + 1)
+        grid = extend_grid(grid, k_extend=self.k)
+        self.grid = nn.Parameter(grid, requires_grad=False)
+
+        range_t = (
+            torch.tensor(grid_range, device=device, dtype=torch.float32)
+            .unsqueeze(0)
+            .expand(self.in_dim, -1)
+            .clone()
+        )
+        self.grid_range = nn.Parameter(range_t, requires_grad=False)
+        self.grid_intervals = nn.Parameter(
+            (self.grid_range[:, 1] - self.grid_range[:, 0]) / self.num,
+            requires_grad=False,
+        )
+
+    def _init_coef(self, noise_scale: float) -> None:
+        noises = (torch.rand(self.num + 1, self.in_dim, self.out_dim) - 0.5)
+        noises = noises * noise_scale / self.num
+        inner_grid = self.grid[:, self.k : -self.k].permute(1, 0)
+        self.coef = nn.Parameter(curve2coef(inner_grid, noises, self.grid, self.k))
+
+    def _init_linears(
+        self,
+        mu: float,
+        sigma: float,
+        scale_sp: float,
+        in_dim: int,
+        out_dim: int,
+        sparse_init: bool,
+        sp_trainable: bool,
+        sb_trainable: bool,
+    ) -> None:
+        inv_sqrt = 1.0 / np.sqrt(in_dim)
+        base_scale = mu * inv_sqrt + sigma * (torch.rand(in_dim, out_dim) * 2 - 1) * inv_sqrt
+
+        mask = sparse_mask(in_dim, out_dim) if sparse_init else torch.ones(in_dim, out_dim)
+        self.mask = nn.Parameter(mask, requires_grad=False)
+
+        self.base_linear = nn.Linear(in_dim, out_dim, bias=False)
+        self.spline_linear = nn.Linear(in_dim, out_dim, bias=False)
+
+        with torch.no_grad():
+            self.base_linear.weight.copy_((base_scale * self.mask).T)
+            self.spline_linear.weight.copy_(
+                (torch.ones(in_dim, out_dim) * scale_sp * self.mask).T
+            )
+
+        self.base_linear.weight.requires_grad_(sb_trainable)
+        self.spline_linear.weight.requires_grad_(sp_trainable)
+
+    # ------------------------------------------------------------------
+    # Weights (pykan-compatible views)
+    # ------------------------------------------------------------------
+
+    @property
+    def scale_base(self) -> torch.Tensor:
+        return self.base_linear.weight.T
+
+    @scale_base.setter
+    def scale_base(self, value: torch.Tensor) -> None:
+        self.base_linear.weight.data.copy_(value.T)
+
+    @property
+    def scale_sp(self) -> torch.Tensor:
+        return self.spline_linear.weight.T
+
+    @scale_sp.setter
+    def scale_sp(self, value: torch.Tensor) -> None:
+        self.spline_linear.weight.data.copy_(value.T)
+
+    @property
+    def edge_weight(self) -> torch.Tensor:
+        """Coef reduced over spline bases, shape (out_dim, in_dim) — matches Linear.weight."""
+        return (self.coef.mean(dim=-1) * self.mask).T
+
+    @property
+    def edge_weight_io(self) -> torch.Tensor:
+        """Same as edge_weight but shape (in_dim, out_dim), for per-edge plots."""
+        return self.coef.mean(dim=-1) * self.mask
+
+    # ------------------------------------------------------------------
+    # Basis matrix 
+    # ------------------------------------------------------------------
+
+    def _compute_basis_matrix(self) -> torch.Tensor:
+        basis = torch.tensor([[1.0]], dtype=torch.float32, device=self.device)
+        scalar = 1.0
+
+        for order in range(2, self.k + 2):
+            prev = order - 1
+            lower = F.pad(basis, (0, 0, 0, 1))
+            upper = F.pad(basis, (0, 0, 1, 0))
+
+            c_pos = torch.zeros(prev, order, device=self.device, dtype=basis.dtype)
+            c_der = torch.zeros(prev, order, device=self.device, dtype=basis.dtype)
+            for i in range(prev):
+                c_pos[i, i] = i + 1
+                c_pos[i, i + 1] = order - (i + 2)
+                c_der[i, i] = -1
+                c_der[i, i + 1] = 1
+
+            basis = lower @ c_pos + upper @ c_der
+            scalar /= prev
+
+        return basis * scalar
+
+    # ------------------------------------------------------------------
+    # Spline evaluation
+    # ------------------------------------------------------------------
+
+    def _power_bases(self, x: torch.Tensor):
+        """Return local monomials [1, u, u², …] and interval membership mask."""
+        x3d = x.unsqueeze(2)
+        grid = self.grid.unsqueeze(0)
+
+        x_intervals = (x3d >= grid[:, :, :-1]) & (x3d < grid[:, :, 1:])
+
+        interval_idx = torch.argmax(x_intervals.int(), dim=-1, keepdim=True)
+        grid_start = self.grid[:, 0].unsqueeze(0).expand(x.shape[0], -1)
+        floor = interval_idx.squeeze(-1) * self.grid_intervals + grid_start
+        ceil = floor + self.grid_intervals
+
+        u1 = ((x - floor) / (ceil - floor)).unsqueeze(-1)
+        powers = torch.cat([u1 ** i for i in range(self.k + 1)], dim=-1)
+        powers[..., 0] = 1.0
+        return powers, x_intervals
+
+    def _b_spline_basis(self, x: torch.Tensor) -> torch.Tensor:
+        """Evaluate all basis functions: (batch, in_dim, num+k)."""
+        power_bases, x_intervals = self._power_bases(x)
+
+        pad = self.k + self.num
+        basis_padded = F.pad(self.basis_matrix, (pad, pad))
+        basis_exp = basis_padded.unsqueeze(0).unsqueeze(0).expand(
+            x.shape[0], self.in_dim, -1, -1
+        )
+
+        oob = torch.zeros(
+            *x_intervals.shape[:2], 1, dtype=torch.bool, device=x.device
+        )
+        intervals = torch.cat([oob, x_intervals], dim=-1)
+        floor_idx = torch.argmax(intervals.int(), dim=-1, keepdim=True)
+
+        start_col = 2 * self.k + self.num - floor_idx + 1
+        col_range = torch.arange(self.k + self.num, device=x.device)
+        col_idx = (start_col.unsqueeze(-2) + col_range.view(1, 1, 1, -1)).expand(
+            -1, -1, self.k + 1, -1
+        )
+
+        gathered = torch.gather(basis_exp, -1, col_idx)
+        out = torch.matmul(power_bases.unsqueeze(-2), gathered).squeeze(-2)
+        return torch.nan_to_num(out)
+
+    def _b_spline_output(self, x: torch.Tensor) -> torch.Tensor:
+        """Mean of basis activations per input dim → (batch, in_dim)."""
+        return self._b_spline_basis(x).mean(dim=-1)
+
+    @staticmethod
+    def _masked_linear(x: torch.Tensor, linear: nn.Linear, mask_out_in: torch.Tensor) -> torch.Tensor:
+        """Linear layer with elementwise mask on weight; mask shape (out_dim, in_dim)."""
+        return F.linear(x, linear.weight * mask_out_in)
+
+    # ------------------------------------------------------------------
+    # Forward
+    # ------------------------------------------------------------------
+
+    def forward(self, x: torch.Tensor):
+        mask_out_in = self.mask.T
+
+        spline_sum = self._b_spline_output(x)
+        base_out = self.base_fun(
+            self._masked_linear(x, self.base_linear, mask_out_in)
+        )
+        spline_out = self._masked_linear(spline_sum, self.spline_linear, self.edge_weight)
+        y = base_out + spline_out
+
+        if not self.save_plot_data:
+            return y, None, None, None
+
+        return y, *self._plot_tensors(x, spline_sum)
+
+    def _plot_tensors(self, x: torch.Tensor, spline_sum: torch.Tensor):
+        """Per-edge tensors for pykan-style interpretability plots."""
+        batch = x.shape[0]
+        base_w = self.scale_base * self.mask
+        spline_w = self.scale_sp * self.edge_weight_io
+
+        preacts = x[:, None, :].expand(batch, self.out_dim, self.in_dim).clone()
+        base_edge = x[:, :, None] * base_w.unsqueeze(0)
+        spline_edge = spline_sum[:, :, None] * spline_w.unsqueeze(0)
+        postspline = spline_edge.permute(0, 2, 1).clone()
+        postacts = (base_edge + spline_edge).permute(0, 2, 1).clone()
+        return preacts, postacts, postspline
+
+    # ------------------------------------------------------------------
+    # Grid updates
+    # ------------------------------------------------------------------
+
+    def _adaptive_grid(self, x_sorted: torch.Tensor, num_interval: int) -> torch.Tensor:
+        batch = x_sorted.shape[0]
+        ids = [int(batch / num_interval * i) for i in range(num_interval)] + [-1]
+        grid_adaptive = x_sorted[ids, :].permute(1, 0)
+
+        step = (grid_adaptive[:, [-1]] - grid_adaptive[:, [0]]) / num_interval
+        grid_uniform = grid_adaptive[:, [0]] + step * torch.arange(
+            num_interval + 1, device=x_sorted.device
+        ).unsqueeze(0)
+
+        return self.grid_eps * grid_uniform + (1 - self.grid_eps) * grid_adaptive
+
+    def _apply_new_grid(
+        self, grid: torch.Tensor, x_pos: torch.Tensor, y_eval: torch.Tensor
+    ) -> None:
+        self.grid_range[:, 0] = grid[:, 0]
+        self.grid_range[:, 1] = grid[:, -1]
+        self.grid_intervals.data = (self.grid_range[:, 1] - self.grid_range[:, 0]) / self.num
+        self.grid.data = extend_grid(grid, k_extend=self.k)
+        self.coef.data = curve2coef(x_pos, y_eval, self.grid, self.k)
+
+    def update_grid_from_samples(self, x: torch.Tensor, mode: str = "sample") -> None:
+        x_sorted = torch.sort(x, dim=0).values
+        y_eval = coef2curve(x_sorted, self.grid, self.coef, self.k)
+        num_interval = self.grid.shape[1] - 1 - 2 * self.k
+
+        grid = self._adaptive_grid(x_sorted, num_interval)
+
+        if mode == "grid":
+            sample_grid = self._adaptive_grid(x_sorted, 2 * num_interval)
+            x_sorted = sample_grid.permute(1, 0)
+            y_eval = coef2curve(x_sorted, self.grid, self.coef, self.k)
+
+        self._apply_new_grid(grid, x_sorted, y_eval)
+
+    def initialize_grid_from_parent(
+        self, parent: MatrixKANLayer, x: torch.Tensor, mode: str = "sample"
+    ) -> None:
+        x_sorted = torch.sort(x, dim=0).values
+        y_eval = coef2curve(x_sorted, parent.grid, parent.coef, parent.k)
+        num_interval = self.grid.shape[1] - 1 - 2 * self.k
+
+        grid = self._adaptive_grid(x_sorted, num_interval)
+
+        if mode == "grid":
+            sample_grid = self._adaptive_grid(x_sorted, 2 * num_interval)
+            x_sorted = sample_grid.permute(1, 0)
+            y_eval = coef2curve(x_sorted, parent.grid, parent.coef, parent.k)
+
+        self._apply_new_grid(grid, x_sorted, y_eval)
+
+    def __getattribute__(self, name: str):
+        if name == "KANLayer":
+            return MatrixKANLayer
+        return super().__getattribute__(name)