Fix Softmax

src/NeuralNet/ActivationFunctions/Softmax.php

@@ -75,26 +75,24 @@ public function activate(NDArray $input) : NDArray
     /**
      * Calculate the derivative of the Softmax activation function.
      *
-     * For Softmax, the derivative can be calculated using only the output:
-     * f'(x) = diag(s) - outer(s, s)
-     * where f(x) is the output of the softmax function and s is the softmax output
+     * Returns the element-wise diagonal of each sample's Softmax Jacobian:
+     * f'(x_i) = f(x_i) * (1 - f(x_i))
      *
-     * Since we typically need this for backpropagation where we multiply by the gradient,
-     * we can simplify by using the Jacobian-vector product directly.
+     * The result has the same shape as the input, preserving the
+     * `[classes, batch]` layout where each sample column is treated
+     * independently. Since the full Softmax Jacobian couples all classes of a
+     * sample together, an exact backward pass through this function requires
+     * the Jacobian-vector product which must be handled by the output layer.
      *
      * @param NDArray $input
      * @param NDArray $output The output from the Softmax activation
      * @return NDArray The derivative
      */
     public function differentiate(NDArray $input, NDArray $output) : NDArray
     {
-        $softmax = NumPower::flatten($output);
+        $oneMinusOutput = NumPower::subtract(1.0, $output);
 
-        $diag = NumPower::diag($softmax);
-
-        $outer = NumPower::outer($softmax, $softmax);
-
-        return NumPower::subtract($diag, $outer);
+        return NumPower::multiply($output, $oneMinusOutput);
     }
 
     /**

src/NeuralNet/Layers/Multiclass.php

@@ -196,6 +196,14 @@ public function back(array $labels, Optimizer $optimizer) : array
     /**
      * Calculate the gradient for the previous layer.
      *
+     * For Cross Entropy, the loss derivative cancels with the Softmax
+     * derivative, so dZ = (output - expected). Otherwise, an exact backward
+     * pass through Softmax is computed as the Jacobian-vector product
+     *
+     * dZ_i = s_i * (g_i - sum_j(s_j * g_j))
+     *
+     * where s is the Softmax output and g is the upstream gradient.
+     *
      * @param NDArray $input
      * @param NDArray $output
      * @param NDArray $expected
@@ -212,14 +220,21 @@ public function gradient(NDArray $input, NDArray $output, NDArray $expected) : N
             );
         }
 
-        $dLoss = NumPower::divide(
+        $gradient = NumPower::divide(
             $this->costFn->differentiate($output, $expected),
             $n
         );
 
+        $batch = $output->shape()[1];
+
+        $dot = NumPower::reshape(
+            NumPower::sum(NumPower::multiply($gradient, $output), axis: 0),
+            [1, $batch]
+        );
+
         return NumPower::multiply(
-            $this->softmax->differentiate($input, $output),
-            $dLoss
+            $output,
+            NumPower::subtract($gradient, $dot)
         );
     }
 

tests/NeuralNet/ActivationFunctions/SoftmaxTest.php

@@ -96,8 +96,8 @@ public static function differentiateProvider() : Generator
                 [0.4],
             ]),
             [
-                [0.24, -0.24],
-                [-0.24, 0.24],
+                [0.24],
+                [0.24],
             ],
         ];
 
@@ -108,9 +108,9 @@ public static function differentiateProvider() : Generator
                 [0.2],
             ]),
             [
-                [0.21, -0.15, -0.06],
-                [-0.15, 0.25, -0.10],
-                [-0.06, -0.10, 0.16],
+                [0.21],
+                [0.25],
+                [0.16],
             ],
         ];
 
@@ -120,8 +120,22 @@ public static function differentiateProvider() : Generator
                 [0.7310585],
             ]),
             [
-                [0.1966119, -0.19661192],
-                [-0.1966119, 0.19661192],
+                [0.1966119],
+                [0.1966120],
+            ],
+        ];
+
+        // A batch of 3 samples must be differentiated independently per column.
+        yield [
+            NumPower::array([
+                [0.3097901, 0.5671766, 0.3127109],
+                [0.4762272, 0.2283023, 0.1768459],
+                [0.2139826, 0.2045210, 0.5104430],
+            ]),
+            [
+                [0.2138202, 0.2454873, 0.2149228],
+                [0.2494349, 0.1761804, 0.1455714],
+                [0.1681940, 0.1626922, 0.2498909],
             ],
         ];
     }
@@ -186,14 +200,16 @@ public function testActivate(NDArray $input, array $expected) : void
     }
 
     #[Test]
-    #[TestDox('Correctly activates the input')]
+    #[TestDox('Correctly differentiates the activation')]
     #[DataProvider('differentiateProvider')]
     public function testDifferentiate(NDArray $output, array $expected) : void
     {
         $input = NumPower::zeros($output->shape());
-        $derivatives = $this->activationFn->differentiate($input, $output)->toArray();
+        $derivatives = $this->activationFn->differentiate($input, $output);
+
+        static::assertEquals($output->shape(), $derivatives->shape());
 
-        $this->assertEqualsWithDelta($expected, $derivatives, 1e-7);
+        $this->assertEqualsWithDelta($expected, $derivatives->toArray(), 1e-7);
     }
 
     #[Test]

tests/NeuralNet/Layers/MulticlassTest.php

@@ -17,6 +17,7 @@
 use Rubix\ML\NeuralNet\Layers\Multiclass;
 use Rubix\ML\NeuralNet\Optimizers\Stochastic;
 use Rubix\ML\NeuralNet\CostFunctions\CrossEntropy;
+use Rubix\ML\NeuralNet\CostFunctions\RelativeEntropy;
 use PHPUnit\Framework\TestCase;
 
 #[Group('Layers')]
@@ -195,6 +196,44 @@ public function testGradient(array $expectedGradient) : void
         self::assertEqualsWithDelta($expectedGradient, $gradient->toArray(), 1e-7);
     }
 
+    #[Test]
+    #[TestDox('Computes exact Softmax Jacobian-vector product for losses other than Cross Entropy')]
+    public function testGradientWithRelativeEntropy() : void
+    {
+        $expectedGradient = [
+            [-0.0920019936, 0.0055337012, 0.0691078631],
+            [0.0856411220, -0.1061040102, 0.0001709579],
+            [0.0063608715, 0.1005703090, -0.0692788210],
+        ];
+
+        $layer = new Multiclass(
+            classes: ['hot', 'cold', 'ice cold'],
+            costFn: new RelativeEntropy()
+        );
+
+        $layer->initialize(3);
+
+        $output = $layer->forward($this->input);
+
+        // Rebuild expected one-hot matrix the same way as Multiclass::back()
+        $expected = [];
+
+        foreach (['hot', 'cold', 'ice cold'] as $class) {
+            $row = [];
+
+            foreach ($this->labels as $label) {
+                $row[] = $class === $label ? 1.0 : 0.0;
+            }
+
+            $expected[] = $row;
+        }
+
+        $gradient = $layer->gradient($this->input, $output, NumPower::array($expected));
+
+        self::assertEquals($this->input->shape(), $gradient->shape());
+        self::assertEqualsWithDelta($expectedGradient, $gradient->toArray(), 1e-7);
+    }
+
     #[Test]
     #[TestDox('Computes infer softmax probabilities')]
     #[DataProvider('inferProvider')]