fix: route activation caches through SaveForBackward (sweep 2/3)

src/Models/Options/DQNOptions.cs

@@ -24,9 +24,20 @@ public class DQNOptions<T> : ModelOptions
     public ILossFunction<T> LossFunction { get; set; } = new MeanSquaredErrorLoss<T>();
 
     /// <summary>
-    /// Optimizer for Q-network training. Default: Adam (per Mnih et al. 2015).
+    /// Gets or sets a factory for the optimizer each of this agent's networks trains with.
     /// </summary>
-    public IGradientBasedOptimizer<T, Tensor<T>, Tensor<T>>? Optimizer { get; set; }
+    /// <remarks>
+    /// <para>
+    /// When this is <c>null</c>, the Q-network and its target train with the network's default Adam
+    /// (Mnih et al. 2015 use Adam).
+    /// </para>
+    /// <para>
+    /// The agent invokes the factory once for every network it builds, so the factory must create a new
+    /// optimizer instance on each call. A shared instance would share its step counter, and Adam's bias
+    /// correction would then advance on every network's step instead of only on that network's own.
+    /// </para>
+    /// </remarks>
+    public Func<IGradientBasedOptimizer<T, Tensor<T>, Tensor<T>>>? OptimizerFactory { get; set; }
 
     public List<int> HiddenLayers { get; set; } = new List<int> { 128, 128 };
 

src/Models/Options/DecisionTransformerOptions.cs

@@ -47,7 +47,17 @@ public class DecisionTransformerOptions<T> : ReinforcementLearningOptions<T>
     public int BufferSize { get; init; } = 1000000;
 
     /// <summary>
-    /// The optimizer used for updating network parameters. If null, Adam optimizer will be used by default.
+    /// Gets or sets a factory for the optimizer each of this agent's networks trains with.
     /// </summary>
-    public IOptimizer<T, Vector<T>, Vector<T>>? Optimizer { get; init; }
+    /// <remarks>
+    /// <para>
+    /// When this is <c>null</c>, the network trains with its default optimizer.
+    /// </para>
+    /// <para>
+    /// The agent invokes the factory once for every network it builds, so the factory must create a new
+    /// optimizer instance on each call. A shared instance would share its step counter, and Adam's bias
+    /// correction would then advance on every network's step instead of only on that network's own.
+    /// </para>
+    /// </remarks>
+    public Func<IGradientBasedOptimizer<T, Tensor<T>, Tensor<T>>>? OptimizerFactory { get; init; }
 }

src/Models/Options/MADDPGOptions.cs

@@ -76,9 +76,19 @@ public class MADDPGOptions<T> : ReinforcementLearningOptions<T>
     public List<int> CriticHiddenLayers { get; init; } = new List<int> { 128, 128 };
 
     /// <summary>
-    /// The optimizer used for updating network parameters. If null, Adam optimizer will be used by default.
+    /// Gets or sets a factory for the optimizer each of this agent's networks trains with.
     /// </summary>
-    public IOptimizer<T, Vector<T>, Vector<T>>? Optimizer { get; init; }
+    /// <remarks>
+    /// <para>
+    /// When this is <c>null</c>, every actor and critic trains with the network's default optimizer.
+    /// </para>
+    /// <para>
+    /// The agent invokes the factory once for every network it builds, so the factory must create a new
+    /// optimizer instance on each call. A shared instance would share its step counter, and Adam's bias
+    /// correction would then advance on every network's step instead of only on that network's own.
+    /// </para>
+    /// </remarks>
+    public Func<IGradientBasedOptimizer<T, Tensor<T>, Tensor<T>>>? OptimizerFactory { get; init; }
 
     /// <summary>
     /// Initializes a new instance with default values.
@@ -144,7 +154,7 @@ public MADDPGOptions(MADDPGOptions<T> other) : this()
                 $"{nameof(other.CriticHiddenLayers)} must not be null.", nameof(other));
         }
         CriticHiddenLayers = new List<int>(other.CriticHiddenLayers);
-        Optimizer = other.Optimizer;
+        OptimizerFactory = other.OptimizerFactory;
     }
 
     /// <summary>

src/Models/Options/QMIXOptions.cs

@@ -43,7 +43,18 @@ public class QMIXOptions<T> : ReinforcementLearningOptions<T>
     public List<int> MixingHiddenLayers { get; init; } = new List<int> { 32 };
 
     /// <summary>
