Add refined MatrixKAN with mean-based spline path and MNIST benchmark

.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__/
+*.pt
+*.pth
+*.ckpt
+
+
+experiments/**/data/
+
+experiments/**/*.log
+
+
+
+.DS_Store
+Thumbs.db
+
+
+*.prof
+/tmp/

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
+    ]
+  }
+}

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.