fix: correct DenseLayer activation gradients

Author: ooples

src/NeuralNetworks/Layers/DenseLayer.cs

@@ -135,6 +135,7 @@ public class DenseLayer<T> : LayerBase<T>
     /// </para>
     /// </remarks>
     private Tensor<T>? _lastInput;
+    private Tensor<T>? _lastPreActivation;
 
     /// <summary>
     /// Gets the total number of trainable parameters in the layer.
@@ -370,16 +371,17 @@ public override Tensor<T> Forward(Tensor<T> input)
         _lastInput = input;
 
         var (input2D, squeezeOutput) = EnsureRank2BatchFirst(input);
-        var output = input2D.Multiply(_weights.Transpose()).Add(_biases);
+        var preActivation = input2D.Multiply(_weights.Transpose()).Add(_biases);
+        _lastPreActivation = preActivation;
 
         if (UsingVectorActivation)
         {
-            var activated = VectorActivation!.Activate(output);
+            var activated = VectorActivation!.Activate(preActivation);
             return squeezeOutput ? activated.Reshape([activated.Shape[1]]) : activated;
         }
         else
         {
-            var activated = ApplyActivation(output);
+            var activated = ApplyActivation(preActivation);
             return squeezeOutput ? activated.Reshape([activated.Shape[1]]) : activated;
         }
     }
@@ -411,7 +413,7 @@ public override Tensor<T> Forward(Tensor<T> input)
     /// </remarks>
     public override Tensor<T> Backward(Tensor<T> outputGradient)
     {
-        if (_lastInput == null)
+        if (_lastInput == null || _lastPreActivation == null)
             throw new InvalidOperationException("Forward pass must be called before backward pass.");
 
         var (lastInput2D, _) = EnsureRank2BatchFirst(_lastInput);
@@ -430,20 +432,7 @@ public override Tensor<T> Backward(Tensor<T> outputGradient)
             outputGradient2D = outputGradient.Reshape([batchSize, outputGradient.Length / batchSize]);
         }
 
-        Tensor<T> activationGradient;
-        if (UsingVectorActivation)
-        {
-            activationGradient = VectorActivation!.Derivative(outputGradient2D);
-        }
-        else
-        {
-            // Apply scalar activation derivative element-wise
-            activationGradient = new Tensor<T>(outputGradient2D.Shape);
-            for (int i = 0; i < outputGradient2D.Length; i++)
-            {
-                activationGradient[i] = ScalarActivation!.Derivative(outputGradient2D[i]);
-            }
-        }
+        var activationGradient = ApplyActivationDerivative(_lastPreActivation, outputGradient2D);
 
         _weightsGradient = activationGradient.Transpose([1, 0]).ToMatrix().Multiply(lastInput2D.ToMatrix());
         _biasesGradient = activationGradient.Sum([0]).ToVector();
@@ -638,6 +627,7 @@ public override void ResetState()
     {
         // Clear cached values from forward and backward passes
         _lastInput = null;
+        _lastPreActivation = null;
         _weightsGradient = null;
         _biasesGradient = null;
     }