-    /// The optimizer used for updating network parameters. If null, Adam optimizer will be used by default.
+    /// Gets or sets a factory for the optimizer of the agent's joint training step.
     /// </summary>
-    public IOptimizer<T, Vector<T>, Vector<T>>? Optimizer { get; init; }
+    /// <remarks>
+    /// <para>
+    /// QMIX trains every agent network and the mixing network on one TD loss in one step (Rashid et al. 2018),
+    /// so the agent invokes the factory once and steps all of them with that optimizer. When this is
+    /// <c>null</c>, the step uses plain SGD at the agent's learning rate.
+    /// </para>
+    /// <para>
+    /// The factory must create a new optimizer instance on each call, so agents built from copies of these
+    /// options never share moment buffers or a step counter.
+    /// </para>
+    /// </remarks>
+    public Func<IGradientBasedOptimizer<T, Tensor<T>, Tensor<T>>>? OptimizerFactory { get; init; }
 }

src/NeuralNetworks/Layers/GatedFeatureLearningUnitLayer.cs

@@ -71,11 +71,8 @@ public partial class GatedFeatureLearningUnitLayer<T> : LayerBase<T>, IShapeCont
     // Gating mechanism
     private readonly FullyConnectedLayer<T> _gateTransform;
 
-    // Cached values
-    [Scratch]
-    private Tensor<T>? _inputCache;
-    [Scratch]
-    private Tensor<T>? _transformedCache;
+    // The last forward's gate values. Not a backward cache: GetGateValues / GetFeatureImportance
+    // read it after an inference forward for interpretability, so it is kept in every mode.
     [Scratch]
     private Tensor<T>? _gateCache;
 
