feat: add support for softmax activation and cross-entropy loss in MLP

Author: Fenriir42

bs.md

@@ -30,7 +30,7 @@ This document explains the full journey of building neural models from first pri
 Given inputs $x \in \mathbb{R}^n$, weights $w \in \mathbb{R}^n$, bias $b$:
 
 $$
-z = \sum_i w_i x_i + b,\quad
+z = \sum_{i=1}^{n} w_i x_i + b,\quad
 \hat y = H(z) = \begin{cases}
 1 & z \ge 0 \\
 0 & \text{otherwise}
@@ -58,7 +58,7 @@ Only misclassified samples update parameters (classic perceptron rule, not gradi
 Layer sizes: $L_0$ (input) … $L_k$ (output). For neuron $j$ in layer $L$:
 
 $$
-z_j^{(L)} = \sum_i w_{j,i}^{(L)}\, a_i^{(L-1)} + b_j^{(L)},\quad
+z_j^{(L)} = \sum_{i=1}^{n_{L-1}} w_{j,i}^{(L)}\, a_i^{(L-1)} + b_j^{(L)},\quad
 a_j^{(L)} = \sigma\big(z_j^{(L)}\big) = \frac{1}{1 + e^{-z_j^{(L)}}}.
 $$
 
@@ -71,7 +71,7 @@ Internal storage:
 ### Loss (Per Sample, MSE)
 
 $$
-\mathcal{L} = \sum_j \tfrac{1}{2}\,\big(a_j^{(\text{out})} - y_j\big)^2.
+\mathcal{L} = \sum_{j=1}^{n_{\text{out}}} \tfrac{1}{2}\,\big(a_j^{(\text{out})} - y_j\big)^2.
 $$
 
 Using MSE for simplicity (cross‑entropy would be more suitable for classification but requires softmax modifications).
@@ -89,7 +89,7 @@ $$
 Hidden layer deltas (chain rule):
 
 $$
-\delta_i^{(L)} = \left( \sum_j w_{j,i}^{(L+1)}\, \delta_j^{(L+1)} \right)\, \sigma'\!\big(a_i^{(L)}\big).
+\delta_i^{(L)} = \left( \sum_{j=1}^{n_{L+1}} w_{j,i}^{(L+1)}\, \delta_j^{(L+1)} \right)\, \sigma'\!\big(a_i^{(L)}\big).
 $$
 
 Weight & bias gradients:

bs/mlp.py

@@ -25,16 +25,44 @@ def sigmoid_derivative(y: float) -> float:
     return y * (1.0 - y)
 
 
+def softmax(vec: List[float]) -> List[float]:
+    m = max(vec)
+    exps = [math.exp(v - m) for v in vec]
+    s = sum(exps)
+    return [e / s for e in exps]
+
+
 class MLP:
     def __init__(
-        self, layer_sizes: List[int], lr: float = 0.5, seed: int | None = None
+        self,
+        layer_sizes: List[int],
+        lr: float = 0.5,
+        seed: int | None = None,
+        output_activation: str = "sigmoid",  # or 'softmax'
+        loss: str = "mse",  # 'mse' or 'ce' (cross-entropy only with softmax)
+        class_weights: List[float] | None = None,
     ):
         if len(layer_sizes) < 2:
             raise ValueError("Need at least input and output layer")
         if any(s <= 0 for s in layer_sizes):
             raise ValueError("All layer sizes must be > 0")
         self.layer_sizes = layer_sizes
         self.lr = lr
+        if output_activation not in ("sigmoid", "softmax"):
+            raise ValueError(
+                "output_activation must be 'sigmoid' or 'softmax'"
+            )
+        if loss not in ("mse", "ce"):
+            raise ValueError("loss must be 'mse' or 'ce'")
+        if output_activation == "softmax" and layer_sizes[-1] == 1:
+            raise ValueError("softmax output requires >1 output neurons")
+        if loss == "ce" and output_activation != "softmax":
+            raise ValueError(
+                "cross-entropy currently only supported with softmax output"
+            )
+        self.output_activation = output_activation
+        self.loss = loss
+        self.class_weights = class_weights
         if seed is not None:
             random.seed(seed)
         # Weights: list of matrices (next_layer_size x current_layer_size)
@@ -61,17 +89,25 @@ def to_dict(self) -> dict:
             "learning_rate": self.lr,
             "weights": self.weights,
             "biases": self.biases,
+            "output_activation": self.output_activation,
+            "loss": self.loss,
+            "class_weights": self.class_weights,
         }
 
     @staticmethod
     def from_dict(data: dict) -> "MLP":
         required = {"layer_sizes", "learning_rate", "weights", "biases"}
         if not required.issubset(data):
             raise ValueError("Invalid MLP model file")
+        output_activation = data.get("output_activation", "sigmoid")
+        loss = data.get("loss", "mse")
         mlp = MLP(
-            data["layer_sizes"], lr=data["learning_rate"]
+            data["layer_sizes"],
+            lr=data["learning_rate"],
+            output_activation=output_activation,
+            loss=loss,
+            class_weights=data.get("class_weights"),
         )  # initializes sizes
-        # Replace weights/biases with stored values (shape consistency assumed)
         if len(mlp.weights) != len(data["weights"]):
             raise ValueError("Weights shape mismatch")
         mlp.weights = data["weights"]
