diff --git a/src/Models/Options/DQNOptions.cs b/src/Models/Options/DQNOptions.cs index ad8419e9d4..61c6486201 100644 --- a/src/Models/Options/DQNOptions.cs +++ b/src/Models/Options/DQNOptions.cs @@ -24,9 +24,20 @@ public class DQNOptions : ModelOptions public ILossFunction LossFunction { get; set; } = new MeanSquaredErrorLoss(); /// - /// 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. /// - public IGradientBasedOptimizer, Tensor>? Optimizer { get; set; } + /// + /// + /// When this is null, the Q-network and its target train with the network's default Adam + /// (Mnih et al. 2015 use Adam). + /// + /// + /// 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. + /// + /// + public Func, Tensor>>? OptimizerFactory { get; set; } public List HiddenLayers { get; set; } = new List { 128, 128 }; diff --git a/src/Models/Options/DecisionTransformerOptions.cs b/src/Models/Options/DecisionTransformerOptions.cs index 1b5f7c2406..f90b1dcc7c 100644 --- a/src/Models/Options/DecisionTransformerOptions.cs +++ b/src/Models/Options/DecisionTransformerOptions.cs @@ -47,7 +47,17 @@ public class DecisionTransformerOptions : ReinforcementLearningOptions public int BufferSize { get; init; } = 1000000; /// - /// 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. /// - public IOptimizer, Vector>? Optimizer { get; init; } + /// + /// + /// When this is null, the network trains with its default optimizer. + /// + /// + /// 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. + /// + /// + public Func, Tensor>>? OptimizerFactory { get; init; } } diff --git a/src/Models/Options/MADDPGOptions.cs b/src/Models/Options/MADDPGOptions.cs index 1659a4470c..48553b5a88 100644 --- a/src/Models/Options/MADDPGOptions.cs +++ b/src/Models/Options/MADDPGOptions.cs @@ -76,9 +76,19 @@ public class MADDPGOptions : ReinforcementLearningOptions public List CriticHiddenLayers { get; init; } = new List { 128, 128 }; /// - /// 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. /// - public IOptimizer, Vector>? Optimizer { get; init; } + /// + /// + /// When this is null, every actor and critic trains with the network's default optimizer. + /// + /// + /// 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. + /// + /// + public Func, Tensor>>? OptimizerFactory { get; init; } /// /// Initializes a new instance with default values. @@ -144,7 +154,7 @@ public MADDPGOptions(MADDPGOptions other) : this() $"{nameof(other.CriticHiddenLayers)} must not be null.", nameof(other)); } CriticHiddenLayers = new List(other.CriticHiddenLayers); - Optimizer = other.Optimizer; + OptimizerFactory = other.OptimizerFactory; } /// diff --git a/src/Models/Options/QMIXOptions.cs b/src/Models/Options/QMIXOptions.cs index ae3034a804..9ee8da4edd 100644 --- a/src/Models/Options/QMIXOptions.cs +++ b/src/Models/Options/QMIXOptions.cs @@ -43,7 +43,18 @@ public class QMIXOptions : ReinforcementLearningOptions public List MixingHiddenLayers { get; init; } = new List { 32 }; /// - /// 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. /// - public IOptimizer, Vector>? Optimizer { get; init; } + /// + /// + /// 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 + /// null, the step uses plain SGD at the agent's learning rate. + /// + /// + /// 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. + /// + /// + public Func, Tensor>>? OptimizerFactory { get; init; } } diff --git a/src/NeuralNetworks/Layers/GatedFeatureLearningUnitLayer.cs b/src/NeuralNetworks/Layers/GatedFeatureLearningUnitLayer.cs index 77a61728e2..5593aea762 100644 --- a/src/NeuralNetworks/Layers/GatedFeatureLearningUnitLayer.cs +++ b/src/NeuralNetworks/Layers/GatedFeatureLearningUnitLayer.cs @@ -71,11 +71,8 @@ public partial class GatedFeatureLearningUnitLayer : LayerBase, IShapeCont // Gating mechanism private readonly FullyConnectedLayer _gateTransform; - // Cached values - [Scratch] - private Tensor? _inputCache; - [Scratch] - private Tensor? _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? _gateCache; @@ -116,11 +113,9 @@ protected override void OnFirstForward(Tensor input) protected override Tensor ForwardTraced(Tensor 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) /// public override void ResetState() { - _inputCache = null; - _transformedCache = null; _gateCache = null; _featureTransform.ResetState(); _gateTransform.ResetState(); diff --git a/src/NeuralNetworks/Layers/GatedLinearUnitLayer.cs b/src/NeuralNetworks/Layers/GatedLinearUnitLayer.cs index f66c8c378c..27050c559e 100644 --- a/src/NeuralNetworks/Layers/GatedLinearUnitLayer.cs +++ b/src/NeuralNetworks/Layers/GatedLinearUnitLayer.cs @@ -167,71 +167,6 @@ public partial class GatedLinearUnitLayer : LayerBase, IShapeContract /// private Tensor _gateBias; - /// - /// The input tensor from the last forward pass, saved for backpropagation. - /// - /// - /// - /// This tensor stores the input received during the last forward pass. It is necessary for computing - /// gradients during the backward pass (backpropagation). - /// - /// For Beginners: 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. - /// - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// The output tensor from the linear path of the last forward pass, saved for backpropagation. - /// - /// - /// - /// 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. - /// - /// For Beginners: 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. - /// - /// - [Scratch] - private Tensor? _lastLinearOutput; - - /// - /// The output tensor from the gating path of the last forward pass, saved for backpropagation. - /// - /// - /// - /// 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. - /// - /// For Beginners: 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. - /// - /// - [Scratch] - private Tensor? _lastGateOutput; - /// /// The gradients for the linear weights, computed during backpropagation. /// @@ -544,7 +479,6 @@ private void InitializeParameters() protected override Tensor ForwardTraced(Tensor 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 ForwardTraced(Tensor 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; diff --git a/src/NeuralNetworks/Layers/GaussianNoiseLayer.cs b/src/NeuralNetworks/Layers/GaussianNoiseLayer.cs index d869f7a521..fbfa67b328 100644 --- a/src/NeuralNetworks/Layers/GaussianNoiseLayer.cs +++ b/src/NeuralNetworks/Layers/GaussianNoiseLayer.cs @@ -92,31 +92,6 @@ public partial class GaussianNoiseLayer : LayerBase, IShapeContract /// private readonly T _standardDeviation; - /// - /// The noise tensor from the last forward pass, saved for potential use in backpropagation. - /// - /// - /// - /// 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. - /// - /// For Beginners: 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. - /// - /// - [Scratch] - private Tensor? _lastNoise; - [Scratch] - private Tensor? _lastInput; - /// /// Gets a value indicating whether this layer supports training. /// @@ -227,11 +202,10 @@ protected override void OnFirstForward(Tensor input) protected override Tensor ForwardTraced(Tensor 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 GenerateNoise(int[] shape) /// 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; } } diff --git a/src/NeuralNetworks/Layers/GlobalPoolingLayer.cs b/src/NeuralNetworks/Layers/GlobalPoolingLayer.cs index 3b38dddc63..e903745d86 100644 --- a/src/NeuralNetworks/Layers/GlobalPoolingLayer.cs +++ b/src/NeuralNetworks/Layers/GlobalPoolingLayer.cs @@ -166,56 +166,6 @@ public partial class GlobalPoolingLayer : LayerBase, IShapeContract /// private readonly PoolingType _poolingType; - /// - /// The input tensor from the last forward pass, saved for backpropagation. - /// - /// - /// - /// 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. - /// - /// For Beginners: 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. - /// - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// The output tensor from the last forward pass, saved for backpropagation. - /// - /// - /// - /// 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. - /// - /// For Beginners: 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. - /// - /// - [Scratch] - private Tensor? _lastOutput; - - /// - /// Stores the indices of the maximum values found during global max pooling. - /// - private int[]? _maxIndices; - // GPU-resident cached tensors for GPU training pipeline [Scratch] private Tensor? _lastOutputGpu; @@ -474,7 +424,6 @@ private static int[] CreateHighDimensionalOutputShape(int[] inputShape) protected override Tensor ForwardTraced(Tensor 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 ForwardTraced(Tensor 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 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(); diff --git a/src/NeuralNetworks/Layers/GraphAttentionLayer.cs b/src/NeuralNetworks/Layers/GraphAttentionLayer.cs index 1ef0f0ae1c..2155487365 100644 --- a/src/NeuralNetworks/Layers/GraphAttentionLayer.cs +++ b/src/NeuralNetworks/Layers/GraphAttentionLayer.cs @@ -140,23 +140,11 @@ private T GetAdjacencyValue(int b, int i, int j) /// private bool _useSparseAggregation = false; - /// - /// Cached input from forward pass for backward computation. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Stores the original input shape for any-rank tensor support. /// private int[]? _originalInputShape; - /// - /// Cached output from forward pass for backward computation. - /// - [Scratch] - private Tensor? _lastOutput; - /// /// Cached attention coefficients from forward pass. /// @@ -389,13 +377,6 @@ protected override Tensor ForwardTraced(Tensor input) processInput = Engine.Reshape(input, [flatBatch, input.Shape[rank - 2], input.Shape[rank - 1]]); } - // #1668: compute the backward-cache decision once up front so every backward-only cache - // in both the sparse and dense paths (and _lastInput here) is gated/cleared consistently. - // Clearing (not just skipping) is required: a prior step's arena-owned tensor must not - // survive into the next denoise-loop Reset(). - bool cacheBwd = ShouldCacheForBackward; - _lastInput = cacheBwd ? processInput : null; - Tensor output; Tensor activatedOutput; @@ -466,9 +447,6 @@ protected override Tensor ForwardTraced(Tensor input) sparseReshaped = Engine.Reshape(activatedOutput, outputShape); } - // #1668: _lastOutput is a backward-only cache; clear it in inference so the - // arena-owned reshaped output is not retained across the next Reset(). - _lastOutput = cacheBwd ? sparseReshaped : null; return sparseReshaped; } @@ -596,9 +574,6 @@ protected override Tensor ForwardTraced(Tensor input) denseReshaped = Engine.Reshape(activatedOutput, outputShape); } - // #1668: _lastOutput is a backward-only cache; clear it in inference so the - // arena-owned reshaped output is not retained across the next Reset(). - _lastOutput = cacheBwd ? denseReshaped : null; return denseReshaped; } @@ -893,8 +868,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; _lastAttentionCoefficients = null; _lastPreSoftmaxScores = null; _weightsGradient = null; diff --git a/src/NeuralNetworks/Layers/GraphConvolutionalLayer.cs b/src/NeuralNetworks/Layers/GraphConvolutionalLayer.cs index dd031fc995..8fa75130c4 100644 --- a/src/NeuralNetworks/Layers/GraphConvolutionalLayer.cs +++ b/src/NeuralNetworks/Layers/GraphConvolutionalLayer.cs @@ -177,12 +177,6 @@ public partial class GraphConvolutionalLayer : LayerBase, IAuxiliaryLossLa private Tensor _bias; - /// - /// Stores the input tensor from the last forward pass for use in the backward pass. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Stores the original input shape for any-rank tensor support. /// @@ -209,12 +203,6 @@ public partial class GraphConvolutionalLayer : LayerBase, IAuxiliaryLossLa [AiDotNet.Attributes.FittedParameter(InputSized = true)] private Tensor? _adjacencyMatrix; - /// - /// Cached reshaped adjacency matrix (3D) for backward pass. - /// - [AiDotNet.Attributes.Scratch] - private Tensor? _adjForBatch; - /// /// Edge source node indices for sparse graph representation. /// When set via SetEdges, enables memory-efficient scatter-based aggregation. @@ -821,7 +809,6 @@ protected override Tensor ForwardTraced(Tensor input) processInput = Engine.Reshape(input, [flatBatch, input.Shape[rank - 2], input.Shape[rank - 1]]); } - _lastInput = processInput; int numNodes = processInput.Shape[1]; int outputFeatures = _weights.Shape[1]; @@ -855,9 +842,6 @@ protected override Tensor ForwardTraced(Tensor input) // Set the aggregated results in output tensor for this batch output.SetSlice(b, aggregated); } - - // Store for backward pass (null for sparse path) - _adjForBatch = null; } else { @@ -892,9 +876,6 @@ protected override Tensor ForwardTraced(Tensor input) adjForBatch = _adjacencyMatrix; } - // Store for backward pass - _adjForBatch = adjForBatch; - // Then: A * (X * W) using batched matmul for 3D @ 3D output = Engine.BatchMatMul(adjForBatch, xw); } @@ -1187,9 +1168,7 @@ public T ComputeAuxiliaryLoss() public override void ResetState() { // Clear cached values from forward and backward passes - _lastInput = null; _lastNodeFeatures = null; - _adjForBatch = null; _weightsGradient = null; _biasGradient = null; } diff --git a/src/NeuralNetworks/Layers/GraphIsomorphismLayer.cs b/src/NeuralNetworks/Layers/GraphIsomorphismLayer.cs index c4e0e764ed..c10d490e01 100644 --- a/src/NeuralNetworks/Layers/GraphIsomorphismLayer.cs +++ b/src/NeuralNetworks/Layers/GraphIsomorphismLayer.cs @@ -122,23 +122,11 @@ public partial class GraphIsomorphismLayer : LayerBase, IGraphConvolutionL [AiDotNet.Attributes.FittedParameter(InputSized = true)] private Tensor? _adjacencyMatrix; - /// - /// Cached input from forward pass. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Stores the original input shape for any-rank tensor support. /// private int[]? _originalInputShape; - /// - /// Cached output from forward pass. - /// - [Scratch] - private Tensor? _lastOutput; - /// /// Cached aggregated features (before MLP). /// @@ -157,12 +145,6 @@ public partial class GraphIsomorphismLayer : LayerBase, IGraphConvolutionL [Scratch] private Tensor? _lastMlpHidden; - /// - /// Cached neighbor sum before applying epsilon. - /// - [Scratch] - private Tensor? _lastNeighborSum; - /// /// Gradients for epsilon. /// @@ -298,8 +280,6 @@ protected override Tensor ForwardTraced(Tensor input) processInput = input; } - _lastInput = processInput; - // Reshape adjacency matrix to match batch dimension if needed Tensor adjForBatch; if (_adjacencyMatrix.Shape.Length == 2 && batchSize == 1) @@ -328,37 +308,33 @@ protected override Tensor ForwardTraced(Tensor input) } // Step 1: Aggregate neighbor features using sum (batched matmul: [batch, nodes, nodes] @ [batch, nodes, features]) - _lastNeighborSum = Engine.BatchMatMul(adjForBatch, processInput); + var neighborSum = Engine.BatchMatMul(adjForBatch, processInput); // Step 2: Combine with self features: (1 + epsilon) * h_v + sum(h_u) T onePlusEpsilon = NumOps.Add(NumOps.One, _epsilon); var scaledSelf = Engine.TensorMultiplyScalar(processInput, onePlusEpsilon); - _lastAggregated = Engine.TensorAdd(scaledSelf, _lastNeighborSum); + var aggregated = Engine.TensorAdd(scaledSelf, neighborSum); + SaveForBackward(ref _lastAggregated, aggregated); // Step 3: Apply MLP - First layer (3D @ 2D requires reshape pattern) - var hidden1 = BatchedMatMul3Dx2D(_lastAggregated, _mlpWeights1, batchSize, numNodes, _inputFeatures, _mlpHiddenDim); + var hidden1 = BatchedMatMul3Dx2D(aggregated, _mlpWeights1, batchSize, numNodes, _inputFeatures, _mlpHiddenDim); var bias1Broadcast = BroadcastBias(_mlpBias1, batchSize, numNodes); - _lastMlpHiddenPreRelu = Engine.TensorAdd(hidden1, bias1Broadcast); - var lastMlpHiddenPreRelu = _lastMlpHiddenPreRelu; + var lastMlpHiddenPreRelu = Engine.TensorAdd(hidden1, bias1Broadcast); + SaveForBackward(ref _lastMlpHiddenPreRelu, lastMlpHiddenPreRelu); // Apply ReLU activation using Engine operations var zeroTensor = new Tensor(lastMlpHiddenPreRelu._shape); zeroTensor.Fill(NumOps.Zero); - _lastMlpHidden = Engine.TensorMax(lastMlpHiddenPreRelu, zeroTensor); + var mlpHidden = Engine.TensorMax(lastMlpHiddenPreRelu, zeroTensor); + SaveForBackward(ref _lastMlpHidden, mlpHidden); // Step 4: Apply MLP - Second layer (3D @ 2D requires reshape pattern) - var hidden2 = BatchedMatMul3Dx2D(_lastMlpHidden, _mlpWeights2, batchSize, numNodes, _mlpHiddenDim, _outputFeatures); + var hidden2 = BatchedMatMul3Dx2D(mlpHidden, _mlpWeights2, batchSize, numNodes, _mlpHiddenDim, _outputFeatures); var bias2Broadcast = BroadcastBias(_mlpBias2, batchSize, numNodes); var mlpOutput = Engine.TensorAdd(hidden2, bias2Broadcast); var result = ApplyActivation(mlpOutput); - // Only store for backward pass during training - skip during inference - if (IsTrainingMode) - { - _lastOutput = result; - } - // Reshape output to match original input shape (except for feature dimension) if (_originalInputShape != null && _originalInputShape.Length == 2) { @@ -489,12 +465,9 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; _lastAggregated = null; _lastMlpHiddenPreRelu = null; _lastMlpHidden = null; - _lastNeighborSum = null; _mlpWeights1Gradient = null; _mlpWeights2Gradient = null; _mlpBias1Gradient = null; diff --git a/src/NeuralNetworks/Layers/GraphSAGELayer.cs b/src/NeuralNetworks/Layers/GraphSAGELayer.cs index 827dbaa844..f6719d9015 100644 --- a/src/NeuralNetworks/Layers/GraphSAGELayer.cs +++ b/src/NeuralNetworks/Layers/GraphSAGELayer.cs @@ -120,24 +120,12 @@ public partial class GraphSAGELayer : LayerBase, IGraphConvolutionLayer /// private int[]? _originalInputShape; - /// - /// Cached output from forward pass. - /// - [Scratch] - private Tensor? _lastOutput; - /// /// Cached aggregated neighbor features. /// [Scratch] private Tensor? _lastAggregated; - /// - /// Cached pre-normalization output for gradient computation. - /// - [Scratch] - private Tensor? _lastPreNorm; - /// /// Cached degrees for each node. /// @@ -294,7 +282,7 @@ protected override Tensor ForwardTraced(Tensor input) processInput = Engine.Reshape(input, [1, 1, input.Shape[0]]); } - _lastInput = processInput; + SaveForBackward(ref _lastInput, processInput); int numNodes = processInput.Shape[1]; // Ensure adjacency matrix has matching rank for batch operation @@ -319,50 +307,45 @@ protected override Tensor ForwardTraced(Tensor input) adjForBatch = _adjacencyMatrix; } - // Store for backward pass - _adjForBatch = adjForBatch; + SaveForBackward(ref _adjForBatch, adjForBatch); // Step 1: Compute degrees for normalization - _lastDegrees = Engine.ReduceSum(adjForBatch, [2], keepDims: false); // [batch, nodes] + var degrees = Engine.ReduceSum(adjForBatch, [2], keepDims: false); // [batch, nodes] + SaveForBackward(ref _lastDegrees, degrees); // Clamp degrees to minimum of 1 to avoid division by zero - var oneTensor = new Tensor(_lastDegrees._shape); + var oneTensor = new Tensor(degrees._shape); oneTensor.Fill(NumOps.One); - var safeDegrees = Engine.TensorMax(_lastDegrees, oneTensor); + var safeDegrees = Engine.TensorMax(degrees, oneTensor); // Step 2: Aggregate neighbor features using vectorized operations - _lastAggregated = ComputeVectorizedAggregation(processInput, safeDegrees, batchSize, numNodes, adjForBatch); + var aggregated = ComputeVectorizedAggregation(processInput, safeDegrees, batchSize, numNodes, adjForBatch); + SaveForBackward(ref _lastAggregated, aggregated); // Step 3: Transform self features (3D @ 2D requires reshape pattern) var selfTransformed = BatchedMatMul3Dx2D(processInput, _selfWeights, batchSize, numNodes, _inputFeatures, _outputFeatures); // Step 4: Transform neighbor features (3D @ 2D requires reshape pattern) - var neighborTransformed = BatchedMatMul3Dx2D(_lastAggregated, _neighborWeights, batchSize, numNodes, _inputFeatures, _outputFeatures); + var neighborTransformed = BatchedMatMul3Dx2D(aggregated, _neighborWeights, batchSize, numNodes, _inputFeatures, _outputFeatures); // Step 5: Combine: self + neighbor + bias var combined = Engine.TensorAdd(selfTransformed, neighborTransformed); var biasBroadcast = BroadcastBias(_bias, batchSize, numNodes); - _lastPreNorm = Engine.TensorAdd(combined, biasBroadcast); + var preNorm = Engine.TensorAdd(combined, biasBroadcast); // Step 6: Apply L2 normalization if enabled Tensor output; if (_normalize) { - output = L2NormalizeVectorized(_lastPreNorm, batchSize, numNodes); + output = L2NormalizeVectorized(preNorm, batchSize, numNodes); } else { - output = _lastPreNorm; + output = preNorm; } var result = ApplyActivation(output); - // Only store for backward pass during training - skip during inference - if (IsTrainingMode) - { - _lastOutput = result; - } - // Restore original shape for any-rank tensor support if (_originalInputShape != null && _originalInputShape.Length != 3) { @@ -657,9 +640,7 @@ public override void ClearGradients() public override void ResetState() { _lastInput = null; - _lastOutput = null; _lastAggregated = null; - _lastPreNorm = null; _lastDegrees = null; _lastMaxIndices = null; _lastMaxInputShape = null; diff --git a/src/NeuralNetworks/Layers/GraphTransformerLayer.cs b/src/NeuralNetworks/Layers/GraphTransformerLayer.cs index 8aa0528cbb..f4ad0e793c 100644 --- a/src/NeuralNetworks/Layers/GraphTransformerLayer.cs +++ b/src/NeuralNetworks/Layers/GraphTransformerLayer.cs @@ -205,12 +205,6 @@ public partial class GraphTransformerLayer : LayerBase, IGraphConvolutionL [AiDotNet.Attributes.FittedParameter(InputSized = true)] private Tensor? _adjacencyMatrix; - /// - /// Cached values for backward pass. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Stores the original input shape for any-rank tensor support. /// @@ -232,11 +226,7 @@ public partial class GraphTransformerLayer : LayerBase, IGraphConvolutionL [Scratch] private Tensor? _lastAttnOutput; [Scratch] - private Tensor? _lastNormed1; - [Scratch] private Tensor? _lastFFNHidden; - [Scratch] - private Tensor? _lastFFNOutput; /// /// Gradients for parameters. @@ -550,7 +540,6 @@ protected override Tensor ForwardTraced(Tensor input) } bool cacheBwd = ShouldCacheForBackward; // #1668: gate backward-only caches (arena safety) - _lastInput = cacheBwd ? processInput : null; int numNodes = processInput.Shape[1]; // Multi-head attention block with residual connection @@ -579,11 +568,9 @@ protected override Tensor ForwardTraced(Tensor input) } var normed1 = Engine.TensorAdd(attended, residualInput); - _lastNormed1 = cacheBwd ? normed1 : null; // Feed-forward network with residual var ffnOutput = FeedForwardNetwork(normed1, batchSize, numNodes); - _lastFFNOutput = cacheBwd ? ffnOutput : null; // Final residual and layer norm var output = Engine.TensorAdd(normed1, ffnOutput); @@ -902,7 +889,6 @@ public override void ClearGradients() public override void ResetState() { - _lastInput = null; _lastOutput = null; _lastQueries = null; _lastKeys = null; @@ -911,9 +897,7 @@ public override void ResetState() _lastHeadOutputs = null; _lastConcatenated = null; _lastAttnOutput = null; - _lastNormed1 = null; _lastFFNHidden = null; - _lastFFNOutput = null; _queryWeightsGradient = null; _keyWeightsGradient = null; _valueWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/GroupNormalizationLayer.cs b/src/NeuralNetworks/Layers/GroupNormalizationLayer.cs index c1d6c7581c..9522e57c20 100644 --- a/src/NeuralNetworks/Layers/GroupNormalizationLayer.cs +++ b/src/NeuralNetworks/Layers/GroupNormalizationLayer.cs @@ -56,12 +56,6 @@ public partial class GroupNormalizationLayer : LayerBase private Tensor _gamma; private Tensor _beta; [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastMean; - [Scratch] - private Tensor? _lastVariance; - [Scratch] private Tensor? _gammaGradient; [Scratch] private Tensor? _betaGradient; @@ -169,11 +163,6 @@ public GroupNormalizationLayer(int numGroups, int numChannels, double epsilon = protected override Tensor ForwardTraced(Tensor input) { - // Retain backward-activation caches (_lastInput / _lastMean / _lastVariance) - // only when an eager manual Backward will read them; skip in inference/under - // tape so the denoise-loop arena can recycle scratch without aliasing (#1668). - bool cacheBwd = ShouldCacheForBackward; - _lastInput = cacheBwd ? input : null; _originalInputShape = input._shape; var shape = input._shape; @@ -234,11 +223,8 @@ protected override Tensor ForwardTraced(Tensor input) _gamma, _beta, NumOps.ToDouble(_epsilon), - out var mean, - out var variance); - - _lastMean = cacheBwd ? mean : null; - _lastVariance = cacheBwd ? variance : null; + out _, + out _); // Restore original tensor rank — via Engine so tape records the restore. if (_originalInputShape.Length > 4) @@ -276,9 +262,6 @@ public override void UpdateParameters(T learningRate) public override void ResetState() { - _lastInput = null; - _lastMean = null; - _lastVariance = null; _gammaGradient = null; _betaGradient = null; _addedBatchDimension = false; diff --git a/src/NeuralNetworks/Layers/GroupedQueryAttentionLayer.cs b/src/NeuralNetworks/Layers/GroupedQueryAttentionLayer.cs index 8564017e5a..1e5b1c11dc 100644 --- a/src/NeuralNetworks/Layers/GroupedQueryAttentionLayer.cs +++ b/src/NeuralNetworks/Layers/GroupedQueryAttentionLayer.cs @@ -132,10 +132,6 @@ public partial class GroupedQueryAttentionLayer : LayerBase, IShapeContrac // Cached values for backward pass [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] private Tensor? _lastProjectedQueries; [Scratch] private Tensor? _lastProjectedKeys; @@ -147,8 +143,6 @@ public partial class GroupedQueryAttentionLayer : LayerBase, IShapeContrac private Tensor? _lastExpandedValues; [Scratch] private Tensor? _lastAttentionWeights; - [Scratch] - private Tensor? _lastAttentionContext; private int[]? _originalInputShape; // Gradients @@ -476,7 +470,6 @@ protected override Tensor ForwardTraced(Tensor input) : Engine.Reshape(input, new[] { batchSize, seqLen, embDim }); bool cacheBwd = ShouldCacheForBackward; // #1668: gate all backward caches (arena safety) - _lastInput = cacheBwd ? input3D : null; // Project Q, K, V var input2D = Engine.Reshape(input3D, new[] { batchSize * seqLen, embDim }); @@ -582,11 +575,6 @@ protected override Tensor ForwardTraced(Tensor input) contextPermuted, new[] { batchSize * seqLen, _numHeads * _headDimension }); - // Cache pre-projection context for output weights gradient - _lastAttentionContext = cacheBwd - ? Engine.Reshape(contextTransposed, new[] { batchSize, seqLen, _numHeads * _headDimension }) - : null; - // Output projection var output = Engine.TensorMatMul(contextTransposed, _outputWeights); var output3D = Engine.Reshape(output, new[] { batchSize, seqLen, _embeddingDimension }); @@ -598,8 +586,6 @@ protected override Tensor ForwardTraced(Tensor input) var outputWithBias = Engine.TensorAdd(output3D, biasBroadcast); var result = ApplyActivation(outputWithBias); - _lastOutput = cacheBwd ? result : null; - // Reshape back to original rank — via Engine for tape recording. if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _embeddingDimension }); @@ -790,15 +776,12 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; _lastProjectedQueries = null; _lastProjectedKeys = null; _lastProjectedValues = null; _lastExpandedKeys = null; _lastExpandedValues = null; _lastAttentionWeights = null; - _lastAttentionContext = null; _originalInputShape = null; _queryWeightsGradient = null; _keyWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/HeterogeneousGraphLayer.cs b/src/NeuralNetworks/Layers/HeterogeneousGraphLayer.cs index f5a1d1a769..cf4b6ef741 100644 --- a/src/NeuralNetworks/Layers/HeterogeneousGraphLayer.cs +++ b/src/NeuralNetworks/Layers/HeterogeneousGraphLayer.cs @@ -227,18 +227,10 @@ public partial class HeterogeneousGraphLayer : LayerBase, IGraphConvolutio /// private Dictionary? _nodeTypeMap; - /// - /// Cached values for backward pass. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Stores the original input shape for any-rank tensor support. /// private int[]? _originalInputShape; - [Scratch] - private Tensor? _lastOutput; /// /// Gradients for weights, self-loop weights, and biases. @@ -487,7 +479,6 @@ protected override Tensor ForwardTraced(Tensor input) processInput = Engine.Reshape(input, [flatBatch, numNodes, features]); } - _lastInput = processInput; int processNumNodes = processInput.Shape[1]; var output = TensorAllocator.Rent([batchSize, processNumNodes, _outputFeatures]); @@ -591,12 +582,6 @@ protected override Tensor ForwardTraced(Tensor input) } } - // Store output AFTER shape restoration so Backward gets matching rank - if (IsTrainingMode) - { - _lastOutput = result; - } - return result; } @@ -1024,8 +1009,6 @@ public void SetParameterTensors(List> parameters) /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; _edgeTypeWeightsGradients = null; _selfLoopWeightsGradients = null; _biasesGradients = null; diff --git a/src/NeuralNetworks/Layers/HighwayLayer.cs b/src/NeuralNetworks/Layers/HighwayLayer.cs index 99ee5a6164..2861624fc3 100644 --- a/src/NeuralNetworks/Layers/HighwayLayer.cs +++ b/src/NeuralNetworks/Layers/HighwayLayer.cs @@ -156,42 +156,12 @@ public partial class HighwayLayer : LayerBase, IAuxiliaryLossLayer, ISh /// private Tensor _gateBias; - /// - /// Stores the input tensor from the last forward pass for use in the backward pass. - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// Stores the output tensor from the last forward pass for use in the backward pass. - /// - [Scratch] - private Tensor? _lastOutput; - - /// - /// Stores the transformed output tensor from the last forward pass for use in the backward pass. - /// - [Scratch] - private Tensor? _lastTransformOutput; - /// /// Stores the gate output tensor from the last forward pass for use in the backward pass. /// [Scratch] private Tensor? _lastGateOutput; - /// - /// Stores the pre-activation transform values from the last forward pass. - /// - [Scratch] - private Tensor? _lastTransformPreActivation; - - /// - /// Stores the pre-activation gate values from the last forward pass. - /// - [Scratch] - private Tensor? _lastGatePreActivation; - /// /// Stores the gradients for the transform weights calculated during the backward pass. /// @@ -516,24 +486,17 @@ private void InitializeTensor(Tensor tensor, T scale) /// protected override Tensor ForwardTraced(Tensor input) { - // #1668: gate every backward-only cache (not just _lastInput) so an InferenceMode - // arena loop holds no recycled-scratch references. Forward uses locals throughout. - bool cacheBwd = ShouldCacheForBackward; - _lastInput = cacheBwd ? input : null; - // Transform path: transform = activation(input @ weights + bias) var transformLinear = Engine.TensorMatMul(input, _transformWeights); var transformWithBias = Engine.TensorAdd(transformLinear, _transformBias); - _lastTransformPreActivation = cacheBwd ? transformWithBias : null; // Store pre-activation for backward pass var transformOutput = ApplyActivation(transformWithBias, _transformActivation, _vectorTransformActivation); - _lastTransformOutput = cacheBwd ? transformOutput : null; // Gate path: gate = sigmoid(input @ weights + bias) var gateLinear = Engine.TensorMatMul(input, _gateWeights); var gateWithBias = Engine.TensorAdd(gateLinear, _gateBias); - _lastGatePreActivation = cacheBwd ? gateWithBias : null; // Store pre-activation for backward pass var gateOutput = ApplyActivation(gateWithBias, _gateActivation, _vectorGateActivation); - _lastGateOutput = cacheBwd ? gateOutput : null; + // The gate values feed ComputeAuxiliaryLoss; nothing else outlives the forward. + SaveForBackward(ref _lastGateOutput, gateOutput); // Highway output: output = gate * transform + (1 - gate) * input // Rewritten as: output = gate * (transform - input) + input @@ -541,7 +504,6 @@ protected override Tensor ForwardTraced(Tensor input) var gatedDiff = Engine.TensorMultiply(gateOutput, transformMinusInput); var output = Engine.TensorAdd(gatedDiff, input); - _lastOutput = cacheBwd ? output : null; return output; } @@ -687,12 +649,7 @@ public override void ClearGradients() public override void ResetState() { // Clear cached values from forward and backward passes - _lastInput = null; - _lastOutput = null; - _lastTransformOutput = null; _lastGateOutput = null; - _lastTransformPreActivation = null; - _lastGatePreActivation = null; _transformWeightsGradient = null; _transformBiasGradient = null; _gateWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/HyperbolicLinearLayer.cs b/src/NeuralNetworks/Layers/HyperbolicLinearLayer.cs index 3e1f0b2b49..ac6b5bb930 100644 --- a/src/NeuralNetworks/Layers/HyperbolicLinearLayer.cs +++ b/src/NeuralNetworks/Layers/HyperbolicLinearLayer.cs @@ -84,23 +84,11 @@ public partial class HyperbolicLinearLayer : LayerBase, IShapeContract /// private readonly T _curvature; - /// - /// Stored input from forward pass for backpropagation. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Stores the original input shape for any-rank tensor support. /// private int[]? _originalInputShape; - /// - /// Stored pre-activation output for gradient computation. - /// - [Scratch] - private Tensor? _lastOutput; - /// /// Gradient for weights, stored during backward pass. /// @@ -258,8 +246,6 @@ protected override Tensor ForwardTraced(Tensor input) $"Input size {inputLen} does not match expected {InputFeatures}."); } - _lastInput = inputTensor; - // Batched Möbius matrix-vector multiply per Ganea et al. 2018, Section 2.3. // M⊗_c x = tanh(||Mx|| * arctanh(√c·||x||) / (√c·||x||)) * Mx / (√c·||Mx||) // All operations are tape-tracked tensor ops for automatic differentiation. @@ -275,7 +261,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.TensorTranspose(_weights) : (_weightsTCache ??= Engine.TensorTranspose(_weights)); var linear = Engine.FusedLinear(inputTensor, euclideanW, _biases, FusedActivationType.None); - _lastOutput = linear; var euclideanOut = ApplyActivation(linear); if (_originalInputShape == null || _originalInputShape.Length == 2) return euclideanOut; if (_originalInputShape.Length == 1) return Engine.Reshape(euclideanOut, [OutputFeatures]); @@ -376,8 +361,6 @@ protected override Tensor ForwardTraced(Tensor input) projScale = Engine.TensorClamp(projScale, NumOps.Zero, NumOps.One); // clamp to [0, 1] output = Engine.TensorMultiply(projScale, output); // [batch, outputFeatures] - _lastOutput = output; - // Apply activation function var activated = ApplyActivation(output); @@ -630,8 +613,6 @@ public override void ClearGradients() public override void ResetState() { - _lastInput = null; - _lastOutput = null; _weightsGradient = null; _biasesGradient = null; } diff --git a/src/NeuralNetworks/Layers/InstanceNormalizationLayer.cs b/src/NeuralNetworks/Layers/InstanceNormalizationLayer.cs index 7664b24b26..6946835575 100644 --- a/src/NeuralNetworks/Layers/InstanceNormalizationLayer.cs +++ b/src/NeuralNetworks/Layers/InstanceNormalizationLayer.cs @@ -82,11 +82,6 @@ public partial class InstanceNormalizationLayer : LayerBase, IShapeContrac [TrainableParameter(Role = PersistentTensorRole.NormalizationParams, Shape = "_numChannels", Condition = nameof(Affine))] private Tensor _beta; - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastMean; - [Scratch] - private Tensor? _lastVariance; [Scratch] private Tensor? _gammaGradient; [Scratch] @@ -249,8 +244,6 @@ protected override Tensor ForwardTraced(Tensor input) if (channels != _numChannels) throw new ArgumentException($"Input has {channels} channels but layer expects {_numChannels} channels."); - _lastInput = flattenedInput; - // Instance Normalization is Group Normalization with numGroups = numChannels // This means each channel is normalized independently var output = Engine.GroupNorm( @@ -259,11 +252,8 @@ protected override Tensor ForwardTraced(Tensor input) _gamma, _beta, NumOps.ToDouble(_epsilon), - out var mean, - out var variance); - - _lastMean = mean; - _lastVariance = variance; + out _, + out _); // Reshape output back to original shape return Engine.Reshape(output, _originalInputShape); @@ -296,9 +286,6 @@ public override void UpdateParameters(T learningRate) public override void ResetState() { - _lastInput = null; - _lastMean = null; - _lastVariance = null; _gammaGradient = null; _betaGradient = null; _gpuMean?.Dispose(); diff --git a/src/NeuralNetworks/Layers/InteractingLayer.cs b/src/NeuralNetworks/Layers/InteractingLayer.cs index a350b6a625..d573928651 100644 --- a/src/NeuralNetworks/Layers/InteractingLayer.cs +++ b/src/NeuralNetworks/Layers/InteractingLayer.cs @@ -83,19 +83,7 @@ public partial class InteractingLayer : LayerBase, IShapeContract // Cached values [Scratch] - private Tensor? _inputCache; - [Scratch] - private Tensor? _queriesCache; - [Scratch] - private Tensor? _keysCache; - [Scratch] - private Tensor? _valuesCache; - [Scratch] private Tensor? _attentionScoresCache; - [Scratch] - private Tensor? _attendedCache; - [Scratch] - private Tensor? _preActivationCache; /// public override bool SupportsTraining => true; @@ -190,8 +178,6 @@ private void InitializeWeights(double scale) /// Feature interactions [batchSize, numFeatures, embeddingDim]. protected override Tensor ForwardTraced(Tensor input) { - _inputCache = input; - int batchSize = input.Shape[0]; int numFeatures = input.Shape[1]; int embDim = input.Shape[2]; @@ -201,13 +187,8 @@ protected override Tensor ForwardTraced(Tensor input) var keys = ProjectInput(input, _keyWeights, batchSize, numFeatures, embDim, _attentionDim); var values = ProjectInput(input, _valueWeights, batchSize, numFeatures, embDim, _attentionDim); - _queriesCache = queries; - _keysCache = keys; - _valuesCache = values; - // Multi-head self-attention using Engine ops var attended = MultiHeadAttention(queries, keys, values, batchSize, numFeatures); - _attendedCache = attended; // Output projection using Engine.TensorMatMul var output = ProjectOutput(attended, batchSize, numFeatures); @@ -218,8 +199,6 @@ protected override Tensor ForwardTraced(Tensor input) output = AddResidual(input, output, batchSize, numFeatures, embDim); } - _preActivationCache = output; - // Apply ReLU activation using Engine output = Engine.ReLU(output); @@ -360,13 +339,7 @@ private static void CopyTensorToVector(Tensor tensor, Vector vector, ref i /// public override void ResetState() { - _inputCache = null; - _queriesCache = null; - _keysCache = null; - _valuesCache = null; _attentionScoresCache = null; - _attendedCache = null; - _preActivationCache = null; Engine.TensorFill(_queryWeightsGrad, NumOps.Zero); Engine.TensorFill(_keyWeightsGrad, NumOps.Zero); diff --git a/src/NeuralNetworks/Layers/IntersampleAttentionLayer.cs b/src/NeuralNetworks/Layers/IntersampleAttentionLayer.cs index c9f2906969..6c6b98e169 100644 --- a/src/NeuralNetworks/Layers/IntersampleAttentionLayer.cs +++ b/src/NeuralNetworks/Layers/IntersampleAttentionLayer.cs @@ -71,10 +71,6 @@ public partial class IntersampleAttentionLayer : LayerBase, IShapeContract private Tensor _layerNormBeta; // Cached values for backward pass - [Scratch] - private Tensor? _inputCache; - [Scratch] - private Tensor? _normalizedCache; /// public override bool SupportsTraining => true; @@ -119,8 +115,6 @@ public IntersampleAttentionLayer( // real forward pass). The probe input is discarded. var probe = new Tensor([1, 1, embeddingDim]); Forward(probe); - _inputCache = null; - _normalizedCache = null; } /// @@ -134,8 +128,6 @@ protected override Tensor ForwardTraced(Tensor input) // EnsureSubLayersRegistered(); they otherwise stay invisible to every structural walker. EnsureInitializedFromInput(input); - _inputCache = input; - int batchSize = input.Shape[0]; int numFeatures = input.Shape[1]; int embDim = input.Shape[2]; @@ -171,7 +163,6 @@ protected override Tensor ForwardTraced(Tensor input) // Residual connection and layer normalization over the embedding axis. var output = AddResidualAndNormalize(input, projected, batchSize, numFeatures, embDim); - _normalizedCache = output; return output; } @@ -228,8 +219,6 @@ internal override Dictionary GetMetadata() /// public override void ResetState() { - _inputCache = null; - _normalizedCache = null; _queryProjection.ResetState(); _keyProjection.ResetState(); diff --git a/src/NeuralNetworks/Layers/LSTMLayer.cs b/src/NeuralNetworks/Layers/LSTMLayer.cs index 96c5522264..64244bece2 100644 --- a/src/NeuralNetworks/Layers/LSTMLayer.cs +++ b/src/NeuralNetworks/Layers/LSTMLayer.cs @@ -412,28 +412,6 @@ public partial class LSTMLayer : LayerBase, IShapeContract [TrainableParameter(Role = PersistentTensorRole.Biases, Shape = "_hiddenSize")] private Tensor _biasO; - /// - /// The input tensor from the last forward pass. - /// - /// - /// - /// This field stores the input tensor from the most recent forward pass. It is needed during the backward pass - /// to compute gradients correctly. This field is reset when ResetState is called. - /// - /// For Beginners: This stores the most recent data that was fed into the layer. - /// - /// The layer needs to remember the input: - /// - To calculate how each input value affected the output - /// - To determine how to propagate gradients during training - /// - To ensure the backward pass works correctly - /// - /// This is like keeping track of what you were given so you can explain - /// how you processed it if asked. - /// - /// - [Scratch] - private Tensor? _lastInput; - /// /// The hidden state from the last forward pass. /// @@ -499,17 +477,6 @@ internal Tensor ForwardFromState( _stateStepping = false; } } - /// - /// Cached hidden states for all time steps (Batch, Time, Hidden). - /// - [Scratch] - private Tensor? _cachedHiddenStates; - - /// - /// Cached cell states for all time steps (Batch, Time, Hidden). - /// - [Scratch] - private Tensor? _cachedCellStates; /// /// Stores the original input shape for any-rank tensor support. @@ -1263,12 +1230,6 @@ protected override Tensor ForwardTraced(Tensor input) input3D = Engine.Reshape(input, [flatBatch, timeSteps, _inputSize]); } - // #1668: _lastInput is a manual-backward cache; gate it (and the per-timestep BPTT - // caches below) on ShouldCacheForBackward so inference/tape/eval keeps no arena-backed - // activation alive across a TensorArena.Reset(). - bool cacheBwd = ShouldCacheForBackward; - _lastInput = cacheBwd ? input3D : null; - // The engine owns CPU/GPU dispatch. Both float and double have fused BPTT; restricting // this shared layer to float left every double-valued model on the per-step graph. // The state-returning overload dispatches sequence and final-state gradients together. @@ -1316,12 +1277,6 @@ protected override Tensor ForwardTraced(Tensor input) // when the sequence primitive is unavailable, for other numeric types, or when the engine declines the primitive. var hiddenStatesList = new System.Collections.Generic.List>(timeSteps); - // #1668: the per-timestep BPTT caches are read only by the manual Backward; skip - // renting+populating them when no eager Backward will run (cacheBwd computed above), - // so an arena loop holds no per-step references across a Reset. - _cachedHiddenStates = cacheBwd ? TensorAllocator.Rent(new int[] { batchSize, timeSteps, _hiddenSize }) : null; - _cachedCellStates = cacheBwd ? TensorAllocator.Rent(new int[] { batchSize, timeSteps, _hiddenSize }) : null; - var currentH = TensorAllocator.Rent(new int[] { batchSize, _hiddenSize }); var currentC = TensorAllocator.Rent(new int[] { batchSize, _hiddenSize }); // TensorAllocator.Rent returns POOLED memory that is not zero-initialized. @@ -1364,17 +1319,16 @@ Tensor Recurrent(Tensor x) => _useVectorActivation : LstmGateDispatch.Forward(Engine, gates, currentC, Activate, Recurrent, out currentC); // Collect the hidden state for the tape-connected output concat below. hiddenStatesList.Add(Engine.Reshape(currentH, new[] { batchSize, 1, _hiddenSize })); - // Caches for the manual backward path (not on the tape). - _cachedHiddenStates?.SetSlice(1, t, currentH); - _cachedCellStates?.SetSlice(1, t, currentC); } // Assemble the [batch, timeSteps, hiddenSize] output on the tape so gradients // flow back through the recurrence to the gate weights. var output = Engine.TensorConcatenate(hiddenStatesList.ToArray(), axis: 1); - _lastHiddenState = currentH; - _lastCellState = currentC; + // The final state is only read back by ForwardFromState. Outside a stepped call it is + // cleared, as the fused path above does, so it never pins this pass's buffers. + _lastHiddenState = _stateStepping ? currentH : null; + _lastCellState = _stateStepping ? currentC : null; // Restore original batch dimensions for any-rank support if (_originalInputShape != null && _originalInputShape.Length > 3) @@ -2010,14 +1964,9 @@ public override void SetTrainingMode(bool isTraining) public override void ResetState() { // Clear cached values from forward and backward passes - _lastInput = null; _lastHiddenState = null; _lastCellState = null; - // Clear per-time-step cached states to prevent stale state leakage between sequences - _cachedHiddenStates = null; - _cachedCellStates = null; - Gradients.Clear(); } diff --git a/src/NeuralNetworks/Layers/LambdaLayer.cs b/src/NeuralNetworks/Layers/LambdaLayer.cs index 4b5a025bf6..b22c01b590 100644 --- a/src/NeuralNetworks/Layers/LambdaLayer.cs +++ b/src/NeuralNetworks/Layers/LambdaLayer.cs @@ -175,18 +175,6 @@ public partial class LambdaLayer : LayerBase, IShapeContract - /// - /// Stores the input tensor from the last forward pass for use in the backward pass. - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// Stores the output tensor from the last forward pass for use in the backward pass. - /// - [Scratch] - private Tensor? _lastOutput; - /// /// Gets a value indicating whether this layer supports training. /// @@ -388,12 +376,10 @@ public LambdaLayer(int[] inputShape, int[] outputShape, /// protected override Tensor ForwardTraced(Tensor input) { - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) var output = _forwardFunction(input); output = ApplyActivation(output); - _lastOutput = ShouldCacheForBackward ? output : null; // #1668: skip in inference (arena safety) return output; } @@ -421,8 +407,6 @@ protected override Tensor ForwardTraced(Tensor input) public override void ResetState() { // Clear cached values from forward pass - _lastInput = null; - _lastOutput = null; } } diff --git a/src/NeuralNetworks/Layers/LocallyConnectedLayer.cs b/src/NeuralNetworks/Layers/LocallyConnectedLayer.cs index 8573cb72a5..299c511f12 100644 --- a/src/NeuralNetworks/Layers/LocallyConnectedLayer.cs +++ b/src/NeuralNetworks/Layers/LocallyConnectedLayer.cs @@ -158,29 +158,11 @@ public partial class LocallyConnectedLayer : LayerBase, IShapeContract private Tensor _biases; - /// - /// Stores the input tensor from the last forward pass for use in the backward pass. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Stores the original input shape for any-rank tensor support. /// private int[]? _originalInputShape; - /// - /// Stores the pre-activation output from the last forward pass for use in the backward pass. - /// - [Scratch] - private Tensor? _lastPreActivation; - - /// - /// Stores the output tensor from the last forward pass for use in the backward pass. - /// - [Scratch] - private Tensor? _lastOutput; - /// /// Stores the gradients for the weights calculated during the backward pass. /// @@ -642,8 +624,6 @@ protected override Tensor ForwardTraced(Tensor input) processInput = Engine.Reshape(input, [flatBatch, height, width, channels]); } - _lastInput = processInput; - // === GPU-Accelerated LocallyConnectedConv2D === // The layer uses NHWC format but Engine expects NCHW, so we transpose. // Every shape op via Engine so the gradient tape records the transformation — @@ -665,13 +645,6 @@ protected override Tensor ForwardTraced(Tensor input) // Apply activation function var result = ApplyActivation(preActivation); - // Only store for backward pass during training - skip during inference - if (IsTrainingMode) - { - _lastPreActivation = preActivation; - _lastOutput = result; - } - // Restore output shape to match original input rank if (_originalInputShape != null && _originalInputShape.Length != 4) { @@ -813,9 +786,6 @@ public override void ClearGradients() public override void ResetState() { // Clear cached values from forward and backward passes - _lastInput = null; - _lastPreActivation = null; - _lastOutput = null; _weightGradients = null; _biasGradients = null; diff --git a/src/NeuralNetworks/Layers/LogVarianceLayer.cs b/src/NeuralNetworks/Layers/LogVarianceLayer.cs index 6d0265f71f..5b6722c508 100644 --- a/src/NeuralNetworks/Layers/LogVarianceLayer.cs +++ b/src/NeuralNetworks/Layers/LogVarianceLayer.cs @@ -97,65 +97,6 @@ public partial class LogVarianceLayer : LayerBase, IShapeContract /// public override bool SupportsTraining => false; - /// - /// The input tensor from the last forward pass. - /// - /// - /// - /// This field stores the input tensor from the most recent forward pass. It is needed during the backward pass - /// to compute gradients correctly. This field is reset when ResetState is called. - /// - /// For Beginners: This stores the most recent data that was fed into the layer. - /// - /// The layer needs to remember the input: - /// - To calculate how each input value affected the output - /// - To determine how to propagate gradients during training - /// - To ensure the backward pass works correctly - /// - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// The output tensor from the last forward pass. - /// - /// - /// - /// This field stores the output tensor from the most recent forward pass. It is needed during the backward pass - /// because the derivative of the logarithm function depends on the output value. This field is reset when ResetState is called. - /// - /// For Beginners: This stores the most recent result that came out of the layer. - /// - /// The layer needs to remember its output: - /// - Because the derivative of log(x) is 1/x, so we need x (our output) during backpropagation - /// - To avoid recalculating values during the backward pass - /// - To make the training process more efficient - /// - /// - [Scratch] - private Tensor? _lastOutput; - - /// - /// The mean values calculated during the last forward pass. - /// - /// - /// - /// This field stores the mean values calculated during the most recent forward pass. These values are needed - /// during both the variance calculation and the backward pass. This field is reset when ResetState is called. - /// - /// For Beginners: This stores the average values calculated during the first step. - /// - /// The layer needs to remember these mean values: - /// - They're used when calculating the variance (which measures deviation from the mean) - /// - They're needed again during the backward pass - /// - Storing them avoids having to recalculate them multiple times - /// - /// Think of it as saving an intermediate result that will be reused later. - /// - /// - [AiDotNet.Attributes.Scratch] - private Tensor? _meanValues; - /// /// Initializes a new instance of the class. /// @@ -261,13 +202,8 @@ private static int[] CalculateOutputShape(int[] inputShape, int axis) protected override Tensor ForwardTraced(Tensor input) { EnsureInitializedFromInput(input); - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) - // Use Engine operations for GPU/CPU acceleration - _meanValues = Engine.ReduceMean(input, [Axis], keepDims: true); - _lastOutput = Engine.ReduceLogVariance(input, [Axis], keepDims: false, epsilon: 1e-8); - - return _lastOutput; + return Engine.ReduceLogVariance(input, [Axis], keepDims: false, epsilon: 1e-8); } /// @@ -315,9 +251,7 @@ internal override Dictionary GetMetadata() /// public override void ResetState() { - // Clear cached values from forward pass - _lastInput = null; - _lastOutput = null; - _meanValues = null; + // Nothing to clear: the forward keeps no per-pass state. + return; } } diff --git a/src/NeuralNetworks/Layers/MaskingLayer.cs b/src/NeuralNetworks/Layers/MaskingLayer.cs index 3e9975a947..c43fdc571f 100644 --- a/src/NeuralNetworks/Layers/MaskingLayer.cs +++ b/src/NeuralNetworks/Layers/MaskingLayer.cs @@ -64,52 +64,6 @@ public partial class MaskingLayer : LayerBase, IShapeContract /// private readonly T _maskValue; - /// - /// The input tensor from the last forward pass. - /// - /// - /// - /// This field stores the input tensor from the most recent forward pass. It is kept for potential future use, - /// though the current implementation doesn't directly use it after storing it. This field is reset when - /// ResetState is called. - /// - /// For Beginners: This stores the most recent data that was fed into the layer. - /// - /// The layer keeps track of: - /// - The input it received in the most recent forward pass - /// - This information might be needed for future operations - /// - It gets cleared when you reset the layer's state - /// - /// Storing the last input is a common pattern in neural network layers, - /// even when not immediately needed for calculations. - /// - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// The mask tensor from the last forward pass. - /// - /// - /// - /// This field stores the binary mask tensor created during the most recent forward pass. It is needed during - /// the backward pass to apply the same masking pattern to the output gradients. This field is reset when - /// ResetState is called. - /// - /// For Beginners: This stores the pattern of 0s and 1s created in the last forward pass. - /// - /// The layer remembers: - /// - Which positions were masked (set to 0) in the forward pass - /// - This same pattern must be applied during the backward pass - /// - This ensures consistency between forward and backward operations - /// - /// By keeping the mask, the layer doesn't need to recalculate it during - /// the backward pass, which saves computation time. - /// - /// - [Scratch] - private Tensor? _lastMask; - /// /// The GPU mask tensor from the last GPU forward pass (for backward pass caching). /// @@ -204,9 +158,7 @@ protected override void OnFirstForward(Tensor input) protected override Tensor ForwardTraced(Tensor input) { EnsureInitializedFromInput(input); - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) var mask = CreateMask(input); - _lastMask = ShouldCacheForBackward ? mask : null; // #1668: cache the mask only for backward return ApplyMask(input, mask); } @@ -289,10 +241,6 @@ private Tensor ApplyMask(Tensor input, Tensor mask) /// public override void ResetState() { - // Clear cached values from forward pass - _lastInput = null; - _lastMask = null; - // Clear GPU mask tensor if (_lastMaskGpu is IDisposable disposable) { diff --git a/src/NeuralNetworks/Layers/MaxPool3DLayer.cs b/src/NeuralNetworks/Layers/MaxPool3DLayer.cs index b29c8eee8a..ed8fdd6bd6 100644 --- a/src/NeuralNetworks/Layers/MaxPool3DLayer.cs +++ b/src/NeuralNetworks/Layers/MaxPool3DLayer.cs @@ -133,17 +133,6 @@ public partial class MaxPool3DLayer : LayerBase, IShapeContract #region Private Fields - /// - /// Stores the indices of maximum values from the forward pass for gradient routing. - /// Shape: [batch, channels, outD, outH, outW, 3] where the last dimension stores [d, h, w] indices. - /// - private int[,,,,,]? _maxIndices; - - /// - /// Cached input from the last forward pass. - /// - [Scratch] - private Tensor? _lastInput; private int[]? _originalInputShape; /// @@ -233,7 +222,6 @@ protected override void OnFirstForward(Tensor input) protected override Tensor ForwardTraced(Tensor input) { EnsureInitializedFromInput(input); - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) _originalInputShape = input._shape; int rank = input.Rank; @@ -265,7 +253,7 @@ protected override Tensor ForwardTraced(Tensor input) batchedInput, [PoolSize, PoolSize, PoolSize], [Stride, Stride, Stride], - out _maxIndices); + out _); // Restore original tensor rank if (_originalInputShape.Length > 5) @@ -322,8 +310,6 @@ internal override Dictionary GetMetadata() /// public override void ResetState() { - _lastInput = null; - _maxIndices = null; _gpuIndicesBuffer?.Dispose(); _gpuIndicesBuffer = null; } diff --git a/src/NeuralNetworks/Layers/MaxPoolingLayer.cs b/src/NeuralNetworks/Layers/MaxPoolingLayer.cs index ae5509b8b8..6030a00381 100644 --- a/src/NeuralNetworks/Layers/MaxPoolingLayer.cs +++ b/src/NeuralNetworks/Layers/MaxPoolingLayer.cs @@ -128,23 +128,12 @@ public int[] GetStride() /// private IGpuBuffer? _gpuIndices; - /// - /// Stores the last input tensor from the forward pass for use in autodiff backward pass. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Tracks whether a batch dimension was added during the forward pass. /// private bool _addedBatchDimension; private int[]? _originalInputShape; - /// - /// Stores the actual output shape for 3D inputs (may differ from pre-computed OutputShape). - /// - private int[]? _lastOutputShape; - /// /// Creates a new max pooling layer with the specified parameters. /// @@ -248,7 +237,6 @@ protected override Tensor ForwardTraced(Tensor input) if (input.Shape.Length < 3) throw new ArgumentException($"MaxPooling layer requires at least 3D tensor [C, H, W]. Got rank {input.Shape.Length}."); - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) _originalInputShape = input._shape; int rank = input.Shape.Length; @@ -294,13 +282,11 @@ protected override Tensor ForwardTraced(Tensor input) outputShape[_originalInputShape.Length - 3] = output4D.Shape[1]; outputShape[_originalInputShape.Length - 2] = output4D.Shape[2]; outputShape[_originalInputShape.Length - 1] = output4D.Shape[3]; - _lastOutputShape = outputShape; return Engine.Reshape(output4D, outputShape); } if (_addedBatchDimension) { var actualOutputShape = new int[] { output4D.Shape[1], output4D.Shape[2], output4D.Shape[3] }; - _lastOutputShape = actualOutputShape; return Engine.Reshape(output4D, actualOutputShape); } return output4D; @@ -344,7 +330,6 @@ public override IEnumerable GetActivationTypes() public override void ResetState() { // Clear cached values from forward pass - _lastInput = null; _addedBatchDimension = false; // Dispose GPU resources diff --git a/src/NeuralNetworks/Layers/MeanLayer.cs b/src/NeuralNetworks/Layers/MeanLayer.cs index 51fabe39e5..a16fd46577 100644 --- a/src/NeuralNetworks/Layers/MeanLayer.cs +++ b/src/NeuralNetworks/Layers/MeanLayer.cs @@ -129,26 +129,6 @@ public partial class MeanLayer : LayerBase, IShapeContract /// public override bool SupportsGpuTraining => true; - /// - /// The input tensor from the most recent forward pass. - /// - /// - /// This field stores the input tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// The output tensor from the most recent forward pass. - /// - /// - /// This field stores the output tensor from the most recent forward pass, which may be - /// useful for certain operations or debugging. - /// - [Scratch] - private Tensor? _lastOutput; - /// /// Initializes a new instance of the class with the specified input shape and axis. /// @@ -263,12 +243,8 @@ private static int[] CalculateOutputShape(int[] inputShape, int axis) protected override Tensor ForwardTraced(Tensor input) { EnsureInitializedFromInput(input); - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) - // Use Engine operation for GPU/CPU acceleration - _lastOutput = Engine.ReduceMean(input, [Axis], keepDims: false); - - return _lastOutput; + return Engine.ReduceMean(input, [Axis], keepDims: false); } /// @@ -322,8 +298,7 @@ internal override Dictionary GetMetadata() /// public override void ResetState() { - // Clear cached values from forward pass - _lastInput = null; - _lastOutput = null; + // Nothing to clear: the forward keeps no per-pass state. + return; } } diff --git a/src/NeuralNetworks/Layers/MeasurementLayer.cs b/src/NeuralNetworks/Layers/MeasurementLayer.cs index 02130679fe..12f67360a4 100644 --- a/src/NeuralNetworks/Layers/MeasurementLayer.cs +++ b/src/NeuralNetworks/Layers/MeasurementLayer.cs @@ -44,25 +44,6 @@ public partial class MeasurementLayer : LayerBase, IShapeContract /// Construction state, retained so the layer can be rebuilt exactly rather than inferred from its shape. private readonly int _size; - /// - /// The input tensor from the most recent forward pass. - /// - /// - /// This field stores the input tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// The output tensor from the most recent forward pass. - /// - /// - /// This field stores the output tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastOutput; private int[]? _originalInputShape; @@ -175,8 +156,6 @@ protected override Tensor ForwardTraced(Tensor input) input2D = Engine.Reshape(input, new[] { batchDim, stateSize }); } - _lastInput = input2D; - int batchSize = input2D.Shape[0]; var probabilities = TensorAllocator.Rent(new[] { batchSize, stateSize }); @@ -205,8 +184,6 @@ protected override Tensor ForwardTraced(Tensor input) } } - _lastOutput = probabilities; - if (input.Rank == 1) { return Engine.Reshape(probabilities, [stateSize]); @@ -246,8 +223,6 @@ protected override Tensor ForwardTraced(Tensor input) public override void ResetState() { // Clear cached values from forward and backward passes - _lastInput = null; - _lastOutput = null; _originalInputShape = null; } diff --git a/src/NeuralNetworks/Layers/MemoryReadLayer.cs b/src/NeuralNetworks/Layers/MemoryReadLayer.cs index 43e8f0901b..d6d12a0fdc 100644 --- a/src/NeuralNetworks/Layers/MemoryReadLayer.cs +++ b/src/NeuralNetworks/Layers/MemoryReadLayer.cs @@ -136,36 +136,6 @@ public partial class MemoryReadLayer : LayerBase, IAuxiliaryLossLayer, private Tensor _outputBias; - /// - /// The input tensor from the most recent forward pass. - /// - /// - /// This field stores the input tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// The memory tensor from the most recent forward pass. - /// - /// - /// This field stores the memory tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastMemory; - - /// - /// The output tensor from the most recent forward pass. - /// - /// - /// This field stores the output tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastOutput; - /// /// The attention scores tensor from the most recent forward pass. /// @@ -176,16 +146,6 @@ public partial class MemoryReadLayer : LayerBase, IAuxiliaryLossLayer, [Scratch] private Tensor? _lastAttentionScores; - /// - /// The transformed tensor from the most recent forward pass (input to output weights). - /// - /// - /// This field stores the transformed tensor (result of readValues × valueWeights) from the most - /// recent forward pass, which is needed during the backward pass for output weights gradient calculation. - /// - [Scratch] - private Tensor? _lastTransformed; - /// /// The gradient of the loss with respect to the key weights. /// @@ -469,9 +429,6 @@ public Tensor Forward(Tensor input, Tensor memory) // initial [0, 0] sentinel — matmul then throws on incompatible dims. EnsureInitializedFromInput(input); - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) - _lastMemory = memory; - // Use Engine operations for matrix multiplications var keys = Engine.TensorMatMul(input, _keyWeights); @@ -482,14 +439,14 @@ public Tensor Forward(Tensor input, Tensor memory) // Apply softmax to get attention weights var softmaxActivation = new SoftmaxActivation(); var attentionWeights = softmaxActivation.Activate(attentionScores); - _lastAttentionScores = attentionWeights; + // Read only by ComputeAuxiliaryLoss, so it is kept only while that loss can be asked for. + _lastAttentionScores = IsTrainingMode && UseAuxiliaryLoss ? attentionWeights : null; // Read from memory: attentionWeights × memory var readValues = Engine.TensorMatMul(attentionWeights, memory); // Apply value and output transformations var transformed = Engine.TensorMatMul(readValues, _valueWeights); - _lastTransformed = transformed; // Cache for output weights gradient computation var projected = Engine.TensorMatMul(transformed, _outputWeights); // Broadcast bias and add @@ -497,15 +454,7 @@ public Tensor Forward(Tensor input, Tensor memory) var biasBroadcast = BroadcastBiases(_outputBias, batchSize); var output = Engine.TensorAdd(projected, biasBroadcast); - var result = ApplyActivation(output); - - // Only store for backward pass during training - skip during inference - if (IsTrainingMode) - { - _lastOutput = result; - } - - return result; + return ApplyActivation(output); } /// @@ -745,11 +694,7 @@ public override void ClearGradients() public override void ResetState() { // Clear cached values from forward and backward passes - _lastInput = null; - _lastMemory = null; - _lastOutput = null; _lastAttentionScores = null; - _lastTransformed = null; _keyWeightsGradient = null; _valueWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/MemoryWriteLayer.cs b/src/NeuralNetworks/Layers/MemoryWriteLayer.cs index d94512a460..6d966d582d 100644 --- a/src/NeuralNetworks/Layers/MemoryWriteLayer.cs +++ b/src/NeuralNetworks/Layers/MemoryWriteLayer.cs @@ -146,36 +146,6 @@ public partial class MemoryWriteLayer : LayerBase, IAuxiliaryLossLayer, private Tensor _outputBias; - /// - /// The input tensor from the most recent forward pass. - /// - /// - /// This field stores the input tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// The memory tensor from the most recent forward pass. - /// - /// - /// This field stores the memory tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastMemory; - - /// - /// The output tensor from the most recent forward pass. - /// - /// - /// This field stores the output tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastOutput; - /// /// The attention scores tensor from the most recent forward pass. /// @@ -186,25 +156,6 @@ public partial class MemoryWriteLayer : LayerBase, IAuxiliaryLossLayer, [Scratch] private Tensor? _lastAttentionScores; - /// - /// The write values tensor from the most recent forward pass. - /// - /// - /// This field stores the write values tensor computed from attention weights, which can be used - /// when updating memory or computing auxiliary objectives. - /// - [Scratch] - private Tensor? _lastWriteValues; - - /// - /// The transformed values tensor from the most recent forward pass. - /// - /// - /// This field stores the input values after value projection, which are used for output gradients. - /// - [Scratch] - private Tensor? _lastValues; - /// /// The gradient of the loss with respect to the query weights. /// @@ -488,9 +439,6 @@ public Tensor Forward(Tensor input, Tensor memory) // but only when the layer has run at least once. Force init here. EnsureInitializedFromInput(input); - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) - _lastMemory = memory; - // Dynamic input dimension adaptation int actualInputDim = input.Shape[^1]; int expectedInputDim = _queryWeights.Shape[0]; @@ -523,7 +471,6 @@ public Tensor Forward(Tensor input, Tensor memory) var queries = Engine.TensorMatMul(input, _queryWeights); var keys = Engine.TensorMatMul(input, _keyWeights); var values = Engine.TensorMatMul(input, _valueWeights); - _lastValues = values; // Compute attention scores: queries × memory^T var memoryTransposed = Engine.TensorTranspose(memory); @@ -536,15 +483,8 @@ public Tensor Forward(Tensor input, Tensor memory) // Apply softmax to get attention weights var softmaxActivation = new SoftmaxActivation(); var attentionWeights = softmaxActivation.Activate(attentionScores); - _lastAttentionScores = attentionWeights; - - // Compute write values: attention-weighted combination of values - // attentionWeights: [batch, numSlots], values: [batch, memoryDim] - // We want: [numSlots, memoryDim] to update memory - // Use: attentionWeights^T @ values = [numSlots, batch] @ [batch, memoryDim] = [numSlots, memoryDim] - var attentionT = Engine.TensorTranspose(attentionWeights); - var writeValues = Engine.TensorMatMul(attentionT, values); - _lastWriteValues = writeValues; // Cache for gradient computation / memory update + // Read only by ComputeAuxiliaryLoss, so it is kept only while that loss can be asked for. + _lastAttentionScores = IsTrainingMode && UseAuxiliaryLoss ? attentionWeights : null; // For the output, we transform the values (what we're writing) through output weights // This maintains [batch, memoryDim] shape for downstream layers @@ -556,15 +496,7 @@ public Tensor Forward(Tensor input, Tensor memory) var biasBroadcast = BroadcastBiases(_outputBias, batchSize); var output = Engine.TensorAdd(projected, biasBroadcast); - var result = ApplyActivation(output); - - // Only store for backward pass during training - skip during inference - if (IsTrainingMode) - { - _lastOutput = result; - } - - return result; + return ApplyActivation(output); } /// @@ -731,12 +663,7 @@ public override void ClearGradients() public override void ResetState() { // Clear cached values from forward and backward passes - _lastInput = null; - _lastMemory = null; - _lastOutput = null; _lastAttentionScores = null; - _lastWriteValues = null; - _lastValues = null; _queryWeightsGradient = null; _keyWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/MeshEdgeConvLayer.cs b/src/NeuralNetworks/Layers/MeshEdgeConvLayer.cs index 3c53cd8c2c..7ecff51fde 100644 --- a/src/NeuralNetworks/Layers/MeshEdgeConvLayer.cs +++ b/src/NeuralNetworks/Layers/MeshEdgeConvLayer.cs @@ -157,29 +157,11 @@ internal override Dictionary GetMetadata() [Scratch] private Tensor? _biasesGradient; - /// - /// Cached input from the last forward pass. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Cached edge adjacency from the last forward pass. /// private int[,]? _lastEdgeAdjacency; - /// - /// Cached output before activation from the last forward pass. - /// - [Scratch] - private Tensor? _lastPreActivation; - - /// - /// Cached output after activation from the last forward pass. - /// - [Scratch] - private Tensor? _lastOutput; - #endregion #region Constructors @@ -349,8 +331,6 @@ protected override Tensor ForwardTraced(Tensor input) "Edge adjacency must be set via SetEdgeAdjacency before calling Forward."); } - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) - int numEdges = input.Shape[0]; int aggregatedFeatures = InputChannels * (1 + NumNeighbors); @@ -360,10 +340,8 @@ protected override Tensor ForwardTraced(Tensor input) var weightsTransposed = Engine.TensorTranspose(_weights); var preActivation = Engine.TensorMatMul(aggregatedInput, weightsTransposed); preActivation = AddBiases(preActivation); - _lastPreActivation = preActivation; var output = ApplyActivation(preActivation); - _lastOutput = output; return output; } @@ -589,9 +567,6 @@ public override void UpdateParameters(T learningRate) /// public override void ResetState() { - _lastInput = null; - _lastPreActivation = null; - _lastOutput = null; _weightsGradient = null; _biasesGradient = null; } diff --git a/src/NeuralNetworks/Layers/MeshPoolLayer.cs b/src/NeuralNetworks/Layers/MeshPoolLayer.cs index 128a14b97a..bba8d0f654 100644 --- a/src/NeuralNetworks/Layers/MeshPoolLayer.cs +++ b/src/NeuralNetworks/Layers/MeshPoolLayer.cs @@ -158,29 +158,11 @@ internal override Dictionary GetMetadata() [Scratch] private Tensor? _importanceWeightsGradient; - /// - /// Cached input from the last forward pass. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Cached edge adjacency from the last forward pass. /// private int[,]? _lastEdgeAdjacency; - /// - /// Cached output from the last forward pass. - /// - [Scratch] - private Tensor? _lastOutput; - - /// - /// Cached importance scores from the last forward pass. - /// - [Scratch] - private Tensor? _lastImportanceScores; - /// /// Number of neighboring edges per edge (default 4 for triangular meshes). /// @@ -274,13 +256,10 @@ protected override Tensor ForwardTraced(Tensor input) "Edge adjacency must be set via SetEdgeAdjacency before calling Forward."); } - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) - int numEdges = input.Shape[0]; int numToKeep = Math.Min(TargetEdges, numEdges); var importanceScores = ComputeImportanceScores(input); - _lastImportanceScores = importanceScores; var sortedIndices = SortEdgesByImportance(importanceScores, numEdges); @@ -304,7 +283,6 @@ protected override Tensor ForwardTraced(Tensor input) UpdatedAdjacency = UpdateAdjacency(_lastEdgeAdjacency, RemainingEdgeIndices, numToKeep); var output = new Tensor(outputData, [numToKeep, InputChannels]); - _lastOutput = output; return output; } @@ -492,9 +470,6 @@ public override void UpdateParameters(T learningRate) /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastImportanceScores = null; _importanceWeightsGradient = null; RemainingEdgeIndices = null; UpdatedAdjacency = null; diff --git a/src/NeuralNetworks/Layers/MessagePassingLayer.cs b/src/NeuralNetworks/Layers/MessagePassingLayer.cs index 0e89fd31f6..66cc0c2280 100644 --- a/src/NeuralNetworks/Layers/MessagePassingLayer.cs +++ b/src/NeuralNetworks/Layers/MessagePassingLayer.cs @@ -211,53 +211,11 @@ public partial class MessagePassingLayer : LayerBase, IGraphConvolutionLay /// private bool _useSparseAggregation = false; - /// - /// Cached input from forward pass. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Stores the original input shape for any-rank tensor support. /// private int[]? _originalInputShape; - /// - /// Cached output from forward pass. - /// - [Scratch] - private Tensor? _lastOutput; - - /// - /// Cached messages for backward pass. - /// - [Scratch] - private Tensor? _lastMessages; - - /// - /// Cached aggregated messages. - /// - [Scratch] - private Tensor? _lastAggregated; - - /// - /// Cached hidden activations from message MLP layer 1. - /// - [Scratch] - private Tensor? _lastMessageHidden; - - /// - /// Cached reset gates. - /// - [Scratch] - private Tensor? _lastResetGate; - - /// - /// Cached update gates. - /// - [Scratch] - private Tensor? _lastUpdateGate; - /// /// Gradients for parameters. /// @@ -635,8 +593,6 @@ protected override Tensor ForwardTraced(Tensor input) } } - _lastInput = processInput; - // Steps 1+2 (message MLP + masked aggregation), bulk-vectorized. // ---------------------------------------------------------------- // Original implementation walked every (b, i, j) edge in scalar: @@ -691,7 +647,6 @@ protected override Tensor ForwardTraced(Tensor input) var hidden = Engine.TensorMatMul(inputFlat, _messageWeights1); hidden = Engine.TensorAdd(hidden, bias1Broadcast); hidden = Engine.ReLU(hidden); - _lastMessageHidden = Engine.Reshape(hidden, [batchSize, numNodes, numNodes, _messageFeatures]); // Layer 2: hidden @ W2 + b2. var bias2Broadcast = Engine.Reshape(_messageBias2, [1, _messageFeatures]); @@ -706,10 +661,10 @@ protected override Tensor ForwardTraced(Tensor input) ? new[] { 1, numNodes, numNodes, 1 } : new[] { batchSize, numNodes, numNodes, 1 }; var adjBroadcastable = Engine.Reshape(_adjacencyMatrix, adjShape); - _lastMessages = Engine.TensorMultiply(messages4D, adjBroadcastable); + var maskedMessages = Engine.TensorMultiply(messages4D, adjBroadcastable); // Step 2: aggregate by summing over the j axis (axis=2 of [B, N(i), N(j), M]). - _lastAggregated = Engine.ReduceSum(_lastMessages, new[] { 2 }, keepDims: false); + var aggregated = Engine.ReduceSum(maskedMessages, new[] { 2 }, keepDims: false); // Step 3: Update node features (GRU-style update), bulk-vectorized. // ---------------------------------------------------------------- @@ -728,7 +683,7 @@ protected override Tensor ForwardTraced(Tensor input) // operands and the broadcast-add on biases stays simple. int bn = batchSize * numNodes; var inputFlatBN = Engine.Reshape(processInput, [bn, _inputFeatures]); - var aggFlatBN = Engine.Reshape(_lastAggregated, [bn, _messageFeatures]); + var aggFlatBN = Engine.Reshape(aggregated, [bn, _messageFeatures]); // Reset gate: σ(input @ W_r + agg @ W_mr + b_r) var resetBiasBcast = Engine.Reshape(_resetBias, [1, _outputFeatures]); @@ -736,7 +691,6 @@ protected override Tensor ForwardTraced(Tensor input) resetLogitsFlat = Engine.TensorAdd(resetLogitsFlat, Engine.TensorMatMul(aggFlatBN, _resetMessageWeights)); resetLogitsFlat = Engine.TensorAdd(resetLogitsFlat, resetBiasBcast); var resetFlat = Engine.Sigmoid(resetLogitsFlat); - _lastResetGate = Engine.Reshape(resetFlat, [batchSize, numNodes, _outputFeatures]); // Update gate: σ(input @ W_u + agg @ W_mu + b_u) var updateBiasBcast = Engine.Reshape(_updateBias, [1, _outputFeatures]); @@ -744,7 +698,6 @@ protected override Tensor ForwardTraced(Tensor input) updateLogitsFlat = Engine.TensorAdd(updateLogitsFlat, Engine.TensorMatMul(aggFlatBN, _updateMessageWeights)); updateLogitsFlat = Engine.TensorAdd(updateLogitsFlat, updateBiasBcast); var updateFlat = Engine.Sigmoid(updateLogitsFlat); - _lastUpdateGate = Engine.Reshape(updateFlat, [batchSize, numNodes, _outputFeatures]); // Output combiner: (1 - update) * input_padded + update * agg_padded. // The original conditional `f < inputFeatures ? processInput : 0` zeros @@ -761,15 +714,7 @@ protected override Tensor ForwardTraced(Tensor input) var outputFlat = Engine.TensorAdd(oldContribution, newContribution); var output = Engine.Reshape(outputFlat, [batchSize, numNodes, _outputFeatures]); - var result = ApplyActivation(output); - - // Only store for backward pass during training - skip during inference - if (IsTrainingMode) - { - _lastOutput = result; - } - - return result; + return ApplyActivation(output); } /// @@ -893,13 +838,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastMessages = null; - _lastAggregated = null; - _lastMessageHidden = null; - _lastResetGate = null; - _lastUpdateGate = null; _messageWeights1Gradient = null; _messageWeights2Gradient = null; _messageBias1Gradient = null; diff --git a/src/NeuralNetworks/Layers/MixtureOfExpertsLayer.cs b/src/NeuralNetworks/Layers/MixtureOfExpertsLayer.cs index 9c5d09e241..000166b61d 100644 --- a/src/NeuralNetworks/Layers/MixtureOfExpertsLayer.cs +++ b/src/NeuralNetworks/Layers/MixtureOfExpertsLayer.cs @@ -244,28 +244,6 @@ public partial class MixtureOfExpertsLayer : LayerBase, IAuxiliaryLossLaye [Scratch] private Tensor? _lastRoutingWeights; - /// - /// Cached expert outputs from the most recent forward pass. - /// - /// - /// - /// This stores the output from each expert before they are combined. During backpropagation, - /// these are used to compute gradients for both the router and the experts themselves. - /// - /// For Beginners: Stored memory of what each expert predicted. - /// - /// We save each expert's output so we can: - /// - Calculate how much each expert contributed to any errors - /// - Adjust the router if it chose the wrong experts - /// - Adjust each expert to improve its predictions - /// - /// It's like keeping a record of each team member's contribution to a group project, - /// so you can give appropriate feedback to each person. - /// - /// - [Scratch] - private List>? _lastExpertOutputs; - /// /// Cached combined output before activation from the most recent forward pass. /// @@ -273,26 +251,6 @@ public partial class MixtureOfExpertsLayer : LayerBase, IAuxiliaryLossLaye private Tensor? _lastPreActivation; - /// - /// Cached routing logits (before softmax) from the most recent forward pass. - /// - /// - /// - /// These are the raw scores from the router before applying softmax. They're needed for - /// computing the gradient of softmax during backpropagation. - /// - /// For Beginners: The router's raw scores before normalization. - /// - /// The router produces raw scores (logits), then we convert them to probabilities (weights). - /// We save the raw scores because: - /// - The conversion process (softmax) is nonlinear - /// - To compute gradients correctly, we need the pre-conversion values - /// - This is a technical requirement for proper backpropagation through softmax - /// - /// - [Scratch] - private Tensor? _lastRoutingLogits; - /// /// Cached top-K indices for sparse routing from the most recent forward pass. /// @@ -692,7 +650,7 @@ protected override Tensor ForwardTraced(Tensor input) } // Cache input for backward pass - _lastInput = ShouldCacheForBackward ? input2D : null; // #1668: skip in inference (arena safety) + SaveForBackward(ref _lastInput, input2D); // Step 1: Compute routing scores var routingLogits = _router.Forward(input2D); @@ -720,8 +678,6 @@ protected override Tensor ForwardTraced(Tensor input) routingLogits = Engine.TensorAdd(routingLogits, Engine.TensorMultiply(noise, noiseScale)); } - _lastRoutingLogits = routingLogits; - // Step 2: Apply softmax to get routing probabilities var routingWeights = ApplySoftmax(routingLogits); @@ -738,7 +694,7 @@ protected override Tensor ForwardTraced(Tensor input) _lastRoutingWeights = routingWeights; // Step 4: Get outputs from all experts (or only top-K if sparse) - _lastExpertOutputs = new List>(); + var expertOutputs = new List>(_experts.Count); _activeExpertsDuringBackward = []; if (_topK > 0) @@ -759,13 +715,13 @@ protected override Tensor ForwardTraced(Tensor input) if (isActive) { - _lastExpertOutputs.Add(_experts[i].Forward(input2D)); + expertOutputs.Add(_experts[i].Forward(input2D)); _activeExpertsDuringBackward.Add(i); } else { // Create a zero tensor for inactive experts - _lastExpertOutputs.Add(new Tensor(new int[] { batchSize }.Concat(OutputShape).ToArray())); + expertOutputs.Add(new Tensor(new int[] { batchSize }.Concat(OutputShape).ToArray())); } } } @@ -774,14 +730,14 @@ protected override Tensor ForwardTraced(Tensor input) // Dense: All experts process all inputs for (int i = 0; i < _experts.Count; i++) { - _lastExpertOutputs.Add(_experts[i].Forward(input2D)); + expertOutputs.Add(_experts[i].Forward(input2D)); _activeExpertsDuringBackward.Add(i); } } // Step 5: Combine expert outputs using routing weights - var combinedOutput = CombineExpertOutputs(_lastExpertOutputs, routingWeights); - _lastPreActivation = combinedOutput; + var combinedOutput = CombineExpertOutputs(expertOutputs, routingWeights); + SaveForBackward(ref _lastPreActivation, combinedOutput); // Step 6: Apply activation function var output = ApplyActivation(combinedOutput); @@ -884,8 +840,6 @@ public override void ResetState() { _lastInput = null; _lastRoutingWeights = null; - _lastRoutingLogits = null; - _lastExpertOutputs = null; _lastPreActivation = null; _lastTopKIndices = null; _originalInputShape = null; diff --git a/src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cs b/src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cs index 43eb57668f..b9e52b20ed 100644 --- a/src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cs +++ b/src/NeuralNetworks/Layers/MultiHeadAttentionLayer.cs @@ -1645,7 +1645,8 @@ static Tensor ApplyQkvTransform(Tensor? produced, Tensor original, stri // per MHA forward; on a 12-layer LayoutLM × 30 training iters // that's 360 dead permute/tape entries per Train_ShouldReduceLoss // run, each one tracked + walked by the gradient tape backward. - if (UseAuxiliaryLoss) + // Training only: an inference forward never asks for the auxiliary loss. + if (UseAuxiliaryLoss && IsTrainingMode) { // via Engine so the tape records the transpose var permutedForCache = Engine.TensorPermute(context_4D, new[] { 1, 0, 2, 3 }); // [H, B, S, D] diff --git a/src/NeuralNetworks/Layers/MultiplyLayer.cs b/src/NeuralNetworks/Layers/MultiplyLayer.cs index 82b234951c..5f9b2cc5ba 100644 --- a/src/NeuralNetworks/Layers/MultiplyLayer.cs +++ b/src/NeuralNetworks/Layers/MultiplyLayer.cs @@ -59,16 +59,6 @@ public partial class MultiplyLayer : LayerBase, IShapeContract /// private Tensor[]? _lastInputs; - /// - /// The output tensor from the most recent forward pass. - /// - /// - /// This field stores the output tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastOutput; - /// /// Gets a value indicating whether this layer supports training. /// @@ -296,21 +286,13 @@ protected override Tensor ForwardTracedMany(params Tensor[] inputs) throw new ArgumentException("MultiplyLayer requires at least two inputs."); } - _lastInputs = inputs; + SaveForBackward(ref _lastInputs, inputs); // Use Engine.TensorMultiplyMany for GPU/CPU accelerated element-wise multiplication of all tensors // This is production-grade: no loops, single optimized call that batches all multiplications var result = Engine.TensorMultiplyMany(inputs); - var activated = ApplyActivation(result); - - // Only store for backward pass during training - skip during inference - if (IsTrainingMode) - { - _lastOutput = activated; - } - - return activated; + return ApplyActivation(result); } /// @@ -405,6 +387,5 @@ public override void ResetState() { // Clear cached values from forward pass _lastInputs = null; - _lastOutput = null; } } diff --git a/src/NeuralNetworks/Layers/ObliviousDecisionTreeLayer.cs b/src/NeuralNetworks/Layers/ObliviousDecisionTreeLayer.cs index 438aed7ede..afbc04c3c3 100644 --- a/src/NeuralNetworks/Layers/ObliviousDecisionTreeLayer.cs +++ b/src/NeuralNetworks/Layers/ObliviousDecisionTreeLayer.cs @@ -99,13 +99,7 @@ public partial class ObliviousDecisionTreeLayer : LayerBase, IShapeContrac // Cached values [Scratch] - private Tensor? _inputCache; - [Scratch] private Tensor? _featureSelectionsCache; - [Scratch] - private Tensor? _splitDecisionsCache; - [Scratch] - private Tensor? _leafProbabilitiesCache; private readonly int _numLeaves; @@ -328,7 +322,6 @@ protected override Tensor ForwardTraced(Tensor input) $"was resolved with _inputDim={_inputDim}, but input has " + $"{input.Shape[1]} features.", nameof(input)); - _inputCache = input; int batchSize = input.Shape[0]; // Compute soft feature selections for each level (entmax/softmax) @@ -337,11 +330,9 @@ protected override Tensor ForwardTraced(Tensor input) // Compute split decisions for each level var splitDecisions = ComputeSplitDecisions(input, featureSelections, batchSize); - _splitDecisionsCache = splitDecisions; // Compute leaf probabilities var leafProbs = ComputeLeafProbabilities(splitDecisions, batchSize); - _leafProbabilitiesCache = leafProbs; // Weighted sum of leaf values var output = ComputeOutput(leafProbs, batchSize); @@ -468,10 +459,7 @@ public override void UpdateParameters(T learningRate) /// public override void ResetState() { - _inputCache = null; _featureSelectionsCache = null; - _splitDecisionsCache = null; - _leafProbabilitiesCache = null; Engine.TensorFill(_featureSelectionGrad, NumOps.Zero); Engine.TensorFill(_thresholdsGrad, NumOps.Zero); diff --git a/src/NeuralNetworks/Layers/OctonionLinearLayer.cs b/src/NeuralNetworks/Layers/OctonionLinearLayer.cs index cf7ed940bf..5c0b26f021 100644 --- a/src/NeuralNetworks/Layers/OctonionLinearLayer.cs +++ b/src/NeuralNetworks/Layers/OctonionLinearLayer.cs @@ -113,25 +113,11 @@ public partial class OctonionLinearLayer : LayerBase, IShapeContract private Tensor _biases; - /// - /// Stored input from forward pass for backpropagation. - /// Shape: [batch, inputFeatures, 8] - /// - [Scratch] - private Tensor? _lastInput; - /// /// Stores the original input shape for any-rank tensor support. /// private int[]? _originalInputShape; - /// - /// Stored pre-activation output for gradient computation. - /// Shape: [batch, outputFeatures, 8] - /// - [Scratch] - private Tensor? _lastOutput; - /// /// Gradient for weights. Shape: [OutputFeatures, InputFeatures, 8] /// @@ -254,8 +240,6 @@ protected override Tensor ForwardTraced(Tensor input) // Reshape flat input [batch, inputFeatures*8] → [batch, inputFeatures, 8] var input3D = Engine.Reshape(inputTensor, [batchSize, InputFeatures, 8]); - _lastInput = input3D; - // Tensor-based octonion matrix multiplication — no Octonion allocation var output3D = Engine.OctonionMatMulTensor(input3D, _weights); @@ -264,8 +248,6 @@ protected override Tensor ForwardTraced(Tensor input) // which made every octonion bias look detached to reverse-mode autodiff. output3D = Engine.TensorAdd(output3D, _biases); - _lastOutput = output3D; - // Flatten back to [batch, outputFeatures*8] for activation and output var outputTensor = Engine.Reshape(output3D, [batchSize, OutputFeatures * 8]); @@ -323,8 +305,6 @@ public override void ClearGradients() public override void ResetState() { - _lastInput = null; - _lastOutput = null; _originalInputShape = null; _weightsGradient = null; _biasesGradient = null; diff --git a/src/NeuralNetworks/Layers/PReLULayer.cs b/src/NeuralNetworks/Layers/PReLULayer.cs index c0df892b51..4fc95d5cea 100644 --- a/src/NeuralNetworks/Layers/PReLULayer.cs +++ b/src/NeuralNetworks/Layers/PReLULayer.cs @@ -55,8 +55,6 @@ public partial class PReLULayer : LayerBase [TrainableParameter(Role = PersistentTensorRole.Weights)] private Tensor _alpha; - private Tensor? _lastInput; - /// /// Gets a value indicating that this layer has trainable parameters (α). /// @@ -160,8 +158,6 @@ protected override void OnFirstForward(Tensor input) protected override Tensor ForwardTraced(Tensor input) { EnsureInitializedFromInput(input); - // #1668: also clear a previously-cached activation when caching is off (don't leave it pinned). - _lastInput = ShouldCacheForBackward ? input : null; // Use the dedicated IEngine primitive where its channel convention exactly matches this // layer: shared slope, a rank-1 elementwise slope, rank-2 [batch, channels], or NCHW. It @@ -191,6 +187,5 @@ internal override Dictionary GetMetadata() } public override void ResetState() { - _lastInput = null; } } diff --git a/src/NeuralNetworks/Layers/PaddingLayer.cs b/src/NeuralNetworks/Layers/PaddingLayer.cs index b34635c852..4c017393e0 100644 --- a/src/NeuralNetworks/Layers/PaddingLayer.cs +++ b/src/NeuralNetworks/Layers/PaddingLayer.cs @@ -95,16 +95,6 @@ public partial class PaddingLayer : LayerBase, IShapeContract return axes; } - /// - /// The input tensor from the most recent forward pass. - /// - /// - /// This field stores the input tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Gets a value indicating whether this layer supports training. /// @@ -296,7 +286,6 @@ private static int[] CalculateOutputShape(int[] inputShape, int[] padding) protected override Tensor ForwardTraced(Tensor input) { EnsureInitializedFromInput(input); - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) if (_padding.Length != input.Shape.Length) throw new ArgumentException("Padding array length must match input dimensions."); @@ -331,6 +320,5 @@ protected override Tensor ForwardTraced(Tensor input) public override void ResetState() { // Clear cached values from forward pass - _lastInput = null; } } diff --git a/src/NeuralNetworks/Layers/PatchEmbeddingLayer.cs b/src/NeuralNetworks/Layers/PatchEmbeddingLayer.cs index 775cc76a6e..0702026586 100644 --- a/src/NeuralNetworks/Layers/PatchEmbeddingLayer.cs +++ b/src/NeuralNetworks/Layers/PatchEmbeddingLayer.cs @@ -171,12 +171,6 @@ public partial class PatchEmbeddingLayer : LayerBase, IShapeContract private Tensor _projectionBias; - /// - /// Cached input from the forward pass for use in the backward pass. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Stores the original input shape for any-rank tensor support. /// @@ -206,12 +200,6 @@ public partial class PatchEmbeddingLayer : LayerBase, IShapeContract [Scratch] private Tensor? _projectionBiasGradient; - /// - /// Cached pre-activation tensor from forward pass for use in activation derivative calculation. - /// - [Scratch] - private Tensor? _lastPreActivation; - // GPU cached tensors for backward pass [ExternalState] private Tensor? _gpuInput; @@ -605,7 +593,6 @@ protected override Tensor ForwardTraced(Tensor input) new[] { batchSize, _channels, cropHeight, cropWidth }); } - _lastInput = processInput; int patchDim = _channels * _patchSize * _patchSize; // Efficient Patchify using Reshape and Transpose (all via Engine). @@ -644,7 +631,6 @@ protected override Tensor ForwardTraced(Tensor input) var biasBroadcast = Engine.Reshape(_projectionBias, new[] { 1, 1, _embeddingDim }); var preActivation = Engine.TensorAdd(projected, biasBroadcast); - _lastPreActivation = preActivation; var output = ApplyActivation(preActivation); // Restore output shape to match original input rank (via Engine). @@ -705,8 +691,6 @@ public override void ClearGradients() public override void ResetState() { - _lastInput = null; - _lastPreActivation = null; _projectionWeightsGradient = null; _projectionBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/PixelShuffleLayer.cs b/src/NeuralNetworks/Layers/PixelShuffleLayer.cs index b8e81cd039..3a6623cab3 100644 --- a/src/NeuralNetworks/Layers/PixelShuffleLayer.cs +++ b/src/NeuralNetworks/Layers/PixelShuffleLayer.cs @@ -63,12 +63,6 @@ public partial class PixelShuffleLayer : LayerBase, IShapeContract /// private readonly int _upscaleFactor; - /// - /// Cached input from the last forward pass for backpropagation. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Cached original input shape for backward pass with higher-rank tensors. /// @@ -310,7 +304,6 @@ protected override Tensor ForwardTraced(Tensor input) // For 4D tensors (batch, channels, height, width), use Engine directly if (shape.Length == 4) { - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) return Engine.PixelShuffle(input, _upscaleFactor); } @@ -318,7 +311,6 @@ protected override Tensor ForwardTraced(Tensor input) if (shape.Length == 3) { var input4D = Engine.Reshape(input, [1, shape[0], shape[1], shape[2]]); - _lastInput = input4D; var output4D = Engine.PixelShuffle(input4D, _upscaleFactor); // Remove batch dimension from output return Engine.Reshape(output4D, [output4D.Shape[1], output4D.Shape[2], output4D.Shape[3]]); @@ -339,7 +331,6 @@ protected override Tensor ForwardTraced(Tensor input) int width = shape[^1]; var input4D = Engine.Reshape(input, [batchSize, channels, height, width]); - _lastInput = input4D; var output4D = Engine.PixelShuffle(input4D, _upscaleFactor); // Restore original batch dimensions with new spatial dimensions @@ -369,7 +360,6 @@ protected override Tensor ForwardTraced(Tensor input) /// public override void ResetState() { - _lastInput = null; _originalInputShape = null; } diff --git a/src/NeuralNetworks/Layers/PoolingLayer.cs b/src/NeuralNetworks/Layers/PoolingLayer.cs index 8e0af102ac..530d76b552 100644 --- a/src/NeuralNetworks/Layers/PoolingLayer.cs +++ b/src/NeuralNetworks/Layers/PoolingLayer.cs @@ -192,16 +192,6 @@ public partial class PoolingLayer : LayerBase, IShapeContract /// public override bool SupportsTraining => true; - /// - /// The input tensor from the most recent forward pass. - /// - /// - /// This field stores the input tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastInput; - /// /// The execution engine for GPU-accelerated pooling operations. /// @@ -376,7 +366,6 @@ protected override Tensor ForwardTraced(Tensor input) // chain-walkers and tests see the resolved input/output shapes. if (!IsShapeResolved) OnFirstForward(input); - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) _originalInputShape = input._shape; // Support any rank >= 3: last 3 dims are [C, H, W], earlier dims are batch-like @@ -469,7 +458,6 @@ protected override Tensor ForwardTraced(Tensor input) public override void ResetState() { // Clear cached values from forward pass - _lastInput = null; _maxIndices = null; } diff --git a/src/NeuralNetworks/Layers/PrimaryCapsuleLayer.cs b/src/NeuralNetworks/Layers/PrimaryCapsuleLayer.cs index 44e090786a..5396aab13b 100644 --- a/src/NeuralNetworks/Layers/PrimaryCapsuleLayer.cs +++ b/src/NeuralNetworks/Layers/PrimaryCapsuleLayer.cs @@ -100,16 +100,6 @@ public partial class PrimaryCapsuleLayer : LayerBase, IShapeContract [Scratch] private Tensor? _convBiasGradient; - /// - /// The input tensor from the most recent forward pass. - /// - /// - /// This field stores the input tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Stores the original input shape for any-rank tensor support. /// @@ -120,18 +110,6 @@ public partial class PrimaryCapsuleLayer : LayerBase, IShapeContract /// private bool _inputWasNCHW; - /// - /// The output tensor from the most recent forward pass. - /// - /// - /// This field stores the output tensor from the most recent forward pass, which is needed - /// during the backward pass for gradient calculation. - /// - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastPreSquash; - /// /// The number of input channels. /// @@ -635,7 +613,6 @@ protected override Tensor ForwardTraced(Tensor input) } _inputWasNCHW = inputIsNCHW; - _lastInput = processInput; // Get spatial dimensions based on format int inputHeight, inputWidth; @@ -679,9 +656,8 @@ protected override Tensor ForwardTraced(Tensor input) // SquashActivation expects 2D input [numCapsules, capsuleDim], so reshape and apply int totalCapsules = batchSize * outputHeight * outputWidth * _capsuleChannels; var flatOutput = Engine.Reshape(output, [totalCapsules, _capsuleDimension]); - _lastPreSquash = flatOutput; var activatedFlat = ApplyActivation(flatOutput); - _lastOutput = Engine.Reshape(activatedFlat, [batchSize, outputHeight, outputWidth, _capsuleChannels, _capsuleDimension]); + var squashed = Engine.Reshape(activatedFlat, [batchSize, outputHeight, outputWidth, _capsuleChannels, _capsuleDimension]); // Restore output shape to match original input rank // Output is 5D [B, OH, OW, capsuleChannels, capsuleDim] @@ -691,7 +667,7 @@ protected override Tensor ForwardTraced(Tensor input) { // Lower-rank: remove batch dimension // 3D input -> 4D output [OH, OW, capsuleChannels, capsuleDim] - return Engine.Reshape(_lastOutput, [outputHeight, outputWidth, _capsuleChannels, _capsuleDimension]); + return Engine.Reshape(squashed, [outputHeight, outputWidth, _capsuleChannels, _capsuleDimension]); } else { @@ -704,11 +680,11 @@ protected override Tensor ForwardTraced(Tensor input) outShape[_originalInputShape.Length - 2] = outputWidth; outShape[_originalInputShape.Length - 1] = _capsuleChannels; outShape[_originalInputShape.Length] = _capsuleDimension; - return Engine.Reshape(_lastOutput, outShape); + return Engine.Reshape(squashed, outShape); } } - return _lastOutput; + return squashed; } @@ -844,8 +820,6 @@ public override void ClearGradients() public override void ResetState() { // Clear cached values from forward and backward passes - _lastInput = null; - _lastOutput = null; _convWeightsGradient = null; _convBiasGradient = null; _originalInputShape = null; diff --git a/src/NeuralNetworks/Layers/PrincipalNeighbourhoodAggregationLayer.cs b/src/NeuralNetworks/Layers/PrincipalNeighbourhoodAggregationLayer.cs index 6be17e0018..5b84135067 100644 --- a/src/NeuralNetworks/Layers/PrincipalNeighbourhoodAggregationLayer.cs +++ b/src/NeuralNetworks/Layers/PrincipalNeighbourhoodAggregationLayer.cs @@ -100,18 +100,12 @@ public partial class PrincipalNeighbourhoodAggregationLayer : LayerBase, I /// private int[]? _originalInputShape; [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastTransformed; - [Scratch] private Tensor? _lastAggregated; [Scratch] private Tensor? _lastMlpHiddenPreRelu; [Scratch] private Tensor? _lastMlpHidden; [Scratch] - private Tensor? _lastMlpOutput; - [Scratch] private Tensor? _lastDegrees; // Gradients - Tensor-based for GPU acceleration @@ -326,30 +320,31 @@ protected override Tensor ForwardTraced(Tensor input) processInput = Engine.Reshape(input, [flatBatch, numNodes, features]); } - _lastInput = processInput; + SaveForBackward(ref _lastInput, processInput); int processNumNodes = processInput.Shape[1]; // Step 1: Pre-transform input features: transformed = input @ preTransformWeights + preTransformBias // Uses Engine.TensorMatMul for batched matrix multiplication var transformed = Engine.TensorMatMul(processInput, _preTransformWeights); var preBiasBroadcast = BroadcastBias(_preTransformBias, batchSize, processNumNodes); - _lastTransformed = Engine.TensorAdd(transformed, preBiasBroadcast); + var preTransformed = Engine.TensorAdd(transformed, preBiasBroadcast); // Step 2: Compute degrees for each node using adjacency matrix row sums // degrees[b,i] = sum_j(A[b,i,j]) - _lastDegrees = Engine.ReduceSum(_adjacencyMatrix, [2], keepDims: false); // [batch, nodes] + var degrees = Engine.ReduceSum(_adjacencyMatrix, [2], keepDims: false); // [batch, nodes] + SaveForBackward(ref _lastDegrees, degrees); // Avoid division by zero - clamp degrees to minimum of 1 - var oneTensor = new Tensor(_lastDegrees._shape); + var oneTensor = new Tensor(degrees._shape); oneTensor.Fill(NumOps.One); - var safeDegrees = Engine.TensorMax(_lastDegrees, oneTensor); + var safeDegrees = Engine.TensorMax(degrees, oneTensor); // Step 3: Apply multiple aggregators and scalers using vectorized operations var aggregatedParts = new List>(); foreach (var aggregator in _aggregators) { - var aggregated = ComputeVectorizedAggregation(_lastTransformed, aggregator, safeDegrees, processNumNodes); + var aggregated = ComputeVectorizedAggregation(preTransformed, aggregator, safeDegrees, processNumNodes); foreach (var scaler in _scalers) { @@ -359,38 +354,35 @@ protected override Tensor ForwardTraced(Tensor input) } // Concatenate all aggregated features along the last axis - _lastAggregated = Engine.TensorConcatenate(aggregatedParts.ToArray(), axis: 2); // [batch, nodes, combinedFeatures] + var allAggregated = Engine.TensorConcatenate(aggregatedParts.ToArray(), axis: 2); // [batch, nodes, combinedFeatures] + SaveForBackward(ref _lastAggregated, allAggregated); // Step 4: Post-aggregation MLP - Layer 1 with ReLU - var hidden = Engine.TensorMatMul(_lastAggregated, _postAggregationWeights1); + var hidden = Engine.TensorMatMul(allAggregated, _postAggregationWeights1); var bias1Broadcast = BroadcastBias(_postAggregationBias1, batchSize, processNumNodes); - _lastMlpHiddenPreRelu = Engine.TensorAdd(hidden, bias1Broadcast); + var mlpHiddenPreRelu = Engine.TensorAdd(hidden, bias1Broadcast); + SaveForBackward(ref _lastMlpHiddenPreRelu, mlpHiddenPreRelu); // ReLU activation using Engine operations - var zeroTensor = new Tensor(_lastMlpHiddenPreRelu._shape); + var zeroTensor = new Tensor(mlpHiddenPreRelu._shape); zeroTensor.Fill(NumOps.Zero); - _lastMlpHidden = Engine.TensorMax(_lastMlpHiddenPreRelu, zeroTensor); + var mlpHidden = Engine.TensorMax(mlpHiddenPreRelu, zeroTensor); + SaveForBackward(ref _lastMlpHidden, mlpHidden); // Step 5: Post-aggregation MLP - Layer 2 - var mlpOutput = Engine.TensorMatMul(_lastMlpHidden, _postAggregationWeights2); + var mlpOutput = Engine.TensorMatMul(mlpHidden, _postAggregationWeights2); var bias2Broadcast = BroadcastBias(_postAggregationBias2, batchSize, processNumNodes); - _lastMlpOutput = Engine.TensorAdd(mlpOutput, bias2Broadcast); + var mlpProjected = Engine.TensorAdd(mlpOutput, bias2Broadcast); // Step 6: Self-loop transformation and final bias var selfContribution = Engine.TensorMatMul(processInput, _selfWeights); var biasBroadcast = BroadcastBias(_bias, batchSize, processNumNodes); - var preActivation = Engine.TensorAdd(_lastMlpOutput, selfContribution); + var preActivation = Engine.TensorAdd(mlpProjected, selfContribution); preActivation = Engine.TensorAdd(preActivation, biasBroadcast); var result = ApplyActivation(preActivation); - // Only store for backward pass during training - skip during inference - if (IsTrainingMode) - { - _lastOutput = result; - } - // Restore output shape to match original input rank if (_originalInputShape != null && _originalInputShape.Length != 3) { @@ -829,12 +821,9 @@ public override void ClearGradients() public override void ResetState() { _lastInput = null; - _lastOutput = null; - _lastTransformed = null; _lastAggregated = null; _lastMlpHidden = null; _lastMlpHiddenPreRelu = null; - _lastMlpOutput = null; _lastDegrees = null; _preTransformWeightsGradient = null; _preTransformBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/QuantumLayer.cs b/src/NeuralNetworks/Layers/QuantumLayer.cs index 3a5def660a..7fe5b9cc6c 100644 --- a/src/NeuralNetworks/Layers/QuantumLayer.cs +++ b/src/NeuralNetworks/Layers/QuantumLayer.cs @@ -86,7 +86,6 @@ public partial class QuantumLayer : LayerBase, IShapeContract [TrainableParameter(Role = PersistentTensorRole.Weights)] private Tensor _circuitImag; - private Tensor? _lastInput; /// /// Stores the original input shape for any-rank tensor support. @@ -106,14 +105,6 @@ public partial class QuantumLayer : LayerBase, IShapeContract [AiDotNet.Attributes.TrainableParameter] private Tensor _angleGradients; - /// - /// Cached result amplitudes from Forward for use in Backward. - /// Shape: [batch, dimension] for real and imaginary parts. - /// - [Scratch] - private Tensor? _lastResultReal; - [Scratch] - private Tensor? _lastResultImag; private readonly INumericOperations> _complexOps; public override bool SupportsTraining => true; @@ -228,7 +219,6 @@ protected override Tensor ForwardTraced(Tensor input) processInput = Engine.Reshape(input, [flatBatch, input.Shape[rank - 1]]); } - _lastInput = processInput; int dimension = 1 << _numQubits; int inputDim = processInput.Shape[1]; // Use processInput which is always 2D @@ -294,10 +284,6 @@ protected override Tensor ForwardTraced(Tensor input) // Apply quantum circuit using complex matrix multiplication var (resultRealT, resultImagT) = Engine.ComplexMatMul(_circuitReal, _circuitImag, normalizedRealT, normalizedImagT); - // Cache amplitudes for Backward (transpose to [batch, dimension]) - _lastResultReal = Engine.TensorTranspose(resultRealT); - _lastResultImag = Engine.TensorTranspose(resultImagT); - // Convert amplitudes to probabilities and transpose back to [batch, dimension] var probabilitiesT = Engine.ComplexMagnitudeSquared(resultRealT, resultImagT); return Engine.TensorTranspose(probabilitiesT); @@ -399,9 +385,6 @@ public override void ClearGradients() public override void ResetState() { // Clear cached values from forward pass - _lastInput = null; - _lastResultReal = null; - _lastResultImag = null; // Reset angle gradients _angleGradients = new Tensor([_numQubits]); diff --git a/src/NeuralNetworks/Layers/RBFLayer.cs b/src/NeuralNetworks/Layers/RBFLayer.cs index c288bda84a..91d759e623 100644 --- a/src/NeuralNetworks/Layers/RBFLayer.cs +++ b/src/NeuralNetworks/Layers/RBFLayer.cs @@ -132,28 +132,6 @@ public partial class RBFLayer : LayerBase, IShapeContract /// private IRadialBasisFunction _rbf; - /// - /// Stores the input tensor from the most recent forward pass for use in backpropagation. - /// - /// - /// This cached input is needed during the backward pass to compute gradients. It holds the - /// batch of input vectors that were processed in the most recent forward pass. The tensor - /// is null before the first forward pass or after a reset. - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// Stores the output tensor from the most recent forward pass for use in backpropagation. - /// - /// - /// This cached output is needed during the backward pass to compute certain derivatives. - /// It holds the batch of output vectors that were produced in the most recent forward pass. - /// The tensor is null before the first forward pass or after a reset. - /// - [Scratch] - private Tensor? _lastOutput; - /// /// Stores the gradients of the loss with respect to the center parameters. /// @@ -368,18 +346,12 @@ protected override Tensor ForwardTraced(Tensor input) PlaceCentersFromFirstBatch(processedInput); - // Store 2D input for backward pass (RBFKernelBackward requires 2D) - _lastInput = ShouldCacheForBackward ? processedInput : null; // #1668: skip in inference (arena safety) - // Use Engine.RBFKernel for GPU/CPU acceleration // This computes exp(-epsilon * ||x - center||²) for Gaussian RBF // Convert widths to epsilons: epsilon = 1 / (2 * width²) var epsilons = ComputeEpsilonsFromWidths(); var output = Engine.RBFKernel(processedInput, _centers, epsilons); - // Store 2D output for backward pass - _lastOutput = ShouldCacheForBackward ? output : null; // #1668: skip in inference (arena safety) - // Remove batch dimension if input was unbatched return wasUnbatched ? Engine.Reshape(output, [output.Shape[1]]) : output; } @@ -502,8 +474,6 @@ public override void ClearGradients() public override void ResetState() { // Clear cached values from forward and backward passes - _lastInput = null; - _lastOutput = null; _centersGradient = null; _widthsGradient = null; } diff --git a/src/NeuralNetworks/Layers/RBMLayer.cs b/src/NeuralNetworks/Layers/RBMLayer.cs index 1c2d590fe0..91be216c0f 100644 --- a/src/NeuralNetworks/Layers/RBMLayer.cs +++ b/src/NeuralNetworks/Layers/RBMLayer.cs @@ -594,7 +594,7 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, [1, _visibleUnits]) : Engine.Reshape(input, [flatBatch, _visibleUnits]); - _lastVisibleInput = visible2D; + SaveForBackward(ref _lastVisibleInput, visible2D); // Recompute the transpose every forward — caching it across forwards is // unsafe in a tape-trained network because the tape-based optimizer's @@ -612,7 +612,7 @@ protected override Tensor ForwardTraced(Tensor input) // negligible relative to the FusedLinear that follows it. Tensor weightsT = Engine.TensorTranspose(_weights); Tensor hiddenProbs = Engine.FusedLinear(visible2D, weightsT, _hiddenBiases, FusedActivationType.Sigmoid); - _lastHiddenOutput = hiddenProbs; + SaveForBackward(ref _lastHiddenOutput, hiddenProbs); if (rank == 1) return Engine.Reshape(hiddenProbs, [_hiddenUnits]); diff --git a/src/NeuralNetworks/Layers/RRDBLayer.cs b/src/NeuralNetworks/Layers/RRDBLayer.cs index fb8236373e..6f30722bda 100644 --- a/src/NeuralNetworks/Layers/RRDBLayer.cs +++ b/src/NeuralNetworks/Layers/RRDBLayer.cs @@ -116,18 +116,6 @@ public partial class RRDBLayer : LayerBase, IShapeContract /// private readonly int _growthChannels; - /// - /// Cached input for backpropagation. - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// Cached output from RDB3 for backpropagation. - /// - [AiDotNet.Attributes.Scratch] - private Tensor? _rdb3Output; - #endregion #region Properties @@ -269,14 +257,11 @@ protected override Tensor ForwardTraced(Tensor input) { if (!IsShapeResolved) OnFirstForward(input); - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) - // Pass through 3 Residual Dense Blocks sequentially var x = _rdbBlocks[0].Forward(input); x = _rdbBlocks[1].Forward(x); x = _rdbBlocks[2].Forward(x); - _rdb3Output = ShouldCacheForBackward ? x : null; // #1668: skip in inference (arena safety) // Global residual: output = RDB3_output * residualScale + input return AddResidual(x, input, _residualScale); @@ -342,8 +327,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _rdb3Output = null; foreach (var rdb in _rdbBlocks) { rdb.ResetState(); diff --git a/src/NeuralNetworks/Layers/RRDBNetGenerator.cs b/src/NeuralNetworks/Layers/RRDBNetGenerator.cs index 97ce288cc3..8b75027228 100644 --- a/src/NeuralNetworks/Layers/RRDBNetGenerator.cs +++ b/src/NeuralNetworks/Layers/RRDBNetGenerator.cs @@ -188,28 +188,6 @@ public partial class RRDBNetGenerator : LayerBase, IShapeContract private readonly int _growthChannels; private readonly double _residualScale; - /// - /// Cached input for backpropagation. - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// Cached conv1 output for trunk residual. - /// - [AiDotNet.Attributes.Scratch] - private Tensor? _conv1Output; - - /// - /// Cached intermediate outputs for backpropagation. - /// - private Tensor[]? _rrdbOutputs; - - /// - /// Cached upsampling outputs for backpropagation. - /// - private Tensor[]? _upsampleOutputs; - #endregion #region Properties @@ -442,41 +420,27 @@ protected override Tensor ForwardTraced(Tensor input) { if (!IsShapeResolved) OnFirstForward(input); - // #1668: gate every backward-only cache (the per-stage outputs + intermediate - // arrays). Forward uses the local `conv1Output`/`x`; the fields/arrays only hold - // references for an eager Backward, so in an InferenceMode arena loop nothing - // survives the per-step Reset. - bool cacheBwd = ShouldCacheForBackward; - _lastInput = cacheBwd ? input : null; - // Initial feature extraction var conv1Output = _convFirst.Forward(input); - _conv1Output = cacheBwd ? conv1Output : null; var x = conv1Output; - // Deep feature extraction through RRDB blocks. The array is always allocated - // (cheap — references only) but the per-block tensors are stored only when an - // eager Backward will read them (#1668). - _rrdbOutputs = new Tensor[_rrdbBlocks.Length]; + // Deep feature extraction through RRDB blocks. Only the running activation is + // held; the per-block outputs are not retained past the forward. for (int i = 0; i < _rrdbBlocks.Length; i++) { x = _rrdbBlocks[i].Forward(x); - if (cacheBwd) _rrdbOutputs[i] = x; } // Trunk convolution + global residual var trunk = _trunkConv.Forward(x); x = AddTensors(trunk, conv1Output); // Global residual connection - // Upsampling (array always allocated; tensors stored only for backward — #1668) - _upsampleOutputs = new Tensor[_upsampleConvs.Length * 2]; // Conv output + PixelShuffle output + // Upsampling for (int i = 0; i < _upsampleConvs.Length; i++) { x = _upsampleConvs[i].Forward(x); - if (cacheBwd) _upsampleOutputs[i * 2] = x; x = _pixelShuffleLayers[i].Forward(x); x = ApplyLeakyReLU(x); - if (cacheBwd) _upsampleOutputs[i * 2 + 1] = x; } // HR convolution + activation @@ -611,10 +575,6 @@ private static int SetLayerParameters(ILayer layer, Vector parameters, int /// public override void ResetState() { - _lastInput = null; - _conv1Output = null; - _rrdbOutputs = null; - _upsampleOutputs = null; _convFirst.ResetState(); foreach (var rrdb in _rrdbBlocks) diff --git a/src/NeuralNetworks/Layers/ReadoutLayer.cs b/src/NeuralNetworks/Layers/ReadoutLayer.cs index 66e502b99c..0acbc16e7e 100644 --- a/src/NeuralNetworks/Layers/ReadoutLayer.cs +++ b/src/NeuralNetworks/Layers/ReadoutLayer.cs @@ -128,29 +128,6 @@ public partial class ReadoutLayer : LayerBase, IShapeContract [AiDotNet.Attributes.Scratch] private Tensor _biasGradients; - /// - /// Stores the input tensor from the most recent forward pass for use in backpropagation. - /// - /// - /// This cached input is needed during the backward pass to compute gradients. It holds the - /// input tensor that was processed in the most recent forward pass. The tensor is null - /// before the first forward pass or after a reset. - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// Stores the output tensor (post-activation) from the most recent forward pass for use in backpropagation. - /// - [Scratch] - private Tensor? _lastOutput; - - /// - /// Stores the pre-activation output tensor from the most recent forward pass. - /// - [Scratch] - private Tensor? _lastPreActivation; - private int[] _originalInputShape = []; public override bool SupportsTraining => true; @@ -280,7 +257,6 @@ public ReadoutLayer(int inputSize, int outputSize, IVectorActivationFunction /// protected override Tensor ForwardTraced(Tensor input) { - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) _originalInputShape = input._shape; int inputSize = input.Shape[^1]; @@ -314,14 +290,10 @@ protected override Tensor ForwardTraced(Tensor input) var matmul = Engine.TensorMatMul(flattenedInput, weightsTransposed); var withBias = Engine.TensorAdd(matmul, _bias); - _lastPreActivation = withBias; - Tensor activated = UsingVectorActivation ? (VectorActivation ?? throw new InvalidOperationException("VectorActivation has not been initialized.")).Activate(withBias) : (ScalarActivation ?? throw new InvalidOperationException("ScalarActivation has not been initialized.")).Activate(withBias); - _lastOutput = activated; - if (input.Rank == 1) { return Engine.Reshape(activated, [OutputShape[0]]); @@ -414,9 +386,6 @@ public override void ClearGradients() public override void ResetState() { // Clear cached CPU values from forward pass - _lastInput = null; - _lastOutput = null; - _lastPreActivation = null; _originalInputShape = []; // Reset gradients using Tensor diff --git a/src/NeuralNetworks/Layers/RecurrentLayer.cs b/src/NeuralNetworks/Layers/RecurrentLayer.cs index b4ed6d0746..9f0ecc59f7 100644 --- a/src/NeuralNetworks/Layers/RecurrentLayer.cs +++ b/src/NeuralNetworks/Layers/RecurrentLayer.cs @@ -159,33 +159,11 @@ public partial class RecurrentLayer : LayerBase, IShapeContract /// private bool _isInitialized; - /// - /// Stores the input tensor from the most recent forward pass for use in backpropagation. - /// - /// - /// This cached input is needed during the backward pass to compute gradients. It holds the - /// sequence of input vectors that were processed in the most recent forward pass. The tensor - /// is null before the first forward pass or after a reset. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Stores the original input shape for any-rank tensor support. /// private int[]? _originalInputShape; - /// - /// Stores the output tensor from the most recent forward pass for use in backpropagation. - /// - /// - /// This cached output is needed during the backward pass to compute certain derivatives. - /// It holds the sequence of output vectors that were produced in the most recent forward pass. - /// The tensor is null before the first forward pass or after a reset. - /// - [Scratch] - private Tensor? _lastOutput; - /// /// Stores the gradients of the loss with respect to the input weight parameters. /// @@ -428,12 +406,6 @@ protected override Tensor ForwardTraced(Tensor input) input3D = Engine.Reshape(input, new[] { sequenceLength, flatBatch, inputSize }); } - // Cache input only if training - if (IsTrainingMode) - { - _lastInput = input3D; - } - int hiddenSize = _inputWeights.Shape[0]; int actualInputSize = input3D.Shape[2]; // [seqLen, batch, inputSize] int expectedInputSize = _inputWeights.Shape[1]; @@ -505,8 +477,6 @@ protected override Tensor ForwardTraced(Tensor input) // there to the per-step matmuls + biases. var output = Engine.TensorStack(stepOutputs, axis: 0); - _lastOutput = output; - // Restore original batch dimensions for any-rank support — via Engine // for tape recording. if (_originalInputShape != null && _originalInputShape.Length > 3) @@ -732,8 +702,6 @@ public override void ClearGradients() public override void ResetState() { // Clear cached values from forward and backward passes - _lastInput = null; - _lastOutput = null; _inputWeightsGradient = null; _hiddenWeightsGradient = null; _biasesGradient = null; diff --git a/src/NeuralNetworks/Layers/ReshapeLayer.cs b/src/NeuralNetworks/Layers/ReshapeLayer.cs index 748ced0244..5f2bb20435 100644 --- a/src/NeuralNetworks/Layers/ReshapeLayer.cs +++ b/src/NeuralNetworks/Layers/ReshapeLayer.cs @@ -149,16 +149,6 @@ public partial class ReshapeLayer : LayerBase, IBatchAwareShapeContract private int[] _outputShape; - /// - /// Stores the input tensor from the most recent forward pass for use in backpropagation. - /// - /// - /// This cached input is needed during the backward pass to compute the appropriate gradients. - /// The tensor is null before the first forward pass or after a reset. - /// - [Scratch] - private Tensor? _lastInput; - /// /// Gets a value indicating whether this layer supports training. /// @@ -290,7 +280,6 @@ public int[] GetTargetShape() protected override Tensor ForwardTraced(Tensor input) { EnsureInitializedFromInput(input); - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) int batchSize = input.Shape[0]; int[] targetShape = new int[_outputShape.Length + 1]; targetShape[0] = batchSize; @@ -325,7 +314,6 @@ protected override Tensor ForwardTraced(Tensor input) public override void ResetState() { // Clear cached values from forward pass - _lastInput = null; } } diff --git a/src/NeuralNetworks/Layers/ResidualDenseBlock.cs b/src/NeuralNetworks/Layers/ResidualDenseBlock.cs index f3aa88bd8f..1b95adcc8b 100644 --- a/src/NeuralNetworks/Layers/ResidualDenseBlock.cs +++ b/src/NeuralNetworks/Layers/ResidualDenseBlock.cs @@ -135,27 +135,6 @@ public partial class ResidualDenseBlock : LayerBase, IShapeContract /// private int _inputWidth; - /// - /// Cached input for backpropagation. - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// Cached intermediate conv outputs (before activation) for backpropagation. - /// - private Tensor[]? _convOutputs; - - /// - /// Cached intermediate activation outputs for backpropagation. - /// - private Tensor[]? _activationOutputs; - - /// - /// Cached concatenated inputs to each conv layer for backpropagation. - /// - private Tensor[]? _concatInputs; - #endregion #region Properties @@ -341,46 +320,29 @@ protected override Tensor ForwardTraced(Tensor input) { if (!IsShapeResolved) OnFirstForward(input); - _lastInput = ShouldCacheForBackward ? input : null; // #1668: skip in inference (arena safety) - _convOutputs = new Tensor[5]; - _activationOutputs = new Tensor[4]; // Only first 4 have activation - _concatInputs = new Tensor[5]; - + // Every intermediate is a local: the tape keeps what backward needs, and fields would hold + // all five convolutions' activations after the forward, in inference too. // x0 = input var x0 = input; // Conv1: x1 = LeakyReLU(Conv(x0)) - _concatInputs[0] = x0; - _convOutputs[0] = _convLayers[0].Forward(x0); - _activationOutputs[0] = ApplyLeakyReLU(_convOutputs[0]); - var x1 = _activationOutputs[0]; + var x1 = ApplyLeakyReLU(_convLayers[0].Forward(x0)); // Conv2: x2 = LeakyReLU(Conv(concat(x0, x1))) var concat1 = ConcatenateChannels(x0, x1); - _concatInputs[1] = concat1; - _convOutputs[1] = _convLayers[1].Forward(concat1); - _activationOutputs[1] = ApplyLeakyReLU(_convOutputs[1]); - var x2 = _activationOutputs[1]; + var x2 = ApplyLeakyReLU(_convLayers[1].Forward(concat1)); // Conv3: x3 = LeakyReLU(Conv(concat(x0, x1, x2))) var concat2 = ConcatenateChannels(concat1, x2); - _concatInputs[2] = concat2; - _convOutputs[2] = _convLayers[2].Forward(concat2); - _activationOutputs[2] = ApplyLeakyReLU(_convOutputs[2]); - var x3 = _activationOutputs[2]; + var x3 = ApplyLeakyReLU(_convLayers[2].Forward(concat2)); // Conv4: x4 = LeakyReLU(Conv(concat(x0, x1, x2, x3))) var concat3 = ConcatenateChannels(concat2, x3); - _concatInputs[3] = concat3; - _convOutputs[3] = _convLayers[3].Forward(concat3); - _activationOutputs[3] = ApplyLeakyReLU(_convOutputs[3]); - var x4 = _activationOutputs[3]; + var x4 = ApplyLeakyReLU(_convLayers[3].Forward(concat3)); // Conv5: x5 = Conv(concat(x0, x1, x2, x3, x4)) - NO activation var concat4 = ConcatenateChannels(concat3, x4); - _concatInputs[4] = concat4; - _convOutputs[4] = _convLayers[4].Forward(concat4); - var x5 = _convOutputs[4]; + var x5 = _convLayers[4].Forward(concat4); // Local residual learning: output = x5 * residualScale + x0 return AddResidual(x5, x0, _residualScale); @@ -566,10 +528,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _convOutputs = null; - _activationOutputs = null; - _concatInputs = null; foreach (var conv in _convLayers) { diff --git a/src/NeuralNetworks/Layers/ResidualLayer.cs b/src/NeuralNetworks/Layers/ResidualLayer.cs index 1a71e6e08d..7a7fb158fe 100644 --- a/src/NeuralNetworks/Layers/ResidualLayer.cs +++ b/src/NeuralNetworks/Layers/ResidualLayer.cs @@ -65,36 +65,6 @@ public partial class ResidualLayer : LayerBase, IShapeContract /// private ILayer? _innerLayer; - /// - /// Stores the input tensor from the most recent forward pass. - /// - /// - /// - /// This field caches the input to the layer during the forward pass, which is needed during the backward pass - /// to compute gradients. It is cleared when ResetState() is called. - /// - /// For Beginners: This is like the layer's short-term memory of what input it received. - /// - /// During training, the layer needs to remember what input it processed so that it can - /// properly calculate how to improve. This temporary storage is cleared between batches - /// or when you explicitly reset the layer. - /// - /// - [Scratch] - private Tensor? _lastInput; - - /// - /// Stores the output tensor from the inner layer during the forward pass. - /// - /// - /// - /// This field caches the inner layer's output during the forward pass, which is needed during the backward pass - /// to avoid recomputing and potentially corrupting the inner layer's state. It is cleared when ResetState() is called. - /// - /// - [Scratch] - private Tensor? _lastInnerOutput; - /// /// Gets a value indicating whether this layer supports training through backpropagation. /// @@ -365,9 +335,6 @@ public void SetInnerLayer(ILayer innerLayer) protected override Tensor ForwardTraced(Tensor input) { EnsureInitializedFromInput(input); - // #1668: gate backward caches; forward uses the local innerOutput (chained). - bool cacheBwd = ShouldCacheForBackward; - _lastInput = cacheBwd ? input : null; var innerOutput = _innerLayer?.Forward(input); // A residual block's inner branch is shape-preserving by contract (out = x + F(x)), so the @@ -383,7 +350,6 @@ protected override Tensor ForwardTraced(Tensor input) innerOutput = Engine.Reshape(innerOutput, input._shape); } - _lastInnerOutput = cacheBwd ? innerOutput : null; var result = innerOutput == null ? input : Engine.TensorAdd(input, innerOutput); return ApplyActivation(result); @@ -520,8 +486,6 @@ internal override Dictionary GetMetadata() /// public override void ResetState() { - _lastInput = null; - _lastInnerOutput = null; _innerLayer?.ResetState(); } diff --git a/src/NeuralNetworks/Layers/SSM/ABCLayer.cs b/src/NeuralNetworks/Layers/SSM/ABCLayer.cs index bd400c8709..63306f8436 100644 --- a/src/NeuralNetworks/Layers/SSM/ABCLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/ABCLayer.cs @@ -132,23 +132,6 @@ public partial class ABCLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached forward pass values - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastForgetGate; - private Tensor? _lastOutputGate; - [Scratch] - private Tensor? _lastOutputGateRaw; - [Scratch] - private Tensor? _lastSlotReadOutput; private int[]? _originalInputShape; // Gradients @@ -339,31 +322,23 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Q, K, V projections var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); var q = Engine.Reshape(Engine.TensorMatMul(inputFlat, _queryWeights), new[] { batchSize, seqLen, _modelDimension }); var k = Engine.Reshape(Engine.TensorMatMul(inputFlat, _keyWeights), new[] { batchSize, seqLen, _modelDimension }); var v = Engine.Reshape(Engine.TensorMatMul(inputFlat, _valueWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; // Step 2: Forget gate var forgetRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _forgetGateWeights), Engine.Reshape(_forgetGateBias, new[] { 1, _numHeads })), new[] { batchSize, seqLen, _numHeads }); var forgetGate = Engine.Sigmoid(forgetRaw); - _lastForgetGate = forgetGate; // Step 3: Output gate var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var outputGate = Engine.Swish(gateRaw); - _lastOutputGate = outputGate; - _lastOutputGateRaw = gateRaw; // Step 4: Slot competition, write, and read. // One fused, differentiable op for the whole recurrence. The scalar loop this replaces was @@ -371,7 +346,6 @@ protected override Tensor ForwardTraced(Tensor input) // never learned - this layer has no Backward override, so the tape is its only gradient path. var slotOutput = Engine.AbcScanForward( q, k, v, forgetGate, _slotKeys, _numHeads, SlotInitScale); - _lastSlotReadOutput = slotOutput; // Step 5: Gated output var gatedOutput = Engine.TensorMultiply(slotOutput, outputGate); @@ -384,7 +358,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -475,15 +448,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastForgetGate = null; - _lastOutputGate = null; - _lastOutputGateRaw = null; - _lastSlotReadOutput = null; _originalInputShape = null; _queryWeightsGradient = null; _keyWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/BASEDLayer.cs b/src/NeuralNetworks/Layers/SSM/BASEDLayer.cs index 4826cbe61a..99a6a714b1 100644 --- a/src/NeuralNetworks/Layers/SSM/BASEDLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/BASEDLayer.cs @@ -131,36 +131,10 @@ public partial class BASEDLayer : LayerBase, IShapeContract // Cached forward pass values [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastLinearQuery; - [Scratch] - private Tensor? _lastLinearKey; - [Scratch] - private Tensor? _lastLinearValue; - [Scratch] - private Tensor? _lastWindowQuery; - [Scratch] - private Tensor? _lastWindowKey; - [Scratch] - private Tensor? _lastWindowValue; - [Scratch] private Tensor? _lastLinearFeatureQ; [Scratch] private Tensor? _lastLinearFeatureK; [Scratch] - private Tensor? _lastLinearOutput; - [Scratch] - private Tensor? _lastWindowOutput; - [Scratch] - private Tensor? _lastMixingAlpha; - [Scratch] - private Tensor? _lastMixingAlphaRaw; - [Scratch] - private Tensor? _lastCombinedOutput; - [Scratch] private Tensor? _lastWindowScores; private int[]? _originalInputShape; @@ -355,8 +329,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Compute projections for both attention mechanisms var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); @@ -364,25 +336,17 @@ protected override Tensor ForwardTraced(Tensor input) var linQ = Engine.Reshape(Engine.TensorMatMul(inputFlat, _linearQueryWeights), new[] { batchSize, seqLen, _modelDimension }); var linK = Engine.Reshape(Engine.TensorMatMul(inputFlat, _linearKeyWeights), new[] { batchSize, seqLen, _modelDimension }); var linV = Engine.Reshape(Engine.TensorMatMul(inputFlat, _linearValueWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastLinearQuery = linQ; - _lastLinearKey = linK; - _lastLinearValue = linV; // Window attention Q, K, V var winQ = Engine.Reshape(Engine.TensorMatMul(inputFlat, _windowQueryWeights), new[] { batchSize, seqLen, _modelDimension }); var winK = Engine.Reshape(Engine.TensorMatMul(inputFlat, _windowKeyWeights), new[] { batchSize, seqLen, _modelDimension }); var winV = Engine.Reshape(Engine.TensorMatMul(inputFlat, _windowValueWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastWindowQuery = winQ; - _lastWindowKey = winK; - _lastWindowValue = winV; // Step 2: Compute mixing gate alpha var alphaRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _mixingGateWeights), Engine.Reshape(_mixingGateBias, new[] { 1, _numHeads })), new[] { batchSize, seqLen, _numHeads }); var alpha = Engine.Sigmoid(alphaRaw); - _lastMixingAlphaRaw = alphaRaw; - _lastMixingAlpha = alpha; // Step 3: Linear attention with a vectorized Taylor feature map. var qHeads = Engine.Reshape( @@ -408,12 +372,10 @@ protected override Tensor ForwardTraced(Tensor input) CausalLinearAttention.NormalizedHeads( Engine, phiQ, phiK, vHeads, NumOps.FromDouble(1e-6)), new[] { batchSize, seqLen, _modelDimension }); - _lastLinearOutput = linearOutput; // Step 4: Sliding window causal attention var windowOutput = CausalLinearAttention.ScaledDotProduct( Engine, winQ, winK, winV, _numHeads, causal: true, windowSize: _windowSize); - _lastWindowOutput = windowOutput; // Step 5: Combine linear and window attention using learned alpha var alphaHeads = Engine.TensorBroadcastTo( @@ -429,7 +391,6 @@ protected override Tensor ForwardTraced(Tensor input) Engine.TensorMultiply( Engine.TensorSubtract(ones, alphaModel), windowOutput)); - _lastCombinedOutput = combined; // Step 6: Output projection var combinedFlat = Engine.Reshape(combined, new[] { batchSize * seqLen, _modelDimension }); @@ -439,7 +400,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -802,21 +762,8 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastLinearQuery = null; - _lastLinearKey = null; - _lastLinearValue = null; - _lastWindowQuery = null; - _lastWindowKey = null; - _lastWindowValue = null; _lastLinearFeatureQ = null; _lastLinearFeatureK = null; - _lastLinearOutput = null; - _lastWindowOutput = null; - _lastMixingAlpha = null; - _lastMixingAlphaRaw = null; - _lastCombinedOutput = null; _lastWindowScores = null; _originalInputShape = null; _linearQueryWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/DeltaFormerLayer.cs b/src/NeuralNetworks/Layers/SSM/DeltaFormerLayer.cs index cf37cb1aff..af381334ff 100644 --- a/src/NeuralNetworks/Layers/SSM/DeltaFormerLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/DeltaFormerLayer.cs @@ -105,26 +105,6 @@ public partial class DeltaFormerLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastGate; - [Scratch] - private Tensor? _lastGateRaw; - [Scratch] - private Tensor? _lastAttentionWeights; - [Scratch] - private Tensor? _lastStates; - [Scratch] - private Tensor? _lastMechanismOutput; private int[]? _originalInputShape; // Gradients @@ -256,24 +236,17 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Q, K, V projections var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); var q = Engine.Reshape(Engine.TensorMatMul(inputFlat, _queryWeights), new[] { batchSize, seqLen, _modelDimension }); var k = Engine.Reshape(Engine.TensorMatMul(inputFlat, _keyWeights), new[] { batchSize, seqLen, _modelDimension }); var v = Engine.Reshape(Engine.TensorMatMul(inputFlat, _valueWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; // Step 2: Output gate var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var gate = Engine.Sigmoid(gateRaw); - _lastGate = gate; - _lastGateRaw = gateRaw; // Step 3: Either delta rule or standard attention Tensor mechanismOutput; @@ -285,7 +258,6 @@ protected override Tensor ForwardTraced(Tensor input) { mechanismOutput = SoftmaxAttentionForward(q, k, v, batchSize, seqLen); } - _lastMechanismOutput = mechanismOutput; // Step 4: Gated output var gatedOutput = Engine.TensorMultiply(mechanismOutput, gate); @@ -298,7 +270,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -326,8 +297,6 @@ private Tensor SoftmaxAttentionForward( Engine.BatchMatMul(qHeads, Engine.TensorPermute(kHeads, new[] { 0, 2, 1 })), NumOps.FromDouble(1.0 / Math.Sqrt(_headDimension))); var weights = Engine.Softmax(scores, axis: -1); - _lastAttentionWeights = Engine.Reshape(weights, - new[] { batchSize, _numHeads, seqLen, seqLen }); return FromHeadMajor(Engine.BatchMatMul(weights, vHeads), batchSize, seqLen); } @@ -363,7 +332,6 @@ private Tensor DeltaRuleForward( new[] { headBatch, 1, _headDimension })); } - _lastStates = state; return FromHeadMajor(Engine.TensorConcatenate(outputs.ToArray(), 1), batchSize, seqLen); } @@ -433,16 +401,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastGate = null; - _lastGateRaw = null; - _lastAttentionWeights = null; - _lastStates = null; - _lastMechanismOutput = null; _originalInputShape = null; _queryWeightsGradient = null; _keyWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/DeltaNetLayer.cs b/src/NeuralNetworks/Layers/SSM/DeltaNetLayer.cs index 8bb883d38f..82c130f5b8 100644 --- a/src/NeuralNetworks/Layers/SSM/DeltaNetLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/DeltaNetLayer.cs @@ -116,22 +116,6 @@ public partial class DeltaNetLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastBeta; - [Scratch] - private Tensor? _lastStates; - [Scratch] - private Tensor? _lastDeltaRuleOutput; private int[]? _originalInputShape; // Gradients @@ -282,8 +266,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Q, K, V projections var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); @@ -302,20 +284,14 @@ protected override Tensor ForwardTraced(Tensor input) Engine.Reshape(_valueBias, new[] { 1, _modelDimension })); var v = Engine.Reshape(vFlat, new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; - // Step 2: Beta (write strength) via sigmoid var betaRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _betaWeights), Engine.Reshape(_betaBias, new[] { 1, _numHeads })), new[] { batchSize, seqLen, _numHeads }); var beta = Engine.Sigmoid(betaRaw); - _lastBeta = beta; // Step 3: Delta rule recurrence per head var output = DeltaRuleForward(q, k, v, beta, batchSize, seqLen); - _lastDeltaRuleOutput = output; // Step 4: Output projection var outputFlat = Engine.TensorMatMul( @@ -326,7 +302,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -397,7 +372,6 @@ private Tensor DeltaRuleForward( outputs.Add(Engine.Reshape(y, new[] { headBatch, 1, _headDimension })); } - _lastStates = state; return FromHeadMajor(Engine.TensorConcatenate(outputs.ToArray(), 1), batchSize, seqLen); } @@ -481,14 +455,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastBeta = null; - _lastStates = null; - _lastDeltaRuleOutput = null; _originalInputShape = null; _queryWeightsGradient = null; _queryBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/DeltaProductLayer.cs b/src/NeuralNetworks/Layers/SSM/DeltaProductLayer.cs index f6032a9c3e..05fffd430d 100644 --- a/src/NeuralNetworks/Layers/SSM/DeltaProductLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/DeltaProductLayer.cs @@ -115,24 +115,6 @@ public partial class DeltaProductLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached forward pass values - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastBeta; - [Scratch] - private Tensor? _lastHouseholderVecs; - [Scratch] - private Tensor? _lastStates; - [Scratch] - private Tensor? _lastRecurrenceOutput; private int[]? _originalInputShape; // Gradients @@ -283,32 +265,25 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); // Step 1: Q, K, V projections var q = Engine.Reshape(Engine.TensorMatMul(inputFlat, _queryWeights), new[] { batchSize, seqLen, _modelDimension }); var k = Engine.Reshape(Engine.TensorMatMul(inputFlat, _keyWeights), new[] { batchSize, seqLen, _modelDimension }); var v = Engine.Reshape(Engine.TensorMatMul(inputFlat, _valueWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; // Step 2: Beta (write strength) var betaRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _betaWeights), Engine.Reshape(_betaBias, new[] { 1, _numHeads })), new[] { batchSize, seqLen, _numHeads }); var beta = Engine.Sigmoid(betaRaw); - _lastBeta = beta; // Step 3: Compute Householder vectors per timestep // householderVecs: [batch, seqLen, numHouseholders, numHeads, headDim] var hVecs = ComputeHouseholderVectors(inputFlat, batchSize, seqLen); - _lastHouseholderVecs = hVecs; // Step 4: DeltaProduct recurrence var recOutput = DeltaProductRecurrence(q, k, v, beta, hVecs, batchSize, seqLen); - _lastRecurrenceOutput = recOutput; // Step 5: Output projection var outputFlat = Engine.TensorMatMul( @@ -318,7 +293,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -465,7 +439,6 @@ private Tensor DeltaProductRecurrence( new[] { headBatch, 1, _headDimension })); } - _lastStates = state; return FromHeadMajor(Engine.TensorConcatenate(outputs.ToArray(), 1), batchSize, seqLen); } @@ -604,15 +577,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastBeta = null; - _lastHouseholderVecs = null; - _lastStates = null; - _lastRecurrenceOutput = null; _originalInputShape = null; _queryWeightsGradient = null; _keyWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/ExtendedLSTMLayer.cs b/src/NeuralNetworks/Layers/SSM/ExtendedLSTMLayer.cs index 603340b62d..d078793647 100644 --- a/src/NeuralNetworks/Layers/SSM/ExtendedLSTMLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/ExtendedLSTMLayer.cs @@ -126,24 +126,6 @@ public partial class ExtendedLSTMLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastInputGates; - [Scratch] - private Tensor? _lastForgetGates; - [Scratch] - private Tensor? _lastOutputGates; - [Scratch] - private Tensor? _lastQ; - [Scratch] - private Tensor? _lastK; - [Scratch] - private Tensor? _lastV; - [Scratch] - private Tensor? _lastHiddenPreProj; private int[]? _originalInputShape; // Fixed (non-trainable) unit-gamma / zero-beta for the mLSTM output @@ -300,8 +282,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Project the full sequence once. The matrix-memory scan accepts scalar gates per head; // preserve this layer's historical first-coordinate selection and log stabilization. var flat = Engine.Reshape(input3D, new[] { batchSize * seqLen, modelDim }); @@ -356,15 +336,7 @@ Tensor HeadGate(Tensor gate) => Engine.Reshape(Engine.TensorNarrow( } output = Engine.LayerNorm(output, _outputNormGamma, _outputNormBeta, 1e-5, out _, out _); - // The scan owns recurrent caches; retain only the projection caches used by the layer. - _lastInputGates = inputLogits; - _lastForgetGates = forget; - _lastOutputGates = outputGate; - _lastQ = q; _lastK = k; _lastV = v; - _lastHiddenPreProj = hidden; - var result = ApplyActivation(output); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -435,15 +407,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastInputGates = null; - _lastForgetGates = null; - _lastOutputGates = null; - _lastQ = null; - _lastK = null; - _lastV = null; - _lastHiddenPreProj = null; _originalInputShape = null; _inputGateWeightsGradient = null; _inputGateBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/GatedDeltaNetLayer.cs b/src/NeuralNetworks/Layers/SSM/GatedDeltaNetLayer.cs index aeb55f3039..793b905ddd 100644 --- a/src/NeuralNetworks/Layers/SSM/GatedDeltaNetLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/GatedDeltaNetLayer.cs @@ -129,30 +129,6 @@ public partial class GatedDeltaNetLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastConvOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastBeta; - [Scratch] - private Tensor? _lastAlpha; - [Scratch] - private Tensor? _lastGate; - [Scratch] - private Tensor? _lastGateRaw; - [Scratch] - private Tensor? _lastSiluConv; - [Scratch] - private Tensor? _lastDeltaRuleOutput; private int[]? _originalInputShape; // Gradients @@ -335,46 +311,34 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Short convolution var convOutput = DepthwiseConv1DForward(input3D, batchSize, seqLen); var siluConv = Engine.Swish(convOutput); - _lastConvOutput = convOutput; - _lastSiluConv = siluConv; // Step 2: Q, K, V projections var siluFlat = Engine.Reshape(siluConv, new[] { batchSize * seqLen, _modelDimension }); var q = Engine.Reshape(Engine.TensorMatMul(siluFlat, _queryWeights), new[] { batchSize, seqLen, _modelDimension }); var k = Engine.Reshape(Engine.TensorMatMul(siluFlat, _keyWeights), new[] { batchSize, seqLen, _modelDimension }); var v = Engine.Reshape(Engine.TensorMatMul(siluFlat, _valueWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; // Step 3: Gates var betaRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(siluFlat, _betaWeights), Engine.Reshape(_betaBias, new[] { 1, _numHeads })), new[] { batchSize, seqLen, _numHeads }); var beta = Engine.Sigmoid(betaRaw); - _lastBeta = beta; var alphaRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(siluFlat, _alphaWeights), Engine.Reshape(_alphaBias, new[] { 1, _numHeads })), new[] { batchSize, seqLen, _numHeads }); var alpha = Engine.Sigmoid(alphaRaw); - _lastAlpha = alpha; var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(siluFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var gate = Engine.Swish(gateRaw); - _lastGate = gate; - _lastGateRaw = gateRaw; // Step 4: Delta rule recurrence per head var output = Engine.GatedDeltaNetScanForward(q, k, v, alpha, beta, _numHeads); - _lastDeltaRuleOutput = output; // Step 5: Gated output var gatedOutput = Engine.TensorMultiply(output, gate); @@ -387,7 +351,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -504,18 +467,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastConvOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastBeta = null; - _lastAlpha = null; - _lastGate = null; - _lastGateRaw = null; - _lastSiluConv = null; - _lastDeltaRuleOutput = null; _originalInputShape = null; _convWeightsGradient = null; _convBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/GatedDeltaProductLayer.cs b/src/NeuralNetworks/Layers/SSM/GatedDeltaProductLayer.cs index b83d6f9505..e3a53cd878 100644 --- a/src/NeuralNetworks/Layers/SSM/GatedDeltaProductLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/GatedDeltaProductLayer.cs @@ -134,28 +134,6 @@ public partial class GatedDeltaProductLayer : LayerBase, IShapeContract // Cached forward pass values [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastBeta; - [Scratch] - private Tensor? _lastAlpha; - [Scratch] - private Tensor? _lastHouseholderVecs; - [Scratch] - private Tensor? _lastRecurrenceOutput; - [Scratch] - private Tensor? _lastOutputGate; - [Scratch] - private Tensor? _lastOutputGateRaw; - [Scratch] private readonly List> _recurrenceInitialStates = []; private int[]? _originalInputShape; @@ -323,16 +301,12 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); // Step 1: Q, K, V projections var q = Engine.Reshape(Engine.TensorMatMul(inputFlat, _queryWeights), new[] { batchSize, seqLen, _modelDimension }); var k = Engine.Reshape(Engine.TensorMatMul(inputFlat, _keyWeights), new[] { batchSize, seqLen, _modelDimension }); var v = Engine.Reshape(Engine.TensorMatMul(inputFlat, _valueWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; // Step 2: Gates var betaProjection = Engine.TensorMatMul(inputFlat, _betaWeights); @@ -342,7 +316,6 @@ protected override Tensor ForwardTraced(Tensor input) var betaRaw = Engine.Reshape( Engine.TensorAdd(betaProjection, betaBias), new[] { batchSize, seqLen, _numHeads }); var beta = Engine.Sigmoid(betaRaw); - _lastBeta = beta; var alphaProjection = Engine.TensorMatMul(inputFlat, _alphaWeights); var alphaBias = Engine.TensorBroadcastTo( @@ -351,7 +324,6 @@ protected override Tensor ForwardTraced(Tensor input) var alphaRaw = Engine.Reshape( Engine.TensorAdd(alphaProjection, alphaBias), new[] { batchSize, seqLen, _numHeads }); var alpha = Engine.Sigmoid(alphaRaw); - _lastAlpha = alpha; var gateProjection = Engine.TensorMatMul(inputFlat, _outputGateWeights); var gateBias = Engine.TensorBroadcastTo( @@ -360,16 +332,12 @@ protected override Tensor ForwardTraced(Tensor input) var gateRaw = Engine.Reshape( Engine.TensorAdd(gateProjection, gateBias), new[] { batchSize, seqLen, _modelDimension }); var gate = Engine.Swish(gateRaw); - _lastOutputGate = gate; - _lastOutputGateRaw = gateRaw; // Step 3: Compute Householder vectors var hVecs = ComputeHouseholderVectors(inputFlat, batchSize, seqLen); - _lastHouseholderVecs = hVecs; // Step 4: Gated DeltaProduct recurrence var recOutput = GatedDeltaProductRecurrence(q, k, v, alpha, beta, hVecs, batchSize, seqLen); - _lastRecurrenceOutput = recOutput; // Step 5: Gated output var gatedOutput = Engine.TensorMultiply(recOutput, gate); @@ -384,7 +352,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -654,17 +621,6 @@ public override void ClearGradients() public override void ResetState() { DisposeRecurrenceInitialStates(); - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastBeta = null; - _lastAlpha = null; - _lastHouseholderVecs = null; - _lastRecurrenceOutput = null; - _lastOutputGate = null; - _lastOutputGateRaw = null; _originalInputShape = null; _queryWeightsGradient = null; _keyWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/GatedLinearAttentionLayer.cs b/src/NeuralNetworks/Layers/SSM/GatedLinearAttentionLayer.cs index ec4bd3c7d6..7896d4f31a 100644 --- a/src/NeuralNetworks/Layers/SSM/GatedLinearAttentionLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/GatedLinearAttentionLayer.cs @@ -105,20 +105,6 @@ public partial class GatedLinearAttentionLayer : LayerBase, IShapeContract private const double ResidualNormEpsilon = 1e-5; // Cached values - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastGate; - [Scratch] - private Tensor? _lastAttnOutput; // Pre-output-projection attention output private int[]? _originalInputShape; // Gradients @@ -251,8 +237,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - int totalDim = _numHeads * _headDimension; // Project Q, K, V, G @@ -266,11 +250,6 @@ protected override Tensor ForwardTraced(Tensor input) gFlat = Engine.TensorAdd(gFlat, gBias); var gate = Engine.Sigmoid(Engine.Reshape(gFlat, new[] { batchSize, seqLen, totalDim })); - _lastQuery = q; - _lastKey = k; - _lastValue = v; - _lastGate = gate; - // The recurrence consumes one scalar gate per head. Extract the first channel from each // head exactly as the original implementation did, but keep the extraction and scan on the // tape. GlaScanForward records one analytic BPTT node instead of detaching through scalar @@ -284,8 +263,6 @@ protected override Tensor ForwardTraced(Tensor input) new[] { batchSize, seqLen, _numHeads }); var output = Engine.GlaScanForward(q, k, v, gateByHead, _numHeads); - _lastAttnOutput = output; - // Output projection var outFlat = Engine.Reshape(output, new[] { batchSize * seqLen, totalDim }); var outputFlat = Engine.TensorMatMul(outFlat, _outputWeights); @@ -307,7 +284,6 @@ protected override Tensor ForwardTraced(Tensor input) _residualNormGamma, _residualNormBeta, ResidualNormEpsilon, out _, out _); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -370,13 +346,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastGate = null; - _lastAttnOutput = null; _originalInputShape = null; _queryWeightsGradient = null; _keyWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/GatedSlotAttentionLayer.cs b/src/NeuralNetworks/Layers/SSM/GatedSlotAttentionLayer.cs index b52f4da471..22d799e80e 100644 --- a/src/NeuralNetworks/Layers/SSM/GatedSlotAttentionLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/GatedSlotAttentionLayer.cs @@ -128,28 +128,6 @@ public partial class GatedSlotAttentionLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastForgetGate; - [Scratch] - private Tensor? _lastInputGate; - [Scratch] - private Tensor? _lastOutputGate; - [Scratch] - private Tensor? _lastOutputGateRaw; - [Scratch] - private Tensor? _lastSlotStates; - [Scratch] - private Tensor? _lastSlotReadOutput; private int[]? _originalInputShape; // Gradients @@ -342,36 +320,27 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Q, K, V projections var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); var q = Engine.Reshape(Engine.TensorMatMul(inputFlat, _queryWeights), new[] { batchSize, seqLen, _modelDimension }); var k = Engine.Reshape(Engine.TensorMatMul(inputFlat, _keyWeights), new[] { batchSize, seqLen, _modelDimension }); var v = Engine.Reshape(Engine.TensorMatMul(inputFlat, _valueWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; // Step 2: Gates var forgetRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _forgetGateWeights), Engine.Reshape(_forgetGateBias, new[] { 1, _numHeads })), new[] { batchSize, seqLen, _numHeads }); var forgetGate = Engine.Sigmoid(forgetRaw); - _lastForgetGate = forgetGate; var inputGateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _inputGateWeights), Engine.Reshape(_inputGateBias, new[] { 1, _numHeads })), new[] { batchSize, seqLen, _numHeads }); var inputGate = Engine.Sigmoid(inputGateRaw); - _lastInputGate = inputGate; var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var outputGate = Engine.Swish(gateRaw); - _lastOutputGate = outputGate; - _lastOutputGateRaw = gateRaw; // Step 3: Slot update recurrence and slot read // The modular slot mapping means all slots with the same key index @@ -427,7 +396,6 @@ protected override Tensor ForwardTraced(Tensor input) Engine.TensorMultiply(initialRead, Engine.TensorExpandDims(decay, axis: 3)), new[] { batchSize, seqLen, _modelDimension }); var slotOutput = Engine.TensorAdd(dynamicOutput, decayedInitialRead); - _lastSlotReadOutput = slotOutput; // Step 4: Gated output var gatedOutput = Engine.TensorMultiply(slotOutput, outputGate); @@ -440,7 +408,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -557,7 +524,6 @@ private Tensor SlotRecurrenceForward( allStates[new[] { bi, t + 1, hi, si, di }] = slotState[new[] { bi, hi, si, di }]; } - _lastSlotStates = allStates; return output; } @@ -623,17 +589,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastForgetGate = null; - _lastInputGate = null; - _lastOutputGate = null; - _lastOutputGateRaw = null; - _lastSlotStates = null; - _lastSlotReadOutput = null; _originalInputShape = null; _queryWeightsGradient = null; _keyWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/HGRN2Layer.cs b/src/NeuralNetworks/Layers/SSM/HGRN2Layer.cs index a09c8b1112..37021ff841 100644 --- a/src/NeuralNetworks/Layers/SSM/HGRN2Layer.cs +++ b/src/NeuralNetworks/Layers/SSM/HGRN2Layer.cs @@ -121,24 +121,6 @@ public partial class HGRN2Layer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastForgetGate; - [Scratch] - private Tensor? _lastGate; - [Scratch] - private Tensor? _lastGateRaw; - [Scratch] - private Tensor? _lastRecurrenceOutput; private int[]? _originalInputShape; // Gradients @@ -287,35 +269,26 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Q, K, V projections var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); var q = Engine.Reshape(Engine.TensorMatMul(inputFlat, _queryWeights), new[] { batchSize, seqLen, _modelDimension }); var k = Engine.Reshape(Engine.TensorMatMul(inputFlat, _keyWeights), new[] { batchSize, seqLen, _modelDimension }); var v = Engine.Reshape(Engine.TensorMatMul(inputFlat, _valueWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; // Step 2: Forget gate g_t = sigmoid(W_g * x_t + b_g), per-head scalar var forgetRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _forgetGateWeights), Engine.Reshape(_forgetGateBias, new[] { 1, _numHeads })), new[] { batchSize, seqLen, _numHeads }); var forgetGate = Engine.Sigmoid(forgetRaw); - _lastForgetGate = forgetGate; // Step 3: Output gate = swish(W_gate * x + b_gate) var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var gate = Engine.Swish(gateRaw); - _lastGate = gate; - _lastGateRaw = gateRaw; // Step 4: Gated outer-product recurrence per head var recurrenceOutput = OuterProductRecurrenceForward(q, k, v, forgetGate, batchSize, seqLen); - _lastRecurrenceOutput = recurrenceOutput; // Step 5: Gated output var gatedOutput = Engine.TensorMultiply(recurrenceOutput, gate); @@ -328,7 +301,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -436,15 +408,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastForgetGate = null; - _lastGate = null; - _lastGateRaw = null; - _lastRecurrenceOutput = null; _originalInputShape = null; _queryWeightsGradient = null; _keyWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/HGRNLayer.cs b/src/NeuralNetworks/Layers/SSM/HGRNLayer.cs index f89ed990f1..0b54ddd3a1 100644 --- a/src/NeuralNetworks/Layers/SSM/HGRNLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/HGRNLayer.cs @@ -107,20 +107,6 @@ public partial class HGRNLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastProjectedInput; - [Scratch] - private Tensor? _lastForgetGate; - [Scratch] - private Tensor? _lastInputGate; - [Scratch] - private Tensor? _lastHiddenStates; - [Scratch] - private Tensor? _lastRecurrenceOutput; private int[]? _originalInputShape; // Gradients @@ -265,15 +251,12 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Input projection var input2D = Engine.Reshape(input3D, new[] { batchSize * seqLen, modelDim }); var projected = Engine.TensorMatMul(input2D, _inputProjectionWeights); var projBias = Engine.Reshape(_inputProjectionBias, new[] { 1, _modelDimension }); projected = Engine.TensorAdd(projected, projBias); var projected3D = Engine.Reshape(projected, new[] { batchSize, seqLen, _modelDimension }); - _lastProjectedInput = projected3D; // Step 2: Compute forget and input gates from projected input var fBias = Engine.Reshape(_forgetGateBias, new[] { 1, _modelDimension }); @@ -287,14 +270,10 @@ protected override Tensor ForwardTraced(Tensor input) Engine.TensorMatMul(projected, _inputGateWeights), iBias)), new[] { batchSize, seqLen, _modelDimension }); - _lastForgetGate = forgetGate3D; - _lastInputGate = inputGate3D; - // Step 3: Element-wise gated recurrence // h_t = f_t * h_{t-1} + i_t * x_t (where x_t is the projected input) var recurrenceOutput = GatedRecurrenceForward( projected3D, forgetGate3D, inputGate3D, batchSize, seqLen); - _lastRecurrenceOutput = recurrenceOutput; // Step 4: Output projection var outFlat = Engine.Reshape(recurrenceOutput, new[] { batchSize * seqLen, _modelDimension }); @@ -304,7 +283,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -354,10 +332,6 @@ private Tensor GatedRecurrenceForward( var output = hiddenSteps.Count == 1 ? hiddenSteps[0] : Engine.TensorConcatenate(hiddenSteps.ToArray(), axis: 1); - var initialState = Tensor.CreateDefault([batchSize, 1, _modelDimension], NumOps.Zero); - _lastHiddenStates = Engine.TensorConcatenate( - [initialState, output], - axis: 1); return output; } @@ -426,13 +400,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastProjectedInput = null; - _lastForgetGate = null; - _lastInputGate = null; - _lastHiddenStates = null; - _lastRecurrenceOutput = null; _originalInputShape = null; _inputProjectionWeightsGradient = null; _inputProjectionBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/HedgehogLayer.cs b/src/NeuralNetworks/Layers/SSM/HedgehogLayer.cs index ac550cea1b..9c289adc0d 100644 --- a/src/NeuralNetworks/Layers/SSM/HedgehogLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/HedgehogLayer.cs @@ -121,21 +121,6 @@ public partial class HedgehogLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached forward pass values - private Tensor? _lastInput; - private Tensor? _lastOutput; - private Tensor? _lastQuery; - private Tensor? _lastKey; - private Tensor? _lastValue; - private Tensor? _lastPhiQ; - private Tensor? _lastPhiK; - private Tensor? _lastPhiQHidden; - private Tensor? _lastPhiKHidden; - private Tensor? _lastPhiQPreActivation; - private Tensor? _lastPhiKPreActivation; - private Tensor? _lastAttnOutput; - private Tensor? _lastAttnDenominators; - private Tensor? _lastGateRaw; - private Tensor? _lastGate; private int[]? _originalInputShape; // Gradients @@ -333,41 +318,26 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Q, K, V projections var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); var q = Engine.Reshape(Engine.TensorMatMul(inputFlat, _queryWeights), new[] { batchSize, seqLen, _modelDimension }); var k = Engine.Reshape(Engine.TensorMatMul(inputFlat, _keyWeights), new[] { batchSize, seqLen, _modelDimension }); var v = Engine.Reshape(Engine.TensorMatMul(inputFlat, _valueWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; // Step 2: Apply the learned per-head feature MLP with batched matrix multiplies. The // previous scalar extraction/fill path detached Q/K and all four feature-map parameters. var (phiQ, phiQPreAct, phiQHidden) = ApplyFeatureMap(q, batchSize, seqLen); var (phiK, phiKPreAct, phiKHidden) = ApplyFeatureMap(k, batchSize, seqLen); - _lastPhiQ = phiQ; - _lastPhiK = phiK; - _lastPhiQHidden = phiQHidden; - _lastPhiKHidden = phiKHidden; - _lastPhiQPreActivation = phiQPreAct; - _lastPhiKPreActivation = phiKPreAct; - // Step 3: Compute output gate var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var gate = Engine.Swish(gateRaw); - _lastGateRaw = gateRaw; - _lastGate = gate; // Step 4: Causal linear attention with learned features var attnOutput = CausalLinearAttention.Normalized( Engine, phiQ, phiK, v, _numHeads, NumOps.FromDouble(1e-6)); - _lastAttnOutput = attnOutput; // Step 5: Gated output var gatedOutput = Engine.TensorMultiply(gate, attnOutput); @@ -380,7 +350,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -515,7 +484,6 @@ private Tensor LinearAttentionForward( } } - _lastAttnDenominators = denominators; return output; } @@ -588,21 +556,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastPhiQ = null; - _lastPhiK = null; - _lastPhiQHidden = null; - _lastPhiKHidden = null; - _lastPhiQPreActivation = null; - _lastPhiKPreActivation = null; - _lastAttnOutput = null; - _lastAttnDenominators = null; - _lastGateRaw = null; - _lastGate = null; _originalInputShape = null; _queryWeightsGradient = null; _keyWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/HybridBlockScheduler.cs b/src/NeuralNetworks/Layers/SSM/HybridBlockScheduler.cs index 444ce390d0..d6ee3d3306 100644 --- a/src/NeuralNetworks/Layers/SSM/HybridBlockScheduler.cs +++ b/src/NeuralNetworks/Layers/SSM/HybridBlockScheduler.cs @@ -115,13 +115,6 @@ public partial class HybridBlockScheduler : LayerBase, IShapeContract private readonly Tensor[] _normBetas; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - private Tensor[]? _lastNormedInputs; - private Tensor[]? _lastBlockOutputs; - private Tensor[]? _lastResiduals; private int[]? _originalInputShape; // Gradients for norm parameters @@ -239,30 +232,20 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = current; - _lastNormedInputs = new Tensor[_blocks.Length]; - _lastBlockOutputs = new Tensor[_blocks.Length]; - _lastResiduals = new Tensor[_blocks.Length]; - // Process through each block: pre-norm -> block -> residual for (int i = 0; i < _blocks.Length; i++) { - _lastResiduals[i] = current; - // Pre-norm: RMSNorm-style normalization var normed = ApplyRMSNorm(current, _normGammas[i], _normBetas[i], batchSize, seqLen); - _lastNormedInputs[i] = normed; // Run through sub-layer var blockOut = _blocks[i].Forward(normed); - _lastBlockOutputs[i] = blockOut; // Residual connection current = Engine.TensorAdd(current, blockOut); } var result = ApplyActivation(current); - _lastOutput = result; // Reshape back to original rank if (rank == 2) @@ -386,11 +369,6 @@ public override void UpdateParameters(T learningRate) /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastNormedInputs = null; - _lastBlockOutputs = null; - _lastResiduals = null; _originalInputShape = null; _normGammaGradients = null; _normBetaGradients = null; diff --git a/src/NeuralNetworks/Layers/SSM/HyenaLayer.cs b/src/NeuralNetworks/Layers/SSM/HyenaLayer.cs index 372ac0d196..7df58a0ead 100644 --- a/src/NeuralNetworks/Layers/SSM/HyenaLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/HyenaLayer.cs @@ -111,18 +111,6 @@ public partial class HyenaLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - private Tensor[]? _lastProjections; // (order+1) projections: v, x_1, ..., x_N - private Tensor[]? _lastProjectionsRaw; // pre-activation raw projections - private Tensor[]? _lastFilters; // generated filter kernels [seqLen, modelDim] - private Tensor[]? _lastFilterHidden; // hidden states from filter MLPs - private Tensor[]? _lastConvOutputs; // results after each convolution - private Tensor[]? _lastGatedOutputs; // results after each gating step - [Scratch] - private Tensor? _lastPreProjection; // result before output projection private int[]? _originalInputShape; // Gradients for input projections @@ -292,8 +280,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Compute (order + 1) linear projections // Projection 0 = v (value), projections 1..order = gates x_1..x_N var projections = new Tensor[_order + 1]; @@ -320,9 +306,6 @@ protected override Tensor ForwardTraced(Tensor input) } } - _lastProjections = projections; - _lastProjectionsRaw = projectionsRaw; - // Step 2: Generate implicit convolution filters using the filter MLP // For each convolution stage, the MLP takes positional encodings [0/L, 1/L, ..., (L-1)/L] // and outputs a filter of shape [seqLen, modelDim] @@ -356,9 +339,6 @@ protected override Tensor ForwardTraced(Tensor input) filters[i] = filter; } - _lastFilters = filters; - _lastFilterHidden = filterHiddenStates; - // Step 3: Apply the Hyena hierarchy // y = x_N * (h_N conv (x_{N-1} * (h_{N-1} conv ... (h_1 conv (x_1 * v)) ...))) // Start with: current = x_1 * v (element-wise gate of v by first gate) @@ -389,10 +369,6 @@ protected override Tensor ForwardTraced(Tensor input) gatedOutputs[i + 1] = current; } - _lastConvOutputs = convOutputs; - _lastGatedOutputs = gatedOutputs; - _lastPreProjection = current; - // Step 4: Output projection var currentFlat = Engine.Reshape(current, new[] { batchSize * seqLen, _modelDimension }); var outputFlat = Engine.TensorAdd( @@ -401,7 +377,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -568,15 +543,6 @@ private Tensor[] GetAllTensors() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastProjections = null; - _lastProjectionsRaw = null; - _lastFilters = null; - _lastFilterHidden = null; - _lastConvOutputs = null; - _lastGatedOutputs = null; - _lastPreProjection = null; _originalInputShape = null; _inputProjectionWeightsGradients = null; _inputProjectionBiasesGradients = null; diff --git a/src/NeuralNetworks/Layers/SSM/KimiLinearAttentionLayer.cs b/src/NeuralNetworks/Layers/SSM/KimiLinearAttentionLayer.cs index 55a6fc2fe3..303af099ce 100644 --- a/src/NeuralNetworks/Layers/SSM/KimiLinearAttentionLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/KimiLinearAttentionLayer.cs @@ -113,27 +113,9 @@ public partial class KimiLinearAttentionLayer : LayerBase, IShapeContract // Cached values for backward pass [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] private Tensor? _lastKVGate; [Scratch] private Tensor? _lastKVGateRaw; - [Scratch] - private Tensor? _lastOutputGate; - [Scratch] - private Tensor? _lastOutputGateRaw; - [Scratch] - private Tensor? _lastStates; - [Scratch] - private Tensor? _lastRecurrenceOutput; private int[]? _originalInputShape; // Gradients @@ -259,24 +241,17 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Q, K, V projections var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); var q = Engine.Reshape(Engine.TensorMatMul(inputFlat, _queryWeights), new[] { batchSize, seqLen, _modelDimension }); var k = Engine.Reshape(Engine.TensorMatMul(inputFlat, _keyWeights), new[] { batchSize, seqLen, _modelDimension }); var v = Engine.Reshape(Engine.TensorMatMul(inputFlat, _valueWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; // Step 2: Output gate var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var outputGate = Engine.Swish(gateRaw); - _lastOutputGate = outputGate; - _lastOutputGateRaw = gateRaw; // Step 3: KV-driven gated linear attention recurrence // Compute the KV-driven scalar gate without leaving the tape, then use @@ -297,8 +272,8 @@ protected override Tensor ForwardTraced(Tensor input) keepDims: false), Engine.Reshape(_gateKVBias, new[] { 1, 1, _numHeads })); var kvGate = Engine.Sigmoid(gateRawHeads); - _lastKVGateRaw = gateRawHeads; - _lastKVGate = kvGate; + SaveForBackward(ref _lastKVGateRaw, gateRawHeads); + SaveForBackward(ref _lastKVGate, kvGate); var recurrenceOutput = Engine.GlaScanForward( q, Engine.Reshape( @@ -307,7 +282,6 @@ protected override Tensor ForwardTraced(Tensor input) v, kvGate, _numHeads); - _lastRecurrenceOutput = recurrenceOutput; // Step 4: Gated output var gatedOutput = Engine.TensorMultiply(recurrenceOutput, outputGate); @@ -320,7 +294,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -420,7 +393,6 @@ private Tensor KVGatedRecurrence( allStates[new[] { bi, t + 1, hi2, di, ki }] = state[new[] { bi, hi2, di, ki }]; } - _lastStates = allStates; return output; } @@ -491,17 +463,8 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; _lastKVGate = null; _lastKVGateRaw = null; - _lastOutputGate = null; - _lastOutputGateRaw = null; - _lastStates = null; - _lastRecurrenceOutput = null; _originalInputShape = null; _queryWeightsGradient = null; _keyWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/LinearRecurrentUnitLayer.cs b/src/NeuralNetworks/Layers/SSM/LinearRecurrentUnitLayer.cs index 000aea17df..b9794b8ec8 100644 --- a/src/NeuralNetworks/Layers/SSM/LinearRecurrentUnitLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/LinearRecurrentUnitLayer.cs @@ -135,23 +135,13 @@ public partial class LinearRecurrentUnitLayer : LayerBase, IShapeContract // Cached values for backward pass [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastProjectedInput; - [Scratch] private Tensor? _lastHiddenStatesReal; [Scratch] private Tensor? _lastHiddenStatesImag; [Scratch] - private Tensor? _lastRecurrenceOutput; - [Scratch] private Tensor? _lastLambdaReal; [Scratch] private Tensor? _lastLambdaImag; - [Scratch] - private Tensor? _lastLambdaMag; private int[]? _originalInputShape; // Gradients @@ -328,19 +318,15 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Input projection [batch*seq, modelDim] -> [batch*seq, modelDim] var input2D = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); var projected = Engine.TensorMatMul(input2D, _inputProjectionWeights); var bias2D = Engine.Reshape(_inputProjectionBias, new[] { 1, _modelDimension }); var projectedWithBias = Engine.TensorAdd(projected, bias2D); var projected3D = Engine.Reshape(projectedWithBias, new[] { batchSize, seqLen, _modelDimension }); - _lastProjectedInput = projected3D; // Step 2: Compute diagonal recurrence with complex eigenvalues var recurrenceOutput = DiagonalComplexRecurrence(projected3D, batchSize, seqLen); - _lastRecurrenceOutput = recurrenceOutput; // Step 3: Output projection [batch*seq, modelDim] -> [batch*seq, modelDim] var recFlat = Engine.Reshape(recurrenceOutput, new[] { batchSize * seqLen, _modelDimension }); @@ -350,7 +336,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputWithBias, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -390,7 +375,6 @@ private Tensor DiagonalComplexRecurrence(Tensor u, int batchSize, int seqL var lambdaReal = Engine.TensorMultiply(lambdaMag, Engine.TensorCos(angle)); var lambdaImag = Engine.TensorMultiply(lambdaMag, Engine.TensorSin(angle)); - _lastLambdaMag = lambdaMag; _lastLambdaReal = lambdaReal; _lastLambdaImag = lambdaImag; var groupedInput = Engine.Reshape(u, new[] { batchSize, seqLen, 1, _modelDimension }); @@ -654,15 +638,10 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastProjectedInput = null; _lastHiddenStatesReal = null; _lastHiddenStatesImag = null; - _lastRecurrenceOutput = null; _lastLambdaReal = null; _lastLambdaImag = null; - _lastLambdaMag = null; _originalInputShape = null; _nuGradient = null; _thetaGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/LogLinearAttentionLayer.cs b/src/NeuralNetworks/Layers/SSM/LogLinearAttentionLayer.cs index b895dc314c..3a446ea15b 100644 --- a/src/NeuralNetworks/Layers/SSM/LogLinearAttentionLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/LogLinearAttentionLayer.cs @@ -131,22 +131,6 @@ public partial class LogLinearAttentionLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastGate; - [Scratch] - private Tensor? _lastGateRaw; - [Scratch] - private Tensor? _lastLogLinearOutput; private int[]? _originalInputShape; // Gradients @@ -331,8 +315,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Q, K, V projections var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); @@ -348,21 +330,14 @@ protected override Tensor ForwardTraced(Tensor input) Engine.TensorMatMul(inputFlat, _valueWeights), Engine.Reshape(_valueBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; - // Step 2: Output gate var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var gate = Engine.Swish(gateRaw); - _lastGate = gate; - _lastGateRaw = gateRaw; // Step 3: Log-linear attention with hierarchical state compression var logLinOutput = LogLinearForward(q, k, v, batchSize, seqLen); - _lastLogLinearOutput = logLinOutput; // Step 4: Gated output var gatedOutput = Engine.TensorMultiply(logLinOutput, gate); @@ -375,7 +350,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -576,14 +550,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastGate = null; - _lastGateRaw = null; - _lastLogLinearOutput = null; _originalInputShape = null; _queryWeightsGradient = null; _queryBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/LonghornLayer.cs b/src/NeuralNetworks/Layers/SSM/LonghornLayer.cs index d28786d37f..5b3c17363d 100644 --- a/src/NeuralNetworks/Layers/SSM/LonghornLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/LonghornLayer.cs @@ -123,24 +123,6 @@ public partial class LonghornLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastAlpha; - [Scratch] - private Tensor? _lastStates; - [Scratch] - private Tensor? _lastRecurrenceOutput; - [Scratch] - private Tensor? _lastNormedOutput; private int[]? _originalInputShape; // Gradients @@ -305,8 +287,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Q, K, V projections var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); @@ -325,24 +305,17 @@ protected override Tensor ForwardTraced(Tensor input) Engine.Reshape(_valueBias, new[] { 1, _modelDimension })); var v = Engine.Reshape(vFlat, new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; - // Step 2: Alpha (learning rate / update gate) via sigmoid var alphaRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _alphaWeights), Engine.Reshape(_alphaBias, new[] { 1, _numHeads })), new[] { batchSize, seqLen, _numHeads }); var alpha = Engine.Sigmoid(alphaRaw); - _lastAlpha = alpha; // Step 3: Online learning recurrence per head var recurrenceOutput = OnlineLearnerForward(q, k, v, alpha, batchSize, seqLen); - _lastRecurrenceOutput = recurrenceOutput; // Step 4: Group normalization per head var normedOutput = ApplyGroupNorm(recurrenceOutput, batchSize, seqLen); - _lastNormedOutput = normedOutput; // Step 5: Output projection var normedFlat = Engine.Reshape(normedOutput, new[] { batchSize * seqLen, _modelDimension }); @@ -352,7 +325,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -404,11 +376,6 @@ private Tensor OnlineLearnerForward( var state = Tensor.CreateDefault( new[] { headBatch, _headDimension, _headDimension }, NumOps.Zero); var outputs = new List>(seqLen); - var states = new List>(seqLen + 1) - { - Engine.Reshape(state, - new[] { batchSize, 1, _numHeads, _headDimension, _headDimension }) - }; for (int t = 0; t < seqLen; t++) { @@ -430,11 +397,8 @@ private Tensor OnlineLearnerForward( outputs.Add(Engine.Reshape(Engine.BatchMatMul(state, qCol), new[] { headBatch, 1, _headDimension })); - states.Add(Engine.Reshape(state, - new[] { batchSize, 1, _numHeads, _headDimension, _headDimension })); } - _lastStates = Engine.TensorConcatenate(states.ToArray(), axis: 1); return FromHeadMajor(Engine.TensorConcatenate(outputs.ToArray(), axis: 1), batchSize, seqLen); } @@ -637,15 +601,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastAlpha = null; - _lastStates = null; - _lastRecurrenceOutput = null; - _lastNormedOutput = null; _originalInputShape = null; _queryWeightsGradient = null; _queryBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/MEGALayer.cs b/src/NeuralNetworks/Layers/SSM/MEGALayer.cs index 9c048e5184..088eb3c089 100644 --- a/src/NeuralNetworks/Layers/SSM/MEGALayer.cs +++ b/src/NeuralNetworks/Layers/SSM/MEGALayer.cs @@ -136,29 +136,9 @@ public partial class MEGALayer : LayerBase, IShapeContract // Cached forward pass values [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] private Tensor? _lastEmaInput; [Scratch] private Tensor? _lastEmaStates; - [Scratch] - private Tensor? _lastEmaProjected; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastAttnScores; - [Scratch] - private Tensor? _lastAttnOutput; - [Scratch] - private Tensor? _lastGateRaw; - [Scratch] - private Tensor? _lastGate; private int[]? _originalInputShape; // Gradients @@ -370,14 +350,12 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Project input to EMA space var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); var emaInput = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _emaProjectInWeights), Engine.Reshape(_emaProjectInBias, new[] { 1, _emaDimension })), new[] { batchSize, seqLen, _emaDimension }); - _lastEmaInput = emaInput; + SaveForBackward(ref _lastEmaInput, emaInput); // Step 2: Apply multi-dimensional EMA through the fused RG-LRU scan. // RG-LRU evolves h_t = a*h_{t-1} + sqrt(1-a^2)*(i_t*v_t). @@ -410,7 +388,6 @@ protected override Tensor ForwardTraced(Tensor input) var emaProjected = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(emaFlat, _emaProjectOutWeights), Engine.Reshape(_emaProjectOutBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); - _lastEmaProjected = emaProjected; // Step 4: Compute Q, K from EMA output, V from original input var emaFlatFull = Engine.Reshape(emaProjected, new[] { batchSize * seqLen, _modelDimension }); @@ -423,22 +400,16 @@ protected override Tensor ForwardTraced(Tensor input) var v = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _valueWeights), Engine.Reshape(_valueBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; // Step 5: Compute gated attention var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var gate = Engine.Sigmoid(gateRaw); - _lastGateRaw = gateRaw; - _lastGate = gate; // Step 6: Multi-head causal attention var attnOutput = CausalLinearAttention.ScaledDotProduct( Engine, q, k, v, _numHeads, causal: true); - _lastAttnOutput = attnOutput; // Step 7: Gated output var gatedOutput = Engine.TensorMultiply(gate, attnOutput); @@ -451,7 +422,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -578,7 +548,6 @@ private Tensor CausalAttentionForward( } } - _lastAttnScores = attnScores; return output; } @@ -685,18 +654,8 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; _lastEmaInput = null; _lastEmaStates = null; - _lastEmaProjected = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastAttnScores = null; - _lastAttnOutput = null; - _lastGateRaw = null; - _lastGate = null; _originalInputShape = null; _emaAlphaLogitGradient = null; _emaProjectInWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/Mamba2Block.cs b/src/NeuralNetworks/Layers/SSM/Mamba2Block.cs index c1a6f7a0a3..0d59a14338 100644 --- a/src/NeuralNetworks/Layers/SSM/Mamba2Block.cs +++ b/src/NeuralNetworks/Layers/SSM/Mamba2Block.cs @@ -127,32 +127,6 @@ public partial class Mamba2Block : LayerBase, IShapeContract private Tensor _normBeta; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastXBranch; - [Scratch] - private Tensor? _lastZBranch; - [Scratch] - private Tensor? _lastConvOutput; - [Scratch] - private Tensor? _lastSiluOutput; - [Scratch] - private Tensor? _lastSsdOutput; - [Scratch] - private Tensor? _lastGatedOutput; - [Scratch] - private Tensor? _lastDelta; - [Scratch] - private Tensor? _lastDeltaPreSoftplus; - [Scratch] - private Tensor? _lastB; - [Scratch] - private Tensor? _lastC; - [Scratch] - private Tensor? _lastNormInput; private int[]? _originalInputShape; // Gradients @@ -432,8 +406,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Input projection -> x and z branches var input2D = Engine.Reshape(input3D, new[] { batchSize * seqLen, modelDim }); var projected = Engine.TensorMatMul(input2D, _inputProjectionWeights); @@ -443,15 +415,11 @@ protected override Tensor ForwardTraced(Tensor input) var xBranch = SliceTensor(projected3D, 2, 0, _innerDimension); var zBranch = SliceTensor(projected3D, 2, _innerDimension, _innerDimension); - _lastXBranch = xBranch; - _lastZBranch = zBranch; // Step 2: Conv1D on x branch var convOutput = DepthwiseConv1DForward(xBranch, batchSize, seqLen); - _lastConvOutput = convOutput; // Step 3: SiLU activation var siluOutput = Engine.Swish(convOutput); - _lastSiluOutput = siluOutput; // Step 4: Project to B, C, delta var siluFlat = Engine.Reshape(siluOutput, new[] { batchSize * seqLen, _innerDimension }); @@ -463,24 +431,17 @@ protected override Tensor ForwardTraced(Tensor input) var dtBias2D = Engine.Reshape(_dtProjectionBias, new[] { 1, _numHeads }); dtParam = Engine.TensorAdd(dtParam, dtBias2D); var dt3D = Engine.Reshape(dtParam, new[] { batchSize, seqLen, _numHeads }); - _lastDeltaPreSoftplus = dt3D; var delta = Engine.Softplus(dt3D); - _lastDelta = delta; - _lastB = bParam; - _lastC = cParam; // Step 5: SSD computation (multi-head selective scan) var ssdOutput = SSDForward(siluOutput, delta, bParam, cParam, batchSize, seqLen); - _lastSsdOutput = ssdOutput; // Step 6: RMS Norm (simplified group norm) - _lastNormInput = ssdOutput; var normedOutput = ApplyRMSNorm(ssdOutput, batchSize, seqLen); // Step 7: Output gating: y = normed * SiLU(z) var zGate = Engine.Swish(zBranch); var gatedOutput = Engine.TensorMultiply(normedOutput, zGate); - _lastGatedOutput = gatedOutput; // Step 8: Output projection var gatedFlat = Engine.Reshape(gatedOutput, new[] { batchSize * seqLen, _innerDimension }); @@ -490,7 +451,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputWithBias, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -1116,19 +1076,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastXBranch = null; - _lastZBranch = null; - _lastConvOutput = null; - _lastSiluOutput = null; - _lastSsdOutput = null; - _lastGatedOutput = null; - _lastDelta = null; - _lastDeltaPreSoftplus = null; - _lastB = null; - _lastC = null; - _lastNormInput = null; _originalInputShape = null; _inputProjectionWeightsGradient = null; _inputProjectionBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/MambaBlock.cs b/src/NeuralNetworks/Layers/SSM/MambaBlock.cs index 870e0f6b6a..cfb70aab4c 100644 --- a/src/NeuralNetworks/Layers/SSM/MambaBlock.cs +++ b/src/NeuralNetworks/Layers/SSM/MambaBlock.cs @@ -138,30 +138,6 @@ public partial class MambaBlock : LayerBase, IShapeContract // Cached values for backward pass [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastXBranch; - [Scratch] - private Tensor? _lastZBranch; - [Scratch] - private Tensor? _lastConvOutput; - [Scratch] - private Tensor? _lastSiluOutput; - [Scratch] - private Tensor? _lastScanOutput; - [Scratch] - private Tensor? _lastGatedOutput; - [Scratch] - private Tensor? _lastDelta; - [Scratch] - private Tensor? _lastDeltaPreSoftplus; - [Scratch] - private Tensor? _lastB; - [Scratch] - private Tensor? _lastC; - [Scratch] private Tensor? _lastHiddenStates; private Tensor? _initialHiddenState; private int[]? _originalInputShape; @@ -537,8 +513,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Input projection -> x branch and z branch var input2D = Engine.Reshape(input3D, new[] { batchSize * seqLen, modelDim }); var projected = Engine.TensorMatMul(input2D, _inputProjectionWeights); @@ -550,16 +524,11 @@ protected override Tensor ForwardTraced(Tensor input) var xBranch = Engine.TensorNarrow(projected3D, 2, 0, _innerDimension); var zBranch = Engine.TensorNarrow(projected3D, 2, _innerDimension, _innerDimension); - _lastXBranch = xBranch; - _lastZBranch = zBranch; - // Step 2: Conv1D on x branch (depthwise, causal) - Engine-accelerated var convOutput = DepthwiseConv1DForward(xBranch, seqLen); - _lastConvOutput = convOutput; // Step 3: SiLU activation via Engine var siluOutput = Engine.Swish(convOutput); - _lastSiluOutput = siluOutput; // Step 4: Project to SSM parameters (delta, B, C) var siluFlat = Engine.Reshape(siluOutput, new[] { batchSize * seqLen, _innerDimension }); @@ -577,11 +546,7 @@ protected override Tensor ForwardTraced(Tensor input) var deltaProjWithBias = Engine.TensorAdd(deltaProjFlat, dtBias2D); var deltaProj3D = Engine.Reshape(deltaProjWithBias, new[] { batchSize, seqLen, _innerDimension }); - _lastDeltaPreSoftplus = deltaProj3D; var delta = Engine.Softplus(deltaProj3D); - _lastDelta = delta; - _lastB = bParam; - _lastC = cParam; // Step 6: Selective scan (core SSM computation). // Fast path (no carried initial state AND caller doesn't need state output): @@ -615,12 +580,10 @@ protected override Tensor ForwardTraced(Tensor input) _lastHiddenStates = hiddenStatesResult; } _initialHiddenState = null; // consumed - _lastScanOutput = scanOutput; // Step 7: Output gating: y = scan_output * SiLU(z) via Engine var zGate = Engine.Swish(zBranch); var gatedOutput = Engine.TensorMultiply(scanOutput, zGate); - _lastGatedOutput = gatedOutput; // Step 8: Output projection var gatedFlat = Engine.Reshape(gatedOutput, new[] { batchSize * seqLen, _innerDimension }); @@ -640,7 +603,6 @@ protected override Tensor ForwardTraced(Tensor input) var residualOutput = Engine.TensorAdd(output3D, input3D); var result = ApplyActivation(residualOutput); - _lastOutput = result; // Reshape back to original rank if (rank == 2) @@ -816,18 +778,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastXBranch = null; - _lastZBranch = null; - _lastConvOutput = null; - _lastSiluOutput = null; - _lastScanOutput = null; - _lastGatedOutput = null; - _lastDelta = null; - _lastDeltaPreSoftplus = null; - _lastB = null; - _lastC = null; _lastHiddenStates = null; _originalInputShape = null; _inputProjectionWeightsGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/MegalodonLayer.cs b/src/NeuralNetworks/Layers/SSM/MegalodonLayer.cs index 69f086da55..c140735779 100644 --- a/src/NeuralNetworks/Layers/SSM/MegalodonLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/MegalodonLayer.cs @@ -149,20 +149,10 @@ public partial class MegalodonLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached forward pass values - private Tensor? _lastInput; - private Tensor? _lastOutput; - private Tensor? _lastEmaInput; private Tensor? _lastEmaStatesReal; private Tensor? _lastEmaStatesImag; - private Tensor? _lastEmaOutputNorm; private Tensor? _lastEmaOutputPreNorm; private Tensor? _lastEmaStdInv; - private Tensor? _lastQuery; - private Tensor? _lastKey; - private Tensor? _lastValue; - private Tensor? _lastGate; - private Tensor? _lastGateRaw; - private Tensor? _lastAttentionOutput; private int[]? _originalInputShape; // Gradients @@ -345,8 +335,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: CEMA (Complex Exponential Moving Average) via kernel convolution var cemaOutput = CEMAKernelForward(input3D, batchSize, seqLen); @@ -361,22 +349,16 @@ protected override Tensor ForwardTraced(Tensor input) var q = Engine.Reshape(Engine.TensorMatMul(inputFlat, _queryWeights), new[] { batchSize, seqLen, _modelDimension }); var k = Engine.Reshape(Engine.TensorMatMul(inputFlat, _keyWeights), new[] { batchSize, seqLen, _modelDimension }); var v = Engine.Reshape(Engine.TensorMatMul(inputFlat, _valueWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; // Step 4: Scaled dot-product attention per head (causal) var attentionOutput = CausalLinearAttention.ScaledDotProduct( Engine, q, k, v, _numHeads, causal: true); - _lastAttentionOutput = attentionOutput; // Step 5: Gating between attention and CEMA outputs var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _gateWeights), Engine.Reshape(_gateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var gate = Engine.Sigmoid(gateRaw); - _lastGate = gate; - _lastGateRaw = gateRaw; // y = gate * attention_out + (1 - gate) * cema_projected var ones = CreateOnesLike(gate); @@ -392,7 +374,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -426,7 +407,6 @@ private Tensor CEMAKernelForward(Tensor input, int batchSize, int seqLen) var emaInput3D = Engine.Reshape( emaInput, new[] { batchSize, seqLen, _emaDimension }); - _lastEmaInput = emaInput3D; // ComplexDiagonalSsmScanForward accepts [B,T,G,W]. With G=emaDim, // W=state=1, its maps are scalar and this is an O(B*T*emaDim) @@ -470,7 +450,6 @@ private Tensor CEMAKernelForward(Tensor input, int batchSize, int seqLen) _lastEmaStdInv = Engine.TensorReciprocal( Engine.TensorSqrt( Engine.TensorAddScalar(variance, NumOps.FromDouble(eps)))); - _lastEmaOutputNorm = normed; return normed; } @@ -481,7 +460,6 @@ private Tensor CEMAForward(Tensor input, int batchSize, int seqLen) var emaInput = Engine.TensorMatMul(inputFlat, _emaInputWeights); emaInput = Engine.TensorAdd(emaInput, _emaInputBias.Reshape(1, _emaDimension)); var emaInput3D = emaInput.Reshape(batchSize, seqLen, _emaDimension); - _lastEmaInput = emaInput3D; // CEMA recurrence with complex coefficients // State is complex: h_real + i * h_imag @@ -552,7 +530,6 @@ private Tensor CEMAForward(Tensor input, int batchSize, int seqLen) var varPlusEps = Engine.TensorAddScalar(variance, NumOps.FromDouble(eps)); var stdInv = Engine.TensorReciprocal(Engine.TensorSqrt(varPlusEps)); - _lastEmaOutputNorm = outputNorm; _lastEmaStdInv = stdInv; return outputNorm; } @@ -931,20 +908,10 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastEmaInput = null; _lastEmaStatesReal = null; _lastEmaStatesImag = null; - _lastEmaOutputNorm = null; _lastEmaOutputPreNorm = null; _lastEmaStdInv = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastGate = null; - _lastGateRaw = null; - _lastAttentionOutput = null; _originalInputShape = null; _emaAlphaRealGradient = null; _emaAlphaImagGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/MesaNetLayer.cs b/src/NeuralNetworks/Layers/SSM/MesaNetLayer.cs index 8c7fcd581a..4f6cdf588a 100644 --- a/src/NeuralNetworks/Layers/SSM/MesaNetLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/MesaNetLayer.cs @@ -130,24 +130,6 @@ public partial class MesaNetLayer : LayerBase, IShapeContract private Tensor _lnBeta; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastNormalized; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastGate; - [Scratch] - private Tensor? _lastGateRaw; - [Scratch] - private Tensor? _lastMesaOutput; private int[]? _originalInputShape; // Gradients @@ -333,11 +315,8 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Layer normalization var normalized = LayerNormForward(input3D, batchSize, seqLen); - _lastNormalized = normalized; // Step 2: Q, K, V projections var normFlat = Engine.Reshape(normalized, new[] { batchSize * seqLen, _modelDimension }); @@ -354,21 +333,14 @@ protected override Tensor ForwardTraced(Tensor input) Engine.TensorMatMul(normFlat, _valueWeights), Engine.Reshape(_valueBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; - // Step 3: Output gate (SiLU activation) var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(normFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var gate = Engine.Swish(gateRaw); - _lastGate = gate; - _lastGateRaw = gateRaw; // Step 4: Mesa (ridge regression) recurrence per head using Woodbury updates var mesaOutput = MesaForward(q, k, v, batchSize, seqLen); - _lastMesaOutput = mesaOutput; // Step 5: Gated output var gatedOutput = Engine.TensorMultiply(mesaOutput, gate); @@ -383,7 +355,6 @@ protected override Tensor ForwardTraced(Tensor input) // Step 7: Residual connection var result = Engine.TensorAdd(output3D, input3D); result = ApplyActivation(result); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -558,15 +529,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastNormalized = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastGate = null; - _lastGateRaw = null; - _lastMesaOutput = null; _originalInputShape = null; _queryWeightsGradient = null; _queryBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/MinGRULayer.cs b/src/NeuralNetworks/Layers/SSM/MinGRULayer.cs index 556cf49983..c3a7a408b6 100644 --- a/src/NeuralNetworks/Layers/SSM/MinGRULayer.cs +++ b/src/NeuralNetworks/Layers/SSM/MinGRULayer.cs @@ -121,22 +121,6 @@ public partial class MinGRULayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastProjectedInput; - [Scratch] - private Tensor? _lastGatePreAct; - [Scratch] - private Tensor? _lastGate; - [Scratch] - private Tensor? _lastCandidate; - [Scratch] - private Tensor? _lastHiddenStates; - [Scratch] - private Tensor? _lastRecurrenceOutput; private int[]? _originalInputShape; // Gradients @@ -285,15 +269,12 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Input projection from modelDim to expandedDim var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); var projected = Engine.TensorMatMul(inputFlat, _inputProjectionWeights); var projBias = Engine.Reshape(_inputProjectionBias, new[] { 1, _expandedDimension }); projected = Engine.TensorAdd(projected, projBias); var projected3D = Engine.Reshape(projected, new[] { batchSize, seqLen, _expandedDimension }); - _lastProjectedInput = projected3D; // Step 2: Gate computation: z_t = sigma(W_z * projected_t + b_z) var gatePreAct = Engine.TensorMatMul(Engine.Reshape(projected, new[] { batchSize * seqLen, _expandedDimension }), _gateWeights); @@ -301,19 +282,15 @@ protected override Tensor ForwardTraced(Tensor input) gatePreAct = Engine.TensorAdd(gatePreAct, gateBias2D); var gatePreAct3D = Engine.Reshape(gatePreAct, new[] { batchSize, seqLen, _expandedDimension }); var gate = Engine.Sigmoid(gatePreAct3D); - _lastGatePreAct = gatePreAct3D; - _lastGate = gate; // Step 3: Candidate computation: h_tilde_t = W_h * projected_t + b_h var candidate = Engine.TensorMatMul(Engine.Reshape(projected, new[] { batchSize * seqLen, _expandedDimension }), _candidateWeights); var candBias2D = Engine.Reshape(_candidateBias, new[] { 1, _expandedDimension }); candidate = Engine.TensorAdd(candidate, candBias2D); var candidate3D = Engine.Reshape(candidate, new[] { batchSize, seqLen, _expandedDimension }); - _lastCandidate = candidate3D; // Step 4: Sequential recurrence: h_t = (1 - z_t) * h_{t-1} + z_t * h_tilde_t var recurrenceOutput = MinGRURecurrenceForward(gate, candidate3D, batchSize, seqLen); - _lastRecurrenceOutput = recurrenceOutput; // Step 5: Output projection from expandedDim back to modelDim var recFlat = Engine.Reshape(recurrenceOutput, new[] { batchSize * seqLen, _expandedDimension }); @@ -323,7 +300,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -366,8 +342,6 @@ private Tensor MinGRURecurrenceForward( var output = hiddenSteps.Count == 1 ? hiddenSteps[0] : Engine.TensorConcatenate(hiddenSteps.ToArray(), axis: 1); - var initialState = Tensor.CreateDefault([batchSize, 1, _expandedDimension], NumOps.Zero); - _lastHiddenStates = Engine.TensorConcatenate([initialState, output], axis: 1); return output; } @@ -436,14 +410,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastProjectedInput = null; - _lastGatePreAct = null; - _lastGate = null; - _lastCandidate = null; - _lastHiddenStates = null; - _lastRecurrenceOutput = null; _originalInputShape = null; _inputProjectionWeightsGradient = null; _inputProjectionBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/MinLSTMLayer.cs b/src/NeuralNetworks/Layers/SSM/MinLSTMLayer.cs index fb4a0d5e1d..4eee9cd010 100644 --- a/src/NeuralNetworks/Layers/SSM/MinLSTMLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/MinLSTMLayer.cs @@ -133,30 +133,6 @@ public partial class MinLSTMLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached forward pass values for backward - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastProjectedInput; - [Scratch] - private Tensor? _lastForgetGateRaw; - [Scratch] - private Tensor? _lastInputGateRaw; - [Scratch] - private Tensor? _lastForgetGateSigmoid; - [Scratch] - private Tensor? _lastInputGateSigmoid; - [Scratch] - private Tensor? _lastForgetGateNorm; - [Scratch] - private Tensor? _lastInputGateNorm; - [Scratch] - private Tensor? _lastCellCandidate; - [Scratch] - private Tensor? _lastCellStates; - [Scratch] - private Tensor? _lastRecurrenceOutput; private int[]? _originalInputShape; // Gradients @@ -319,15 +295,12 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Input projection -- [batch*seq, modelDim] x [modelDim, expandedDim] var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); var projected = Engine.TensorMatMul(inputFlat, _inputProjectionWeights); var projBias = Engine.Reshape(_inputProjectionBias, new[] { 1, _expandedDimension }); projected = Engine.TensorAdd(projected, projBias); var projected3D = Engine.Reshape(projected, new[] { batchSize, seqLen, _expandedDimension }); - _lastProjectedInput = projected3D; // Step 2: Compute forget gate f_t = sigma(W_f * x_proj + b_f) var projFlat = Engine.Reshape(projected3D, new[] { batchSize * seqLen, _expandedDimension }); @@ -336,8 +309,6 @@ protected override Tensor ForwardTraced(Tensor input) forgetRaw = Engine.TensorAdd(forgetRaw, fBias); var forgetRaw3D = Engine.Reshape(forgetRaw, new[] { batchSize, seqLen, _expandedDimension }); var forgetSigmoid = Engine.Sigmoid(forgetRaw3D); - _lastForgetGateRaw = forgetRaw3D; - _lastForgetGateSigmoid = forgetSigmoid; // Step 3: Compute input gate i_t = sigma(W_i * x_proj + b_i) var inputGateRaw = Engine.TensorMatMul(projFlat, _inputGateWeights); @@ -345,8 +316,6 @@ protected override Tensor ForwardTraced(Tensor input) inputGateRaw = Engine.TensorAdd(inputGateRaw, iBias); var inputGateRaw3D = Engine.Reshape(inputGateRaw, new[] { batchSize, seqLen, _expandedDimension }); var inputGateSigmoid = Engine.Sigmoid(inputGateRaw3D); - _lastInputGateRaw = inputGateRaw3D; - _lastInputGateSigmoid = inputGateSigmoid; // Step 4: Normalize gates -- f' = f/(f+i), i' = i/(f+i) var gateSum = Engine.TensorAdd(forgetSigmoid, inputGateSigmoid); @@ -357,20 +326,16 @@ protected override Tensor ForwardTraced(Tensor input) var forgetNorm = Engine.TensorDivide(forgetSigmoid, gateSum); var inputNorm = Engine.TensorDivide(inputGateSigmoid, gateSum); - _lastForgetGateNorm = forgetNorm; - _lastInputGateNorm = inputNorm; // Step 5: Cell candidate c_tilde = W_c * x_proj + b_c (no activation -- pure linear) var cellCandRaw = Engine.TensorMatMul(projFlat, _cellCandidateWeights); var cBias = Engine.Reshape(_cellCandidateBias, new[] { 1, _expandedDimension }); cellCandRaw = Engine.TensorAdd(cellCandRaw, cBias); var cellCandidate = Engine.Reshape(cellCandRaw, new[] { batchSize, seqLen, _expandedDimension }); - _lastCellCandidate = cellCandidate; // Step 6: Gated recurrence -- c_t = f'_t * c_{t-1} + i'_t * c_tilde_t, h_t = c_t var recurrenceOutput = MinLSTMRecurrenceForward( forgetNorm, inputNorm, cellCandidate, batchSize, seqLen); - _lastRecurrenceOutput = recurrenceOutput; // Step 7: Output projection -- [batch*seq, expandedDim] x [expandedDim, modelDim] var recFlat = Engine.Reshape(recurrenceOutput, new[] { batchSize * seqLen, _expandedDimension }); @@ -380,7 +345,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -414,8 +378,6 @@ private Tensor MinLSTMRecurrenceForward( // This avoids RG-LRU's square-root normalization and preserves MinLSTM's gradients. var output = Engine.GlaScanForward(CreateOnesLike(cellCandidate), cellCandidate, inputNorm, forgetNorm, _expandedDimension); - var initialState = Tensor.CreateDefault([batchSize, 1, _expandedDimension], NumOps.Zero); - _lastCellStates = Engine.TensorConcatenate([initialState, output], axis: 1); return output; } @@ -485,18 +447,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastProjectedInput = null; - _lastForgetGateRaw = null; - _lastInputGateRaw = null; - _lastForgetGateSigmoid = null; - _lastInputGateSigmoid = null; - _lastForgetGateNorm = null; - _lastInputGateNorm = null; - _lastCellCandidate = null; - _lastCellStates = null; - _lastRecurrenceOutput = null; _originalInputShape = null; _inputProjectionWeightsGradient = null; _inputProjectionBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/MixtureOfMambaLayer.cs b/src/NeuralNetworks/Layers/SSM/MixtureOfMambaLayer.cs index ba2e15056d..12bdae2fbf 100644 --- a/src/NeuralNetworks/Layers/SSM/MixtureOfMambaLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/MixtureOfMambaLayer.cs @@ -121,23 +121,6 @@ public partial class MixtureOfMambaLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastRouterLogits; - [Scratch] - private Tensor? _lastRouterWeightsResult; - private int[,]? _lastTopKIndices; - [Scratch] - private Tensor? _lastExpertOutputs; - [Scratch] - private Tensor? _lastGate; - [Scratch] - private Tensor? _lastGateRaw; - [Scratch] - private Tensor? _lastMoEOutput; private int[]? _originalInputShape; // Gradients @@ -311,14 +294,11 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Router - compute expert selection per token var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); var routerLogits = Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _routerWeights), Engine.Reshape(_routerBias, new[] { 1, _numExperts })); - _lastRouterLogits = routerLogits; // Softmax over experts, then use top-K only to construct the constant sparse mask. The // selected probabilities and their renormalization remain ordinary tape-tracked ops. @@ -338,8 +318,6 @@ protected override Tensor ForwardTraced(Tensor input) var weightSums = Engine.ReduceSum(maskedWeights, new[] { 1 }, keepDims: true); var routerWeightsResult = Engine.TensorDivide( maskedWeights, Engine.TensorAddScalar(weightSums, NumOps.FromDouble(1e-10))); - _lastRouterWeightsResult = routerWeightsResult; - _lastTopKIndices = topKIndices; // Step 2: one operation-specific recurrent node owns the routed expert scan. This exactly // preserves the legacy reset semantics: an expert not selected at t has zero state at t+1. @@ -347,22 +325,18 @@ protected override Tensor ForwardTraced(Tensor input) var transition = Engine.TensorExp(_expertA); var expertOutputs = Engine.RoutedDiagonalSsmScanForward( input3D, activeMask, transition, _expertB, _expertC, _expertD); - _lastExpertOutputs = expertOutputs; // Step 3: Combine expert outputs weighted by router scores var routing4D = Engine.Reshape( routerWeightsResult, new[] { batchSize, seqLen, _numExperts, 1 }); var weightedExperts = Engine.TensorMultiply(expertOutputs, routing4D); var moeOutput = Engine.ReduceSum(weightedExperts, new[] { 2 }, keepDims: false); - _lastMoEOutput = moeOutput; // Step 4: Output gate var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var gate = Engine.Swish(gateRaw); - _lastGate = gate; - _lastGateRaw = gateRaw; // Step 5: Gated output var gatedOutput = Engine.TensorMultiply(moeOutput, gate); @@ -375,7 +349,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -520,15 +493,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastRouterLogits = null; - _lastRouterWeightsResult = null; - _lastTopKIndices = null; - _lastExpertOutputs = null; - _lastGate = null; - _lastGateRaw = null; - _lastMoEOutput = null; _originalInputShape = null; _routerWeightsGradient = null; _routerBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/MixtureOfMemoriesLayer.cs b/src/NeuralNetworks/Layers/SSM/MixtureOfMemoriesLayer.cs index 8997febdd7..4b53acd63a 100644 --- a/src/NeuralNetworks/Layers/SSM/MixtureOfMemoriesLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/MixtureOfMemoriesLayer.cs @@ -146,30 +146,6 @@ public partial class MixtureOfMemoriesLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached values for backward pass - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastGate; - [Scratch] - private Tensor? _lastGateRaw; - [Scratch] - private Tensor? _lastMoMOutput; - [Scratch] - private Tensor? _lastWriteWeights; // [batch, seqLen, numMemories] - [Scratch] - private Tensor? _lastReadWeights; // [batch, seqLen, numMemories] - [Scratch] - private Tensor? _lastForgetGates; // [batch, seqLen, numMemories] - [Scratch] - private Tensor? _lastForgetGatesRaw; private int[]? _originalInputShape; // Gradients @@ -352,8 +328,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - // Step 1: Q, K, V projections var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); @@ -369,10 +343,6 @@ protected override Tensor ForwardTraced(Tensor input) Engine.TensorMatMul(inputFlat, _valueWeights), Engine.Reshape(_valueBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; - _lastKey = k; - _lastValue = v; - // Step 2: Router computations var writeLogits = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _writeRouterWeights), @@ -389,22 +359,14 @@ protected override Tensor ForwardTraced(Tensor input) Engine.Reshape(_gateRouterBias, new[] { 1, _numMemories })), new[] { batchSize, seqLen, _numMemories }); var forgetGates = Engine.Sigmoid(forgetGatesRaw); - _lastWriteWeights = writeWeights; - _lastReadWeights = readWeights; - _lastForgetGates = forgetGates; - _lastForgetGatesRaw = forgetGatesRaw; - // Step 3: Output gate var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var gate = Engine.Swish(gateRaw); - _lastGate = gate; - _lastGateRaw = gateRaw; // Step 4: MoM recurrence var momOutput = MoMForward(q, k, v, writeWeights, readWeights, forgetGates, batchSize, seqLen); - _lastMoMOutput = momOutput; // Step 5: Gated output var gatedOutput = Engine.TensorMultiply(momOutput, gate); @@ -417,7 +379,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -634,18 +595,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastGate = null; - _lastGateRaw = null; - _lastMoMOutput = null; - _lastWriteWeights = null; - _lastReadWeights = null; - _lastForgetGates = null; - _lastForgetGatesRaw = null; _originalInputShape = null; _queryWeightsGradient = null; _queryBiasGradient = null; diff --git a/src/NeuralNetworks/Layers/SSM/MultiLatentAttentionLayer.cs b/src/NeuralNetworks/Layers/SSM/MultiLatentAttentionLayer.cs index f434e1fc7e..ca617f9f70 100644 --- a/src/NeuralNetworks/Layers/SSM/MultiLatentAttentionLayer.cs +++ b/src/NeuralNetworks/Layers/SSM/MultiLatentAttentionLayer.cs @@ -111,26 +111,6 @@ public partial class MultiLatentAttentionLayer : LayerBase, IShapeContract private Tensor _outputProjectionBias; // Cached forward pass values - [Scratch] - private Tensor? _lastInput; - [Scratch] - private Tensor? _lastOutput; - [Scratch] - private Tensor? _lastLatent; - [Scratch] - private Tensor? _lastQuery; - [Scratch] - private Tensor? _lastKey; - [Scratch] - private Tensor? _lastValue; - [Scratch] - private Tensor? _lastAttnWeights; - [Scratch] - private Tensor? _lastAttnOutput; - [Scratch] - private Tensor? _lastOutputGate; - [Scratch] - private Tensor? _lastOutputGateRaw; private int[]? _originalInputShape; // Gradients @@ -280,8 +260,6 @@ protected override Tensor ForwardTraced(Tensor input) ? Engine.Reshape(input, new[] { 1, seqLen, modelDim }) : Engine.Reshape(input, new[] { batchSize, seqLen, modelDim }); - _lastInput = input3D; - var inputFlat = Engine.Reshape(input3D, new[] { batchSize * seqLen, _modelDimension }); // Step 1: Compress input to latent c_t = W_c * x_t + b_c @@ -289,31 +267,24 @@ protected override Tensor ForwardTraced(Tensor input) Engine.TensorMatMul(inputFlat, _compressWeights), Engine.Reshape(_compressBias, new[] { 1, _latentDimension })); var latent = Engine.Reshape(latentFlat, new[] { batchSize, seqLen, _latentDimension }); - _lastLatent = latent; // Step 2: Reconstruct K and V from latent var kFlat = Engine.TensorMatMul(latentFlat, _keyUpWeights); var vFlat = Engine.TensorMatMul(latentFlat, _valueUpWeights); var k = Engine.Reshape(kFlat, new[] { batchSize, seqLen, _modelDimension }); var v = Engine.Reshape(vFlat, new[] { batchSize, seqLen, _modelDimension }); - _lastKey = k; - _lastValue = v; // Step 3: Q projection (full rank) var q = Engine.Reshape(Engine.TensorMatMul(inputFlat, _queryWeights), new[] { batchSize, seqLen, _modelDimension }); - _lastQuery = q; // Step 4: Output gate var gateRaw = Engine.Reshape(Engine.TensorAdd( Engine.TensorMatMul(inputFlat, _outputGateWeights), Engine.Reshape(_outputGateBias, new[] { 1, _modelDimension })), new[] { batchSize, seqLen, _modelDimension }); var gate = Engine.Swish(gateRaw); - _lastOutputGate = gate; - _lastOutputGateRaw = gateRaw; // Step 5: Causal multi-head attention var attnOutput = CausalMultiHeadAttention(q, k, v, batchSize, seqLen); - _lastAttnOutput = attnOutput; // Step 6: Gated output var gatedOutput = Engine.TensorMultiply(attnOutput, gate); @@ -326,7 +297,6 @@ protected override Tensor ForwardTraced(Tensor input) var output3D = Engine.Reshape(outputFlat, new[] { batchSize, seqLen, _modelDimension }); var result = ApplyActivation(output3D); - _lastOutput = result; if (rank == 2) return Engine.Reshape(result, new[] { seqLen, _modelDimension }); @@ -378,7 +348,6 @@ private Tensor CausalMultiHeadAttention( scale: 1.0 / Math.Sqrt(_headDimension), out var attentionWeights); - _lastAttnWeights = attentionWeights; return MergeHeads(attended, batchSize, seqLen); } @@ -483,16 +452,6 @@ public override void ClearGradients() /// public override void ResetState() { - _lastInput = null; - _lastOutput = null; - _lastLatent = null; - _lastQuery = null; - _lastKey = null; - _lastValue = null; - _lastAttnWeights = null; - _lastAttnOutput = null; - _lastOutputGate = null; - _lastOutputGateRaw = null; _originalInputShape = null; _compressWeightsGradient = null; _compressBiasGradient = null; diff --git a/src/ReinforcementLearning/Agents/DQNAgent.cs b/src/ReinforcementLearning/Agents/DQNAgent.cs index 720fee9898..4379843610 100644 --- a/src/ReinforcementLearning/Agents/DQNAgent.cs +++ b/src/ReinforcementLearning/Agents/DQNAgent.cs @@ -181,7 +181,7 @@ private NeuralNetwork BuildQNetwork() ); return new NeuralNetwork(architecture, - optimizer: _dqnOptions.Optimizer, + optimizer: _dqnOptions.OptimizerFactory?.Invoke(), lossFunction: _dqnOptions.LossFunction); } diff --git a/src/ReinforcementLearning/Agents/DecisionTransformerAgent.cs b/src/ReinforcementLearning/Agents/DecisionTransformerAgent.cs index 1d935c3f8f..014a9ce821 100644 --- a/src/ReinforcementLearning/Agents/DecisionTransformerAgent.cs +++ b/src/ReinforcementLearning/Agents/DecisionTransformerAgent.cs @@ -81,8 +81,6 @@ protected override void RegisterComponents() } private DecisionTransformerOptions _options; - private IOptimizer, Vector> _optimizer; - private INeuralNetwork _transformerNetwork; private List<(Vector state, Vector action, T reward, T returnToGo, Vector previousAction)> _trajectoryBuffer; private int _updateCount; @@ -91,19 +89,12 @@ protected override void RegisterComponents() public DecisionTransformerAgent() : this(new DecisionTransformerOptions()) { } - public DecisionTransformerAgent(DecisionTransformerOptions options, IOptimizer, Vector>? optimizer = null) + public DecisionTransformerAgent(DecisionTransformerOptions options) : base(options) { Guard.NotNull(options); _options = options; Options = _options; - _optimizer = optimizer ?? options.Optimizer ?? new AdamOptimizer, Vector>(this, new AdamOptimizerOptions, Vector> - { - InitialLearningRate = 0.001, - Beta1 = 0.9, - Beta2 = 0.999, - Epsilon = 1e-8 - }); _updateCount = 0; _trajectoryBuffer = new List<(Vector, Vector, T, T, Vector)>(); _currentContext = new SequenceContext(); @@ -149,7 +140,7 @@ public DecisionTransformerAgent(DecisionTransformerOptions options, IOptimize outputSize: _options.ActionSize, layers: layers ); - _transformerNetwork = new NeuralNetwork(architecture, lossFunction: _options.LossFunction); + _transformerNetwork = new NeuralNetwork(architecture, _options.OptimizerFactory?.Invoke(), _options.LossFunction); // Register network with base class Networks.Add(_transformerNetwork); diff --git a/src/ReinforcementLearning/Agents/MADDPGAgent.cs b/src/ReinforcementLearning/Agents/MADDPGAgent.cs index ddffbe5d66..e239d7a2be 100644 --- a/src/ReinforcementLearning/Agents/MADDPGAgent.cs +++ b/src/ReinforcementLearning/Agents/MADDPGAgent.cs @@ -81,8 +81,6 @@ protected override void RegisterComponents() } private MADDPGOptions _options; - private IOptimizer, Vector> _optimizer; - // Networks for each agent private List> _actorNetworks; [Buffer] @@ -106,22 +104,13 @@ protected override void RegisterComponents() public MADDPGAgent() : this(new MADDPGOptions()) { } - public MADDPGAgent(MADDPGOptions options, IOptimizer, Vector>? optimizer = null) + public MADDPGAgent(MADDPGOptions options) : base(options) { Guard.NotNull(options); _options = options; Options = _options; _options.Validate(); - // Issue #3 fix: Use configured actor learning rate for default optimizer - _optimizer = optimizer ?? options.Optimizer ?? PaperOptimizerFactory.VerifyHandBuilt(this, - new AdamOptimizer, Vector>(this, new AdamOptimizerOptions, Vector> - { - InitialLearningRate = NumOps.ToDouble(_options.ActorLearningRate), - Beta1 = 0.9, - Beta2 = 0.999, - Epsilon = 1e-8 - })); _stepCount = 0; _replayBuffer = new UniformReplayBuffer, Vector>(_options.ReplayBufferSize); _perAgentRewards = new Dictionary>(); @@ -184,7 +173,7 @@ private INeuralNetwork CreateActorNetwork() outputSize: _options.ActionSize, layers: layers); - return new NeuralNetwork(architecture, lossFunction: _options.LossFunction); + return new NeuralNetwork(architecture, _options.OptimizerFactory?.Invoke(), _options.LossFunction); } private INeuralNetwork CreateCriticNetwork() @@ -215,7 +204,7 @@ private INeuralNetwork CreateCriticNetwork() outputSize: 1, layers: layers); - return new NeuralNetwork(architecture, lossFunction: _options.LossFunction); + return new NeuralNetwork(architecture, _options.OptimizerFactory?.Invoke(), _options.LossFunction); } private void InitializeReplayBuffer() diff --git a/src/ReinforcementLearning/Agents/QMIXAgent.cs b/src/ReinforcementLearning/Agents/QMIXAgent.cs index 8a23e68ad8..17390699fe 100644 --- a/src/ReinforcementLearning/Agents/QMIXAgent.cs +++ b/src/ReinforcementLearning/Agents/QMIXAgent.cs @@ -77,8 +77,6 @@ protected override void RegisterComponents() } private QMIXOptions _options; - private IOptimizer, Vector> _optimizer; - // Per-agent Q-networks private List> _agentNetworks; [Buffer] @@ -101,19 +99,12 @@ public QMIXAgent() { } - public QMIXAgent(QMIXOptions options, IOptimizer, Vector>? optimizer = null) + public QMIXAgent(QMIXOptions options) : base(options) { Guard.NotNull(options); _options = options; Options = _options; - _optimizer = optimizer ?? options.Optimizer ?? new AdamOptimizer, Vector>(this, new AdamOptimizerOptions, Vector> - { - InitialLearningRate = 0.001, - Beta1 = 0.9, - Beta2 = 0.999, - Epsilon = 1e-8 - }); _epsilon = options.EpsilonStart; _replayBuffer = new UniformReplayBuffer, Vector>(_options.ReplayBufferSize); _stepCount = 0; @@ -419,9 +410,10 @@ public override T Train() } private Training.TapeTrainingStepper? _jointStepper; - private AiDotNet.Optimizers.StochasticGradientDescentOptimizer, Tensor>? _jointSgd; + // The joint step's optimizer, built on first use: Options.OptimizerFactory's when set, else plain SGD. + private IGradientBasedOptimizer, Tensor>? _jointOptimizer; - /// One plain-SGD step on the TD loss over the agent and mixing networks, from a packed step. + /// One joint-optimizer step on the TD loss over the agent and mixing networks, from a packed step. private void TrainJointStep(Tensor packed, int globalStateLength) { var agentBases = _agentNetworks.Cast>().ToList(); @@ -453,7 +445,7 @@ Tensor TdLoss(Tensor step) return Engine.ReduceMean(Engine.TensorMultiply(diff, diff), new[] { 0, 1 }, keepDims: false); } - _jointSgd ??= new AiDotNet.Optimizers.StochasticGradientDescentOptimizer, Tensor>(null, + _jointOptimizer ??= _options.OptimizerFactory?.Invoke() ?? new AiDotNet.Optimizers.StochasticGradientDescentOptimizer, Tensor>(null, new StochasticGradientDescentOptimizerOptions, Tensor> { InitialLearningRate = Convert.ToDouble(LearningRate), @@ -467,7 +459,7 @@ Tensor TdLoss(Tensor step) Target = packed, Forward = TdLoss, ComputeLoss = (loss, _) => loss, - Optimizer = _jointSgd, + Optimizer = _jointOptimizer, ExtraParameters = parameters, }); } diff --git a/tests/AiDotNet.Tests/UnitTests/ReinforcementLearning/RLOptimizerFactoryTests.cs b/tests/AiDotNet.Tests/UnitTests/ReinforcementLearning/RLOptimizerFactoryTests.cs new file mode 100644 index 0000000000..77a9958709 --- /dev/null +++ b/tests/AiDotNet.Tests/UnitTests/ReinforcementLearning/RLOptimizerFactoryTests.cs @@ -0,0 +1,172 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Reflection; +using AiDotNet.Interfaces; +using AiDotNet.Models.Options; +using AiDotNet.NeuralNetworks; +using AiDotNet.Optimizers; +using AiDotNet.ReinforcementLearning.Agents; +using AiDotNet.ReinforcementLearning.Agents.DecisionTransformer; +using AiDotNet.ReinforcementLearning.Agents.DQN; +using AiDotNet.ReinforcementLearning.Agents.MADDPG; +using AiDotNet.ReinforcementLearning.Agents.QMIX; +using AiDotNet.Tensors.LinearAlgebra; +using Xunit; + +namespace AiDotNet.Tests.UnitTests.ReinforcementLearning; + +/// +/// OptimizerFactory gives every network an agent trains its own optimizer (issue #2090). One shared instance +/// would share its step counter, so Adam's bias correction would advance on every network's step instead of only on +/// that network's own. +/// +/// The QMIX test trains through the fused step, which shares global optimizer state with the other fused tests. +[Collection("FusedOptimizerGlobalState")] +public class RLOptimizerFactoryTests +{ + /// A factory that records every optimizer it hands out. + private sealed class CountingFactory + { + public List, Tensor>> Created { get; } = new(); + + public IGradientBasedOptimizer, Tensor> Create() + { + var optimizer = new AdamOptimizer, Tensor>(null, + new AdamOptimizerOptions, Tensor> { InitialLearningRate = 1e-3 }); + Created.Add(optimizer); + return optimizer; + } + } + + /// Every network the agent registered, read from the base class's protected list. + private static IReadOnlyList> NetworksOf(DeepReinforcementLearningAgentBase agent) => + (typeof(DeepReinforcementLearningAgentBase) + .GetField("Networks", BindingFlags.NonPublic | BindingFlags.Instance) + ?? throw new InvalidOperationException("DeepReinforcementLearningAgentBase has no Networks field.")) + .GetValue(agent) as List> + ?? throw new InvalidOperationException("Networks is not a List>."); + + /// + /// The factory ran once per network, and each network trains with the instance its own call returned: its + /// TrainingOptimizer, which Train steps with. + /// + private static void AssertOneOptimizerPerNetwork(string agent, CountingFactory factory, IReadOnlyList> networks) + { + Assert.NotEmpty(networks); + Assert.Equal(networks.Count, factory.Created.Count); + var used = networks.Select(n => ((NeuralNetworkBase)n).TrainingOptimizer).ToList(); + Assert.All(used, optimizer => Assert.Contains(factory.Created, created => ReferenceEquals(created, optimizer))); + // ReferenceEqualityComparer is not available on net471, so compare pairwise. + for (int i = 0; i < used.Count; i++) + for (int j = i + 1; j < used.Count; j++) + Assert.False(ReferenceEquals(used[i], used[j]), $"{agent}: networks {i} and {j} share one optimizer instance."); + } + + [Fact] + public void DQN_GivesTheQNetworkAndItsTargetTheirOwnOptimizers() + { + var factory = new CountingFactory(); + var agent = new DQNAgent(new DQNOptions + { + StateSize = 4, + ActionSize = 2, + HiddenLayers = new List { 8 }, + OptimizerFactory = factory.Create, + }); + AssertOneOptimizerPerNetwork("DQN", factory, NetworksOf(agent)); + } + + [Fact] + public void MADDPG_GivesEveryActorAndCriticItsOwnOptimizer() + { + var factory = new CountingFactory(); + var agent = new MADDPGAgent(new MADDPGOptions + { + NumAgents = 2, + StateSize = 4, + ActionSize = 2, + ActorHiddenLayers = new List { 8 }, + CriticHiddenLayers = new List { 8 }, + OptimizerFactory = factory.Create, + }); + var networks = NetworksOf(agent); + Assert.Equal(8, networks.Count); // actor, target actor, critic and target critic for each of two agents + AssertOneOptimizerPerNetwork("MADDPG", factory, networks); + } + + [Fact] + public void DecisionTransformer_TrainsItsNetworkWithTheFactorysOptimizer() + { + var factory = new CountingFactory(); + var agent = new DecisionTransformerAgent(new DecisionTransformerOptions + { + StateSize = 4, + ActionSize = 2, + EmbeddingDim = 8, + NumLayers = 1, + OptimizerFactory = factory.Create, + }); + AssertOneOptimizerPerNetwork("DecisionTransformer", factory, NetworksOf(agent)); + } + + [Fact] + public void QMIX_BuildsOneJointOptimizerPerAgent_EvenFromSharedOptions() + { + // QMIX trains its agent networks and mixing network on one TD loss in one step, so each agent asks for + // exactly one optimizer however many steps it takes; two agents built from the same options get one each. + var factory = new CountingFactory(); + var options = new QMIXOptions + { + NumAgents = 2, + StateSize = 3, + ActionSize = 2, + GlobalStateSize = 1, + BatchSize = 1, + ReplayBufferSize = 1, + AgentHiddenLayers = new List { 4 }, + MixingHiddenLayers = new List { 4 }, + OptimizerFactory = factory.Create, + }; + + var agents = new[] { new QMIXAgent(options), new QMIXAgent(options) }; + Assert.Empty(factory.Created); + foreach (var agent in agents) + { + for (int step = 0; step < 2; step++) + { + agent.StoreMultiAgentExperience( + new List> { State(step, 0), State(step, 1) }, + new List> { OneHot(step), OneHot(step + 1) }, + teamReward: Math.Cos(step), + nextAgentStates: new List> { State(step + 1, 0), State(step + 1, 1) }, + globalState: new Vector(new[] { 0.1 * step }), + nextGlobalState: new Vector(new[] { 0.1 * (step + 1) }), + done: false); + agent.Train(); + } + } + + Assert.Equal(2, factory.Created.Count); + Assert.NotSame(factory.Created[0], factory.Created[1]); + // And each agent's joint step holds the optimizer the factory gave it, not a fallback SGD. + var joint = typeof(QMIXAgent).GetField("_jointOptimizer", BindingFlags.NonPublic | BindingFlags.Instance); + Assert.NotNull(joint); + Assert.Same(factory.Created[0], joint.GetValue(agents[0])); + Assert.Same(factory.Created[1], joint.GetValue(agents[1])); + } + + private static Vector State(int step, int agent) + { + var v = new Vector(3); + for (int i = 0; i < 3; i++) v[i] = Math.Sin(0.7 * i + step + 1.3 * agent); + return v; + } + + private static Vector OneHot(int index) + { + var v = new Vector(2); + v[index % 2] = 1.0; + return v; + } +}