@@ -116,11 +113,9 @@ protected override void OnFirstForward(Tensor<T> input)
     protected override Tensor<T> ForwardTraced(Tensor<T> input)
     {
         EnsureInitializedFromInput(input);
-        _inputCache = input;
 
         // Transform features
         var transformed = _featureTransform.Forward(input);
-        _transformedCache = transformed;
 
         // Compute gate values with sigmoid
         var gateLogits = _gateTransform.Forward(input);
@@ -166,8 +161,6 @@ public override void UpdateParameters(T learningRate)
     /// <inheritdoc/>
     public override void ResetState()
     {
-        _inputCache = null;
-        _transformedCache = null;
         _gateCache = null;
         _featureTransform.ResetState();
         _gateTransform.ResetState();

src/NeuralNetworks/Layers/GatedLinearUnitLayer.cs

@@ -167,71 +167,6 @@ public partial class GatedLinearUnitLayer<T> : LayerBase<T>, IShapeContract
     /// </remarks>
     private Tensor<T> _gateBias;
 
-    /// <summary>
-    /// The input tensor from the last forward pass, saved for backpropagation.
-    /// </summary>
-    /// <remarks>
-    /// <para>
-    /// This tensor stores the input received during the last forward pass. It is necessary for computing
-    /// gradients during the backward pass (backpropagation).
-    /// </para>
-    /// <para><b>For Beginners:</b> This remembers what input data was processed most recently.
-    /// 
-    /// During training:
-    /// - The layer needs to remember what input values it processed
-    /// - This helps when calculating how to improve the weights and biases
-    /// - It's like keeping your work when solving a math problem
-    /// 
-    /// This value is automatically cleared between training batches to save memory.
-    /// </para>
-    /// </remarks>
-    [Scratch]
-    private Tensor<T>? _lastInput;
-
-    /// <summary>
-    /// The output tensor from the linear path of the last forward pass, saved for backpropagation.
-    /// </summary>
-    /// <remarks>
-    /// <para>
-    /// This tensor stores the output from the linear transformation (before gating) during the last forward pass.
-    /// It is used during backpropagation to compute gradients for the gating path.
-    /// </para>
-    /// <para><b>For Beginners:</b> This remembers the values produced by the linear path before gating.
-    /// 
-    /// During training:
-    /// - The layer needs to know what values were being gated
-    /// - This helps calculate how the gate affected the final output
-    /// - It's essential for computing how to improve both paths
-    /// 
-    /// This intermediate result is saved because it's needed when computing
-    /// gradients during the backward pass.
-    /// </para>
-    /// </remarks>
-    [Scratch]
-    private Tensor<T>? _lastLinearOutput;
-
-    /// <summary>
-    /// The output tensor from the gating path of the last forward pass, saved for backpropagation.
-    /// </summary>
-    /// <remarks>
-    /// <para>
-    /// This tensor stores the output from the gating path (after activation) during the last forward pass.
-    /// It is used during backpropagation to compute gradients for the linear path and the gate activation.
-    /// </para>
-    /// <para><b>For Beginners:</b> This remembers how open or closed each gate was in the latest calculation.
-    /// 
-    /// During training:
-    /// - The layer needs to know what gate values were applied
-    /// - This helps calculate how the gates affected the output
-    /// - It's needed for computing gradients for both paths
-    /// 
-    /// These gate values (typically between 0 and 1) determined how much
-    /// of each linear output value passed through to the final output.
-    /// </para>
-    /// </remarks>
-    [Scratch]
-    private Tensor<T>? _lastGateOutput;
-
     /// <summary>
     /// The gradients for the linear weights, computed during backpropagation.
     /// </summary>
@@ -544,7 +479,6 @@ private void InitializeParameters()
     protected override Tensor<T> ForwardTraced(Tensor<T> input)
     {
         EnsureInitializedFromInput(input);
-        _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety)
 
         // Linear path: linear = input @ weights^T + bias
         var linearWeightsT = Engine.TensorPermute(_linearWeights, [1, 0]);
@@ -556,11 +490,9 @@ protected override Tensor<T> ForwardTraced(Tensor<T> input)
         var gateOutput = input.MatrixMultiply(gateWeightsT);
         gateOutput = Engine.TensorAdd(gateOutput, _gateBias); // Broadcasting
 
-        // #1668: cache the linear/gate activations for backward only; forward uses locals.
+        // Locals only: backward runs on the tape, so nothing here needs to outlive the forward.
         var linOut = linearOutput;
         var gateOut = ApplyActivation(gateOutput);
-        _lastLinearOutput = ShouldCacheForBackward ? linOut : null;
-        _lastGateOutput = ShouldCacheForBackward ? gateOut : null;
 
         // GLU output: output = linear * gate
         var output = Engine.TensorMultiply(linOut, gateOut);
@@ -635,9 +567,6 @@ public override void ClearGradients()
     public override void ResetState()
     {
         // Clear cached values from forward and backward passes
-        _lastInput = null;
-        _lastLinearOutput = null;
-        _lastGateOutput = null;
         _linearWeightsGradient = null;
         _gateWeightsGradient = null;
         _linearBiasGradient = null;

src/NeuralNetworks/Layers/GaussianNoiseLayer.cs

@@ -92,31 +92,6 @@ public partial class GaussianNoiseLayer<T> : LayerBase<T>, IShapeContract
     /// </remarks>
     private readonly T _standardDeviation;
 
-    /// <summary>
-    /// The noise tensor from the last forward pass, saved for potential use in backpropagation.
-    /// </summary>
-    /// <remarks>
-    /// <para>
-    /// This field stores the noise tensor that was generated and used during the last forward pass.
-    /// It is saved in case it's needed for the backward pass, though in simple additive noise,
-    /// the gradient flows through unchanged.
-    /// </para>
-    /// <para><b>For Beginners:</b> This remembers what noise was added in the latest calculation.
-    /// 
-    /// During training:
-    /// - The layer needs to keep track of exactly what noise was added
-    /// - This helps ensure consistent behavior during backpropagation
-    /// - It's like keeping a record of what "challenges" were introduced
-    /// 
-    /// This tensor has the same shape as the input data, with each position
-    /// containing the specific noise value that was added at that position.
-    /// </para>
-    /// </remarks>
-    [Scratch]
-    private Tensor<T>? _lastNoise;
-    [Scratch]
-    private Tensor<T>? _lastInput;
-
     /// <summary>
     /// Gets a value indicating whether this layer supports training.
     /// </summary>
@@ -227,11 +202,10 @@ protected override void OnFirstForward(Tensor<T> input)
     protected override Tensor<T> ForwardTraced(Tensor<T> input)
     {
         EnsureInitializedFromInput(input);
-        _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety)
         if (IsTrainingMode)
         {
-            _lastNoise = GenerateNoise(input._shape);
-            return Engine.TensorAdd(input, _lastNoise);
+            var noise = GenerateNoise(input._shape);
+            return Engine.TensorAdd(input, noise);
         }
 
         return input; // During inference, no noise is added
@@ -297,8 +271,7 @@ private Tensor<T> GenerateNoise(int[] shape)
     /// </remarks>
     public override void ResetState()
     {
-        // Clear cached values from forward pass
-        _lastNoise = null;
-        _lastInput = null;
+        // Nothing to clear: the forward keeps no per-pass state.
+        return;
     }
 }

src/NeuralNetworks/Layers/GlobalPoolingLayer.cs

@@ -166,56 +166,6 @@ public partial class GlobalPoolingLayer<T> : LayerBase<T>, IShapeContract
     /// </remarks>
     private readonly PoolingType _poolingType;
 
-    /// <summary>
-    /// The input tensor from the last forward pass, saved for backpropagation.
-    /// </summary>
-    /// <remarks>
-    /// <para>
-    /// This tensor stores the input received during the last forward pass. It is necessary for computing
-    /// gradients during the backward pass, particularly for max pooling which needs to know which
-    /// input values were selected as the maximum.
-    /// </para>
-    /// <para><b>For Beginners:</b> This remembers what input data was processed most recently.
-    /// 
-    /// During training:
-    /// - The layer needs to remember what input values it processed
-    /// - For max pooling, it needs to know which value was the maximum
-    /// - For average pooling, it needs the input dimensions
-    /// - This helps when calculating how to distribute gradients during backpropagation
-    /// 
-    /// This is automatically cleared between training batches to save memory.
-    /// </para>
-    /// </remarks>
-    [Scratch]
-    private Tensor<T>? _lastInput;
-
-    /// <summary>
-    /// The output tensor from the last forward pass, saved for backpropagation.
-    /// </summary>
-    /// <remarks>
-    /// <para>
-    /// This tensor stores the output produced during the last forward pass. It is used during
-    /// backpropagation to compute gradients, particularly when an activation function is applied
-    /// after pooling.
-    /// </para>
-    /// <para><b>For Beginners:</b> This stores what the layer output after its most recent calculation.
-    /// 
-    /// During training:
-    /// - The network needs to remember what pooled values it produced
-    /// - This helps calculate how to improve the network
-    /// - It's especially important if an activation function was applied after pooling
-    /// 
-    /// This is also cleared after each training batch to save memory.
-    /// </para>
-    /// </remarks>
-    [Scratch]
-    private Tensor<T>? _lastOutput;
-
-    /// <summary>
-    /// Stores the indices of the maximum values found during global max pooling.
-    /// </summary>
-    private int[]? _maxIndices;
-
     // GPU-resident cached tensors for GPU training pipeline
     [Scratch]
     private Tensor<T>? _lastOutputGpu;
@@ -474,7 +424,6 @@ private static int[] CreateHighDimensionalOutputShape(int[] inputShape)
     protected override Tensor<T> ForwardTraced(Tensor<T> input)
     {
         EnsureInitializedFromInput(input);
-        _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety)
 
         // Industry-standard: dynamically determine axes based on input rank
         // For any rank, reduce all dimensions except first (batch) and last (features/channels)
@@ -485,30 +434,22 @@ protected override Tensor<T> ForwardTraced(Tensor<T> input)
         {
             // No dimensions to reduce (1D or 2D input)
             output = input;
-            _maxIndices = null;
         }
         else if (_poolingType == PoolingType.Average)
         {
             // Use GPU-accelerated ReduceMean
             // keepDims=true for rank >= 4 to preserve dimension structure for subsequent reshape
             output = Engine.ReduceMean(input, axes, keepDims: input.Shape.Length >= 4);
-            _maxIndices = null;
         }
         else // Max pooling
         {
             // Use GPU-accelerated ReduceMax
             // keepDims=true for rank >= 4 to preserve dimension structure for subsequent reshape
-            output = Engine.ReduceMax(input, axes, keepDims: input.Shape.Length >= 4, out _maxIndices);
+            output = Engine.ReduceMax(input, axes, keepDims: input.Shape.Length >= 4, out _);
         }
 
         var result = ApplyActivation(output);
 
-        // Only store for backward pass during training - skip during inference
-        if (IsTrainingMode)
-        {
-            _lastOutput = result;
-        }
-
         return result;
     }
 
@@ -608,9 +549,6 @@ private Tensor<T> ComputeActivationBackwardGpu(DirectGpuTensorEngine gpuEngine,
     public override void ResetState()
     {
         // Clear cached values from forward and backward passes
-        _lastInput = null;
-        _lastOutput = null;
-        _maxIndices = null;
 
         // Clear GPU cached data
         _lastOutputGpu?.Dispose();