@@ -90,15 +126,21 @@ def forward(
             w_mat = self.weights[layer_idx]
             b_vec = self.biases[layer_idx]
             z_layer: List[float] = []
-            a_next: List[float] = []
             for neuron_idx in range(len(w_mat)):
                 w = w_mat[neuron_idx]
                 z = (
                     sum(w_j * a_j for w_j, a_j in zip(w, a))
                     + b_vec[neuron_idx]
                 )
                 z_layer.append(z)
-                a_next.append(sigmoid(z))
+            # Activation choice: last layer may use softmax
+            if (
+                layer_idx == len(self.weights) - 1
+                and self.output_activation == "softmax"
+            ):
+                a_next = softmax(z_layer)
+            else:
+                a_next = [sigmoid(z) for z in z_layer]
             zs.append(z_layer)
             activations.append(a_next)
             a = a_next
@@ -112,7 +154,22 @@ def _backprop_sample(
         self, activations: List[List[float]], target: List[float]
     ) -> Tuple[List[List[List[float]]], List[List[float]], float, bool]:
         output = activations[-1]
-        loss = sum(0.5 * (o - t) ** 2 for o, t in zip(output, target))
+        if self.loss == "mse":
+            loss = sum(0.5 * (o - t) ** 2 for o, t in zip(output, target))
+        else:  # cross-entropy with softmax output
+            # Add small epsilon for numerical stability
+            eps = 1e-12
+            base = -sum(t * math.log(o + eps) for o, t in zip(output, target))
+            if self.class_weights:
+                # weight by true class
+                w_true = 0.0
+                for i, t in enumerate(target):
+                    if t > 0.0:
+                        w_true = self.class_weights[i]
+                        break
+                loss = w_true * base
+            else:
+                loss = base
         is_correct = False
         if len(target) == 1:
             pred_bin = 1 if output[0] >= 0.5 else 0
@@ -126,8 +183,18 @@ def _backprop_sample(
         out_acts = activations[-1]
         delta_out: List[float] = []
         for i in range(len(out_acts)):
-            error = out_acts[i] - target[i]
-            delta_out.append(error * sigmoid_derivative(out_acts[i]))
+            if self.output_activation == "softmax" and self.loss == "ce":
+                # Softmax + cross-entropy simplifies gradient; apply class weighting if provided
+                scale = 1.0
+                if self.class_weights:
+                    for k, t in enumerate(target):
+                        if t > 0.0:
+                            scale = self.class_weights[k]
+                            break
+                delta_out.append(scale * (out_acts[i] - target[i]))
+            else:
+                error = out_acts[i] - target[i]
+                delta_out.append(error * sigmoid_derivative(out_acts[i]))
         deltas[last_layer_idx] = delta_out
         for layer_idx in range(last_layer_idx - 1, -1, -1):
             layer_deltas: List[float] = []

scripts/infer_mlp.py

@@ -39,7 +39,9 @@ def main():
     mlp = MLP.from_dict(data)
     x = args.input
     out = mlp.predict(x)
-    print(f"Input={x}\nRaw Output={out}")
+    print(
+        f"Input={x}\nRaw Output={out}\nOutputActivation={getattr(mlp,'output_activation','sigmoid')}"
+    )
     if len(out) == 1:
         print(f"Binary Thresholded={(1 if out[0] >= 0.5 else 0)}")
     else:

scripts/train_tictactoe_mlp.py

@@ -56,6 +56,16 @@ def main():
     p.add_argument("--mode", choices=["binary", "multi"], default="binary")
     p.add_argument("--hidden", type=int, default=27, help="Hidden layer size")
     p.add_argument("--epochs", type=int, default=4000)
+    p.add_argument(
+        "--softmax",
+        action="store_true",
+        help="Use softmax + cross-entropy for multi-class (ignored in binary)",
+    )
+    p.add_argument(
+        "--no-class-weights",
+        action="store_true",
+        help="Disable class-weighted loss for multi-class (softmax)",
+    )
     p.add_argument("--lr", type=float, default=0.3)
     p.add_argument(
         "--batch", type=int, default=64, help="Mini-batch size (0=full batch)"
@@ -86,16 +96,39 @@ def main():
     if args.mode == "binary":
         data = build_binary_dataset()
         output_size = 1
+        output_activation = "sigmoid"
+        loss = "mse"
     else:
         data = build_multiclass_dataset()
         output_size = 3
+        output_activation = "softmax" if args.softmax else "sigmoid"
+        loss = "ce" if args.softmax else "mse"
+        class_weights = None
+        if args.softmax and not args.no_class_weights:
+            # Compute inverse-frequency normalized weights: N / (K * count_c)
+            counts = [0, 0, 0]
+            for _, y in data:
+                idx = max(range(len(y)), key=lambda i: y[i])
+                counts[idx] += 1
+            N = len(data)
+            K = 3
+            class_weights = [N / (K * c) if c > 0 else 0.0 for c in counts]
+            print(
+                f"Class counts={counts} -> weights={','.join(f'{w:.3f}' for w in class_weights)}"
+            )
 
     train_set, val_set = split_dataset(data, val_ratio=args.val)
     print(
         f"Dataset size: total={len(data)} train={len(train_set)} val={len(val_set)} mode={args.mode}"
     )
 
-    mlp = MLP([9, args.hidden, output_size], lr=args.lr)
+    mlp = MLP(
+        [9, args.hidden, output_size],
+        lr=args.lr,
+        output_activation=output_activation,
+        loss=loss,
+        class_weights=(class_weights if args.mode == "multi" else None),
+    )
     history = mlp.train(
         train_set,
         epochs=args.epochs,