diff --git a/src/Diffusion/DiffusionModelBase.cs b/src/Diffusion/DiffusionModelBase.cs index e71a02f016..25fafce2b1 100644 --- a/src/Diffusion/DiffusionModelBase.cs +++ b/src/Diffusion/DiffusionModelBase.cs @@ -126,6 +126,29 @@ protected virtual IEnumerable EnumerateDisposableComponents() /// public virtual ModelOptions GetOptions() => _options; + /// + /// Opt-in flag: does this subclass's / + /// path use ONLY traceable engine ops (no + /// host-side .Data.Span loops, no per-step class-field mutation)? + /// When true, attempts a + /// fused-plan step + /// with (noisySample, noise) as persistent slots — the compiled plan replays + /// the forward per training call so PredictNoise must be free of + /// per-call side effects that would freeze at trace time. + /// + /// Base default is false so existing subclasses (which may still use + /// .Data.Span internally) run through the eager tape as before. Override + /// after auditing your forward path. See ooples/AiDotNet#1846. + /// + /// + protected virtual bool SupportsFusedDenoising => false; + + private static void RestoreShadow(Tensor param, Vector shadow) + { + var span = param.Data.Span; + for (int k = 0; k < span.Length && k < shadow.Length; k++) span[k] = shadow[k]; + } + /// /// The optional neural network architecture blueprint for custom layer configuration. /// @@ -1132,6 +1155,70 @@ public virtual void Train(Tensor input, Tensor expectedOutput) noisySampleTensor = new Tensor(input._shape, noisySample); } + // Preferred fused path: MultiSlotFusedStep with (noisySample, noise, + // timestepsPerElement) as persistent slots. Only engages when the + // concrete subclass has certified its PredictNoise/PredictNoiseBatched + // path as fully traceable by opting in via SupportsFusedDenoising — + // eager PredictNoise implementations that use host-side .Data.Span + // loops would freeze at trace time. See ooples/AiDotNet#1846. + // Only try the fused path when an optimizer has already been resolved + // (avoids double-construction with a slightly-different config). + var trainingOptimizerForFused = _trainingOptimizer; + if (SupportsFusedDenoising + && trainingOptimizerForFused is not null + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + trainingOptimizerForFused, + out var mfsOptType, out var mfsLr, out var mfsB1, out var mfsB2, + out var mfsEps, out var mfsWd, out _, out _)) + { + var trainableForFused = CollectTrainableParameters(); + var noiseSlotT = new Tensor(noisySampleTensor._shape, noiseVector); + var slots = new Tensor[] + { + noisySampleTensor, + noiseSlotT, + }; + using var multiSlotStep = new AiDotNet.Training.MultiSlotFusedStep(); + var timestepsSnapshot = timesteps; + var timestepSnapshotSingle = timestep; + var isBatchedSnapshot = isBatched; + Tensor ForwardFromSlots(IReadOnlyList> s) + { + return isBatchedSnapshot + ? PredictNoiseBatched(s[0], timestepsSnapshot) + : PredictNoise(s[0], timestepSnapshotSingle); + } + Tensor ComputeLossFromSlots(Tensor pred, IReadOnlyList> s) + { + var diff = Engine.TensorSubtract(pred, s[1]); + var sq = Engine.TensorMultiply(diff, diff); + var axes = Enumerable.Range(0, sq.Shape.Length).ToArray(); + return Engine.ReduceMean(sq, axes, keepDims: false); + } + if (trainableForFused.Length > 0 + && multiSlotStep.TryStep( + parameters: trainableForFused, + zeroGradAction: null, + freshSlotData: slots, + forward: ForwardFromSlots, + computeLoss: ComputeLossFromSlots, + optimizerType: mfsOptType, + learningRate: mfsLr, + beta1: mfsB1, + beta2: mfsB2, + epsilon: mfsEps, + weightDecay: mfsWd, + out T _)) + { + if (qatShadows is not null && qatParams is not null) + { + for (int i = 0; i < qatParams.Length; i++) + RestoreShadow(qatParams[i], qatShadows[i]); + } + return; + } + } + using var tape = new GradientTape(); // Forward pass — triggers lazy layer initialization, then we walk for diff --git a/src/Finance/Forecasting/Foundation/CSDI.cs b/src/Finance/Forecasting/Foundation/CSDI.cs index 309e982117..dfd3f08e64 100644 --- a/src/Finance/Forecasting/Foundation/CSDI.cs +++ b/src/Finance/Forecasting/Foundation/CSDI.cs @@ -287,16 +287,51 @@ public override void Train(Tensor input, Tensor target) var trainableParams = Training.TapeTrainingStep.CollectParameters(Layers).ToArray(); - // CSDI's stochastic denoising step cannot be traced through the compiled - // fused plan on the current Tensors build. ComputeDenoisingPairTape samples - // a fresh (timestep, noise) pair each call, and calling it twice — once in - // the forward closure, once in the loss closure — produces two independent - // samples that don't correspond to each other. Since the compiled plan - // replays the SAME captured graph every step, we can't refresh the RNG - // sample per replay either. This path stays on the eager tape until - // Tensors' PersistentInputRegistry (PR #763) lands, at which point the - // (timestep, noise) pair can be sampled per step and passed in as data. + // Preferred fused path: MultiSlotFusedStep with the sampled (noise, + // xt-scale, sinT) tuple passed as persistent slots. Refreshes per step + // by host-sampling a fresh (t, ε) pair and copying values into the + // slot tensors — the compiled forward reads the CURRENT slot data on + // every replay. See ooples/AiDotNet#1846. + if (trainableParams.Length > 0 + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _optimizer, + out var mfsOptType, out var mfsLr, out var mfsB1, out var mfsB2, + out var mfsEps, out var mfsWd, out _, out _)) + { + var slots = BuildCsdiSlots(input, target); + if (slots is not null) + { + using var multiSlotStep = new AiDotNet.Training.MultiSlotFusedStep(); + Tensor ForwardFromSlots(IReadOnlyList> s) => DenoiserForwardFromSlots(s); + Tensor ComputeLossFromSlots(Tensor pred, IReadOnlyList> s) + { + // s[1] = ε_true (noise slot). Loss = MSE(ε_pred, ε_true) via + // the model's LossFunctionBase so custom losses are respected. + return loss.ComputeTapeLoss(pred, s[1]); + } + if (multiSlotStep.TryStep( + parameters: trainableParams, + zeroGradAction: null, + freshSlotData: slots, + forward: ForwardFromSlots, + computeLoss: ComputeLossFromSlots, + optimizerType: mfsOptType, + learningRate: mfsLr, + beta1: mfsB1, + beta2: mfsB2, + epsilon: mfsEps, + weightDecay: mfsWd, + out T fusedLoss)) + { + LastLoss = fusedLoss; + return; + } + } + } + // Eager fallback: samples (t, ε) inline and runs the denoising pair on + // the tape. Same objective; used when the fused path can't engage + // (non-fuse-able optimizer, non-GPU host, etc.). using var tape = new GradientTape(); var (epsilonPred, epsilonTarget) = ComputeDenoisingPairTape(input, target); @@ -352,6 +387,112 @@ public override void Train(Tensor input, Tensor target) } } + /// + /// Assembles the persistent-slot data for the MultiSlotFusedStep wire-up: + /// samples (t, ε) host-side then packs everything the compiled forward + /// needs as tensor slots. Slot layout: + /// + /// [0] target — the clean values to denoise (Ho 2020 x_0). + /// [1] ε_true — freshly-sampled Gaussian noise, same shape as target. + /// [2] sqrtAlphaBar_scalar — cumulative α scaling at the sampled t, shape [1]. + /// [3] sqrtOneMinusAlphaBar_scalar — noise scaling at t, shape [1]. + /// [4] sinT_scalar — sin(2π·t/T) timestep encoding, shape [1]. + /// [5] conditioned — RevIN-normalized observed input, shape [1, seqLen]. + /// + /// Returns null when the target is degenerate (zero length) so the caller + /// falls back to eager training. + /// + private IReadOnlyList>? BuildCsdiSlots(Tensor input, Tensor target) + { + int targetLen = target.Length; + if (targetLen <= 0) return null; + + var rand = RandomHelper.CreateSecureRandom(); + int t = rand.Next(_numDiffusionSteps); + var noiseData = new T[targetLen]; + for (int i = 0; i < targetLen; i++) noiseData[i] = SampleStandardNormal(rand); + var epsilonTrue = new Tensor(target._shape, new Vector(noiseData)); + + T sqrtAlphaBar = NumOps.Sqrt(_alphasCumprod[t]); + T sqrtOneMinus = NumOps.Sqrt(NumOps.Subtract(NumOps.One, _alphasCumprod[t])); + T sinT = NumOps.FromDouble(Math.Sin(2.0 * Math.PI * t / Math.Max(1, _numDiffusionSteps - 1))); + + var sqrtAlphaBarT = new Tensor(new[] { 1 }); + sqrtAlphaBarT[0] = sqrtAlphaBar; + var sqrtOneMinusT = new Tensor(new[] { 1 }); + sqrtOneMinusT[0] = sqrtOneMinus; + var sinTT = new Tensor(new[] { 1 }); + sinTT[0] = sinT; + + var conditioned = ApplyInstanceNormalization(input); + if (conditioned.Rank == 1) + conditioned = Engine.Reshape(conditioned, new[] { 1, conditioned.Length }); + return new[] { target, epsilonTrue, sqrtAlphaBarT, sqrtOneMinusT, sinTT, conditioned }; + } + + /// + /// Fused-path denoiser forward: reads the slot tuple from + /// , builds x_t via traceable engine + /// arithmetic (no .Data.Span pack — the previous inline + /// implementation froze the trace batch's x_t into the compiled plan), + /// concatenates [x_t | conditioning_slice | sin(t)] via + /// , and runs the projection + + /// residual stack + output projection. Returns predicted noise aligned + /// with the noise slot's shape. + /// + private Tensor DenoiserForwardFromSlots(IReadOnlyList> slots) + { + var target = slots[0]; + var epsilonTrue = slots[1]; + var sqrtAlphaBarT = slots[2]; + var sqrtOneMinusT = slots[3]; + var sinTT = slots[4]; + var conditioned = slots[5]; + + int targetLen = target.Length; + // x_t = sqrt(α̅_t) * target + sqrt(1-α̅_t) * ε — all traceable. + var scaledTarget = Engine.TensorMultiplyScalar(target, sqrtAlphaBarT[0]); + var scaledNoise = Engine.TensorMultiplyScalar(epsilonTrue, sqrtOneMinusT[0]); + var xt = Engine.TensorAdd(scaledTarget, scaledNoise); + var xt1d = xt.Rank == 1 ? xt : Engine.Reshape(xt, new[] { targetLen }); + + // Conditioning slice: first Min(condLen, hiddenDimension) elements, + // flattened to 1-D so we can concat with xt1d + sinT. + var condFlat = conditioned.Rank == 1 + ? conditioned + : Engine.Reshape(conditioned, new[] { conditioned.Length }); + int condLen = Math.Min(condFlat.Length, _hiddenDimension); + var condSlice = condLen == condFlat.Length + ? condFlat + : Engine.TensorSlice(condFlat, new[] { 0 }, new[] { condLen }); + + // Pack [xt | conditioning | sin(t)] via TensorConcatenate — replaces + // the .Data.Span[i] fill that froze at trace time. + var packed1D = Engine.TensorConcatenate(new[] { xt1d, condSlice, sinTT }, axis: 0); + var denoisingInput = Engine.Reshape(packed1D, new[] { 1, targetLen + condLen + 1 }); + + var eps = (Tensor)denoisingInput; + if (_inputProjection is not null) + eps = _inputProjection.Forward(eps); + foreach (var layer in _residualLayers) + eps = layer.Forward(eps); + if (_outputProjection is not null) + eps = _outputProjection.Forward(eps); + + if (eps.Rank == 2 && eps.Shape[1] > epsilonTrue.Length) + eps = Engine.TensorNarrow(eps, dim: 1, start: 0, length: epsilonTrue.Length); + if (eps.Length != epsilonTrue.Length) + { + throw new InvalidOperationException( + $"CSDI denoising pair: predicted-noise length ({eps.Length}, shape=[" + + $"{string.Join(",", eps._shape)}]) does not match true-noise length (" + + $"{epsilonTrue.Length}, shape=[{string.Join(",", epsilonTrue._shape)}])."); + } + if (!eps._shape.AsEnumerable().SequenceEqual(epsilonTrue._shape)) + eps = Engine.Reshape(eps, epsilonTrue._shape); + return eps; + } + /// /// Builds the (predicted-noise, true-noise) pair for one DDPM /// training step. Samples a timestep and a fresh noise tensor, @@ -535,23 +676,28 @@ public override Dictionary Evaluate(Tensor predictions, Tensor return new Dictionary { ["MSE"] = mse, ["MAE"] = mae, ["RMSE"] = NumOps.Sqrt(mse) }; } + /// + /// + /// Traceable RevIN forward — same pattern as + /// . Uses + /// / + /// / / broadcast subtract+divide so the + /// per-batch stats re-execute on every replay under a compiled fused plan. + /// public override Tensor ApplyInstanceNormalization(Tensor input) { int batchSize = input.Rank > 1 ? input.Shape[0] : 1; int seqLen = input.Rank > 1 ? input.Shape[1] : input.Length; - var result = new Tensor(input._shape); - for (int b = 0; b < batchSize; b++) - { - T mean = NumOps.Zero; - for (int t = 0; t < seqLen; t++) { int idx = b * seqLen + t; if (idx < input.Length) mean = NumOps.Add(mean, input[idx]); } - mean = NumOps.Divide(mean, NumOps.FromDouble(seqLen)); - T variance = NumOps.Zero; - for (int t = 0; t < seqLen; t++) { int idx = b * seqLen + t; if (idx < input.Length) { var diff = NumOps.Subtract(input[idx], mean); variance = NumOps.Add(variance, NumOps.Multiply(diff, diff)); } } - variance = NumOps.Divide(variance, NumOps.FromDouble(seqLen)); - T std = NumOps.Sqrt(NumOps.Add(variance, NumOps.FromDouble(1e-5))); - for (int t = 0; t < seqLen; t++) { int idx = b * seqLen + t; if (idx < input.Length && idx < result.Length) result.Data.Span[idx] = NumOps.Divide(NumOps.Subtract(input[idx], mean), std); } - } - return result; + if (seqLen <= 0) return input; + + bool reshaped = input.Rank != 2; + var flat = reshaped ? Engine.Reshape(input, new[] { batchSize, seqLen }) : input; + var mean = Engine.ReduceMean(flat, new[] { 1 }, keepDims: true); + var variance = Engine.ReduceVariance(flat, new[] { 1 }, keepDims: true); + var std = Engine.TensorSqrt(Engine.TensorAddScalar(variance, NumOps.FromDouble(1e-5))); + var centered = Engine.TensorBroadcastSubtract(flat, mean); + var normalized = Engine.TensorBroadcastDivide(centered, std); + return reshaped ? Engine.Reshape(normalized, input._shape) : normalized; } public override Dictionary GetFinancialMetrics() diff --git a/src/NeuralNetworks/NeuralNetworkBase.cs b/src/NeuralNetworks/NeuralNetworkBase.cs index 1b0e9bd698..e8fd372d38 100644 --- a/src/NeuralNetworks/NeuralNetworkBase.cs +++ b/src/NeuralNetworks/NeuralNetworkBase.cs @@ -8800,7 +8800,7 @@ private static void EmitFusedPathEventIfEnabled(bool hit, string? reason) /// configure-once contract: adaptive learning rates, attached LR schedulers, /// or AMSGrad mode (which the fused kernel doesn't model). /// - private static bool TryMapToFusedOptimizerConfig( + internal static bool TryMapToFusedOptimizerConfig( IGradientBasedOptimizer, Tensor> optimizer, out AiDotNet.Tensors.Engines.Compilation.OptimizerType optimizerType, out float learningRate, diff --git a/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs index 957a71e140..46ce9adb6f 100644 --- a/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CTGANGenerator.cs @@ -614,7 +614,46 @@ private void TrainDiscriminatorStepBatched(Matrix transformedData, int numPac var (realPacked, fakePacked) = BuildPackedRealAndFakeBatches(transformedData, numPacks); - // GPU-RESIDENT WGAN-GP disc via Tensors PR #763. + // Preferred fused path: WganGpFusedStep runs the full WGAN-GP objective + // (Wasserstein + λ·GP with createGraph=true GP) in one compiled plan, + // with persistent (real, fake, ε) slots refreshed each Step. Falls + // through to the GpuResidentFusedStep path (which uses this critic's + // per-batch epsilon sampled inline) when the optimizer has no fused- + // kernel mapping (adaptive LR, non-fuse-able type) or the primitive + // itself declines. See ooples/AiDotNet#1845. + var discParams = TapeTrainingStep.CollectParameters(_discLayers); + if (discParams.Count > 0 + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _discriminatorOptimizer, + out var wganOptType, out var wganLr, out var wganB1, + out var wganB2, out var wganEps, out var wganWd, + out _, out _)) + { + using var wganStep = new AiDotNet.Training.WganGpFusedStep(); + Tensor DiscFwd(Tensor inp) => DiscriminatorForwardBatched(inp, isTraining: true); + Tensor EpsilonSampler(int bs) => + Engine.TensorRandomUniformRange(new[] { bs, 1 }, NumOps.Zero, NumOps.One); + if (wganStep.TryStep( + discParameters: discParams, + realBatch: realPacked, + fakeBatch: fakePacked, + discForward: DiscFwd, + epsilonSampler: EpsilonSampler, + gradientPenaltyWeight: _options.GradientPenaltyWeight, + optimizerType: wganOptType, + learningRate: wganLr, + beta1: wganB1, + beta2: wganB2, + epsilon: wganEps, + weightDecay: wganWd, + out T _)) + { + return; + } + } + + // Secondary fused path: GpuResidentFusedStep with the loss composed via + // this class's ComputeGradientPenalty (createGraph=true GP fix, #1844). var trainableDiscLayers = _discLayers.OfType>().ToList(); if (trainableDiscLayers.Count > 0) { @@ -650,7 +689,8 @@ Tensor Loss(Tensor allScores, Tensor _) } using var tape = new GradientTape(); - var discParams = TapeTrainingStep.CollectParameters(_discLayers); + // discParams already collected above for the WganGpFusedStep attempt; + // reuse it here to avoid a redundant Layers → parameter scan. var realScores = DiscriminatorForwardBatched(realPacked, isTraining: true); var fakeScores = DiscriminatorForwardBatched(fakePacked, isTraining: true); diff --git a/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs index 87ab0a4dcc..130fa68710 100644 --- a/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CausalGANGenerator.cs @@ -604,11 +604,40 @@ private void TrainDiscriminatorStepBatched(Matrix transformedData, int batchS var (realBatch, fakeBatch) = BuildRealAndFakeBatches(transformedData, batchSize); var discLayerList = _discLayers.Cast>().ToList(); - // GPU-RESIDENT WGAN-GP disc via Tensors PR #763 (createGraph=true on - // compiled backward). The gradient-penalty inner tape's ops now record - // on the outer tape, so the fused plan can compile the full - // Wasserstein + λ·GP objective as one graph and backprop the penalty - // through the disc weights. + // Preferred fused path: WganGpFusedStep (see ooples/AiDotNet#1845). + var discParams = TapeTrainingStep.CollectParameters(discLayerList); + if (discParams.Count > 0 + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _discriminatorOptimizer, + out var wganOptType, out var wganLr, out var wganB1, + out var wganB2, out var wganEps, out var wganWd, + out _, out _)) + { + using var wganStep = new AiDotNet.Training.WganGpFusedStep(); + Tensor DiscFwd(Tensor inp) => DiscriminatorForwardBatched(inp, isTraining: true); + Tensor EpsilonSampler(int bs) => + Engine.TensorRandomUniformRange(new[] { bs, 1 }, NumOps.Zero, NumOps.One); + if (wganStep.TryStep( + discParameters: discParams, + realBatch: realBatch, + fakeBatch: fakeBatch, + discForward: DiscFwd, + epsilonSampler: EpsilonSampler, + gradientPenaltyWeight: _options.GradientPenaltyWeight, + optimizerType: wganOptType, + learningRate: wganLr, + beta1: wganB1, + beta2: wganB2, + epsilon: wganEps, + weightDecay: wganWd, + out T _)) + { + return; + } + } + + // Secondary fused path: GpuResidentFusedStep with the loss composed via + // this class's ComputeGradientPenalty (createGraph=true GP fix, #1844). var trainableDiscLayers = discLayerList.OfType>().ToList(); if (trainableDiscLayers.Count > 0) { @@ -647,7 +676,7 @@ Tensor Loss(Tensor allScores, Tensor _) } using var tape = new GradientTape(); - var discParams = TapeTrainingStep.CollectParameters(_discLayers.Cast>()); + // discParams already collected above for the WganGpFusedStep attempt. var realScores = DiscriminatorForwardBatched(realBatch, isTraining: true); var fakeScores = DiscriminatorForwardBatched(fakeBatch, isTraining: true); diff --git a/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs b/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs index 1482278083..137907eab9 100644 --- a/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/CopulaGANGenerator.cs @@ -668,9 +668,40 @@ private void TrainDiscriminatorStepBatched(Matrix transformedData, int numPac var (realPacked, fakePacked) = BuildPackedRealAndFakeBatches(transformedData, numPacks); - // GPU-RESIDENT WGAN-GP disc via Tensors PR #763. Same pack-real-fake pattern - // as CausalGAN; loss closure re-runs the gradient penalty on the outer tape - // via the fixed inner-tape createGraph=true (see AiDotNet #1844). + // Preferred fused path: WganGpFusedStep (see ooples/AiDotNet#1845). + var discParams = TapeTrainingStep.CollectParameters(_discLayers); + if (discParams.Count > 0 + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _discriminatorOptimizer, + out var wganOptType, out var wganLr, out var wganB1, + out var wganB2, out var wganEps, out var wganWd, + out _, out _)) + { + using var wganStep = new AiDotNet.Training.WganGpFusedStep(); + Tensor DiscFwd(Tensor inp) => DiscriminatorForwardBatched(inp, isTraining: true); + Tensor EpsilonSampler(int bs) => + Engine.TensorRandomUniformRange(new[] { bs, 1 }, NumOps.Zero, NumOps.One); + if (wganStep.TryStep( + discParameters: discParams, + realBatch: realPacked, + fakeBatch: fakePacked, + discForward: DiscFwd, + epsilonSampler: EpsilonSampler, + gradientPenaltyWeight: _options.GradientPenaltyWeight, + optimizerType: wganOptType, + learningRate: wganLr, + beta1: wganB1, + beta2: wganB2, + epsilon: wganEps, + weightDecay: wganWd, + out T _)) + { + return; + } + } + + // Secondary fused path: GpuResidentFusedStep with the loss composed via + // this class's ComputeGradientPenalty (createGraph=true GP fix, #1844). var trainableDiscLayers = _discLayers.OfType>().ToList(); if (trainableDiscLayers.Count > 0) { @@ -706,7 +737,7 @@ Tensor Loss(Tensor allScores, Tensor _) } using var tape = new GradientTape(); - var discParams = TapeTrainingStep.CollectParameters(_discLayers); + // discParams already collected above for the WganGpFusedStep attempt. var realScores = DiscriminatorForwardBatched(realPacked, isTraining: true); var fakeScores = DiscriminatorForwardBatched(fakePacked, isTraining: true); diff --git a/src/NeuralNetworks/SyntheticData/TabDDPMGenerator.cs b/src/NeuralNetworks/SyntheticData/TabDDPMGenerator.cs index 1345e86970..ddc59ae99a 100644 --- a/src/NeuralNetworks/SyntheticData/TabDDPMGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TabDDPMGenerator.cs @@ -665,6 +665,12 @@ private void TrainBatch(Matrix numericalData, Matrix categoricalData, { if (_gaussianDiffusion is null || _multinomialDiffusion is null) return; + // MultiSlotFusedStep cache: reused across rows so the compiled plan + // is compiled once (on the first row) and replayed per subsequent row + // by refreshing slot data. See ooples/AiDotNet#1846. + AiDotNet.Training.MultiSlotFusedStep? multiSlotStep = null; + try + { for (int row = startRow; row < endRow; row++) { int t = _gaussianDiffusion.SampleTimestep(); @@ -695,22 +701,205 @@ private void TrainBatch(Matrix numericalData, Matrix categoricalData, catNoisy = catClean; } - // Tape-connected diffusion training step: run the denoiser forward on - // the tape, build the TabDDPM hybrid loss (ε-prediction MSE for the - // Gaussian-diffused numerical features + softmax cross-entropy for the - // multinomial-diffused categorical features, Kotelnikov et al. 2023), - // and backpropagate through the MLP + output heads + timestep - // projection in one optimizer step. The previous hand-rolled gradient - // path never backpropagated through the denoiser, so - // layer.UpdateParameters threw for want of computed gradients. + // Preferred fused path: MultiSlotFusedStep with the raw sinusoidal + // timestep encoding as a persistent slot. The learnable + // _timestepProjection stays INSIDE the compiled forward closure + // (which _plan.Step() replays per row) so its weights participate + // in the backward pass. See ooples/AiDotNet#1846. + var trainableParams = Training.TapeTrainingStep.CollectParameters(Layers).ToArray(); + if (trainableParams.Length > 0 + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _optimizer, + out var mfsOptType, out var mfsLr, out var mfsB1, out var mfsB2, + out var mfsEps, out var mfsWd, out _, out _)) + { + var slots = BuildTabDDPMSlots(numNoisy, actualNoise, catNoisy, catClean, t); + if (slots is not null) + { + multiSlotStep ??= new AiDotNet.Training.MultiSlotFusedStep(); + Tensor ForwardFromSlots(IReadOnlyList> s) + { + // s[0] numNoisy, s[1] actualNoise, s[2] catNoisy, + // s[3] catClean, s[4] rawSinusoidalTimeEmbed. + var timeEmbed = _timestepProjection is not null + ? _timestepProjection.Forward(s[4]) + : s[4]; + var (noisePred, catLogits) = DenoiserForwardFromTensors(s[0], s[2], timeEmbed); + // Concat both heads into a single output tensor so the + // fused-step signature (Tensor forward output) is + // satisfied. Loss closure splits it back. + return Engine.TensorConcatenate(new[] { noisePred, catLogits }, axis: 0); + } + Tensor ComputeLossFromSlots(Tensor pred, IReadOnlyList> s) + { + // Split forward output back into (noisePred, catLogits). + int noisePredLen = _numNumericalFeatures > 0 && _numericalOutputHead is not null + ? s[1].Length : 0; + int catLogitsLen = _totalCategoricalWidth > 0 && _categoricalOutputHead is not null + ? s[3].Length : 0; + Tensor noisePredT, catLogitsT; + if (noisePredLen > 0 && catLogitsLen > 0) + { + noisePredT = Engine.TensorSlice(pred, new[] { 0 }, new[] { noisePredLen }); + catLogitsT = Engine.TensorSlice(pred, new[] { noisePredLen }, new[] { catLogitsLen }); + } + else if (noisePredLen > 0) + { + noisePredT = pred; + catLogitsT = new Tensor(new[] { 0 }); + } + else + { + noisePredT = new Tensor(new[] { 0 }); + catLogitsT = pred; + } + return ComputeDiffusionLossTapeFromTensors(noisePredT, s[1], catLogitsT, s[3]); + } + if (multiSlotStep.TryStep( + parameters: trainableParams, + zeroGradAction: null, + freshSlotData: slots, + forward: ForwardFromSlots, + computeLoss: ComputeLossFromSlots, + optimizerType: mfsOptType, + learningRate: mfsLr, + beta1: mfsB1, + beta2: mfsB2, + epsilon: mfsEps, + weightDecay: mfsWd, + out T _)) + { + continue; + } + } + } + + // Eager fallback: tape-connected diffusion training step: run the + // denoiser forward on the tape, build the TabDDPM hybrid loss + // (ε-prediction MSE for the Gaussian-diffused numerical features + + // softmax cross-entropy for the multinomial-diffused categorical + // features, Kotelnikov et al. 2023), and backpropagate through the + // MLP + output heads + timestep projection in one optimizer step. using var tape = new GradientTape(); - // Compute the timestep embedding INSIDE the tape so the learnable - // _timestepProjection participates in the backward pass. - var timeEmbed = CreateTimestepEmbeddingTensor(t); - var (predictedNoise, predictedLogits) = DenoiserForwardTensors(numNoisy, catNoisy, timeEmbed); + var timeEmbed2 = CreateTimestepEmbeddingTensor(t); + var (predictedNoise, predictedLogits) = DenoiserForwardTensors(numNoisy, catNoisy, timeEmbed2); var loss = ComputeDiffusionLossTape(predictedNoise, actualNoise, predictedLogits, catClean); BackwardAndStepOnPrecomputedLoss(tape, loss, _optimizer); } + } + finally + { + multiSlotStep?.Dispose(); + } + } + + /// Builds the raw sinusoidal timestep encoding (no learnable projection + /// applied — that runs INSIDE the compiled forward). Matches the first-half + /// of . + private Tensor BuildRawTimestepEmbedding(int timestep) + { + int dim = _options.TimestepEmbeddingDimension; + var embedding = new Vector(dim); + int halfDim = dim / 2; + for (int i = 0; i < halfDim; i++) + { + double freq = Math.Exp(-Math.Log(10000.0) * i / halfDim); + double angle = timestep * freq; + embedding[i] = NumOps.FromDouble(Math.Sin(angle)); + if (i + halfDim < dim) + embedding[i + halfDim] = NumOps.FromDouble(Math.Cos(angle)); + } + return VectorToTensor(embedding); + } + + /// Assembles the persistent-slot data list for the MultiSlotFusedStep + /// wire-up. Returns null when any required feature vector's length is zero + /// (indicates a degenerate row) so the caller falls back to eager training. + private IReadOnlyList>? BuildTabDDPMSlots( + Vector numNoisy, Vector actualNoise, + Vector catNoisy, Vector catClean, int timestep) + { + // Degenerate: no features at all — can't run the denoiser. + if (numNoisy.Length == 0 && catNoisy.Length == 0) return null; + + // Use zero-length tensors as placeholders for absent modality (matches + // the DenoiserForwardTensors branching pattern). MultiSlotFusedStep's + // shape-key includes zero-length slots so a shape change in either + // modality triggers a plan recompile. + return new[] + { + VectorToTensor(numNoisy), + VectorToTensor(actualNoise), + VectorToTensor(catNoisy), + VectorToTensor(catClean), + BuildRawTimestepEmbedding(timestep), + }; + } + + /// Tensor-input variant of . Used by the + /// MultiSlotFusedStep forward closure so the numerical/categorical/timestep inputs + /// are all persistent slots (their references are captured once at trace, their + /// data is refreshed per replay). + private (Tensor NoisePred, Tensor CatLogits) DenoiserForwardFromTensors( + Tensor numericalNoisy, Tensor categoricalNoisy, Tensor projectedTimeEmbed) + { + var timeT = projectedTimeEmbed.Rank == 1 + ? projectedTimeEmbed + : Engine.Reshape(projectedTimeEmbed, new[] { projectedTimeEmbed.Length }); + + Tensor current; + if (numericalNoisy.Length > 0 && categoricalNoisy.Length > 0) + current = Engine.TensorConcatenate(new[] { numericalNoisy, categoricalNoisy, timeT }, 0); + else if (numericalNoisy.Length > 0) + current = Engine.TensorConcatenate(new[] { numericalNoisy, timeT }, 0); + else if (categoricalNoisy.Length > 0) + current = Engine.TensorConcatenate(new[] { categoricalNoisy, timeT }, 0); + else + current = timeT; + + foreach (var layer in Layers) + current = layer.Forward(current); + + var noisePred = _numericalOutputHead is not null && _numNumericalFeatures > 0 + ? _numericalOutputHead.Forward(current) + : new Tensor(new[] { 0 }); + var catLogits = _categoricalOutputHead is not null && _totalCategoricalWidth > 0 + ? _categoricalOutputHead.Forward(current) + : new Tensor(new[] { 0 }); + return (noisePred, catLogits); + } + + /// Tensor-input variant of — target + /// noise and target categorical values come in as slot tensors instead of Vector<T>. + private Tensor ComputeDiffusionLossTapeFromTensors( + Tensor noisePred, Tensor actualNoise, Tensor catLogits, Tensor catClean) + { + Tensor? loss = null; + + if (_numNumericalFeatures > 0 && actualNoise.Length > 0) + { + var pred = noisePred.Rank == 1 ? noisePred : Engine.Reshape(noisePred, new[] { noisePred.Length }); + var diff = Engine.TensorSubtract(pred, actualNoise); + loss = ReduceToScalar(Engine.TensorSquare(diff)); + } + + if (_totalCategoricalWidth > 0 && catClean.Length > 0) + { + var logits = catLogits.Rank == 1 ? catLogits : Engine.Reshape(catLogits, new[] { catLogits.Length }); + int offset = 0; + for (int j = 0; j < _numCategoricalFeatures; j++) + { + int numCats = _categoricalColumnWidths[j]; + if (numCats > 0) + { + var ce = SoftmaxCrossEntropy(logits, catClean, offset, numCats); + loss = loss is null ? ce : Engine.TensorAdd(loss, ce); + offset += numCats; + } + } + } + + return loss ?? ReduceToScalar(Engine.TensorSquare(VectorToTensor(new Vector(1)))); } /// diff --git a/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs b/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs index 8b852deadc..39326a3f0e 100644 --- a/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TabSynGenerator.cs @@ -941,19 +941,88 @@ private void TrainDiffusionBatch(Matrix latentCodes, int startRow, int endRow { if (_latentDiffusion is null) return; - for (int row = startRow; row < endRow; row++) + // MultiSlotFusedStep cache: reused across rows so the compiled plan is + // built once (on the first row) and replayed per subsequent row by + // refreshing slot data. See ooples/AiDotNet#1846. + AiDotNet.Training.MultiSlotFusedStep? multiSlotStep = null; + try { - int t = _latentDiffusion.SampleTimestep(); - var clean = GetRow(latentCodes, row); - var (noisy, actualNoise) = _latentDiffusion.AddNoise(clean, t); - var timeEmbed = CreateTimestepEmbedding(t); + var trainableParams = Training.TapeTrainingStep.CollectParameters(_diffMLPLayers).ToArray(); + AiDotNet.Tensors.Engines.Compilation.OptimizerType mfsOptType = default; + float mfsLr = 0f, mfsB1 = 0f, mfsB2 = 0f, mfsEps = 0f, mfsWd = 0f; + bool fusedEligible = false; + if (trainableParams.Length > 0) + { + fusedEligible = NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _optimizer, + out mfsOptType, out mfsLr, out mfsB1, out mfsB2, + out mfsEps, out mfsWd, out _, out _); + } - using var tape = new GradientTape(); - var predictedNoise = DiffusionMLPForwardOnTape(noisy, timeEmbed); - // ε-prediction MSE (Ho et al. 2020) on the VAE latent codes (Zhang et al. 2024 TabSyn). - var diff = Engine.TensorSubtract(predictedNoise, VectorToTensor(actualNoise)); - var loss = ReduceToScalar(Engine.TensorSquare(diff)); - TapeStepOver(tape, loss, _diffMLPLayers); + for (int row = startRow; row < endRow; row++) + { + int t = _latentDiffusion.SampleTimestep(); + var clean = GetRow(latentCodes, row); + var (noisy, actualNoise) = _latentDiffusion.AddNoise(clean, t); + var timeEmbed = CreateTimestepEmbedding(t); + + if (fusedEligible) + { + // Slots: [noisyLatent, actualNoise, projectedTimeEmbed]. + // _timestepProjection is NOT in _diffMLPLayers (parity with the + // eager path — it's detached there too), so we precompute the + // projected embedding host-side and pass it as slot data. + var slots = new[] + { + VectorToTensor(noisy), + VectorToTensor(actualNoise), + VectorToTensor(timeEmbed), + }; + multiSlotStep ??= new AiDotNet.Training.MultiSlotFusedStep(); + Tensor ForwardFromSlots(IReadOnlyList> s) + { + // s[0] noisyLatent, s[1] actualNoise (unused here — read + // in Loss), s[2] projectedTimeEmbed. Concat + MLP forward. + int totalLen = s[0].Length + s[2].Length; + var input = Engine.TensorConcatenate(new[] { s[0], s[2] }, axis: 0); + var current = input; + foreach (var layer in _diffMLPLayers) current = layer.Forward(current); + return current.Rank == 1 ? current : Engine.Reshape(current, new[] { current.Length }); + } + Tensor ComputeLossFromSlots(Tensor pred, IReadOnlyList> s) + { + var diff = Engine.TensorSubtract(pred, s[1]); + return ReduceToScalar(Engine.TensorSquare(diff)); + } + if (multiSlotStep.TryStep( + parameters: trainableParams, + zeroGradAction: null, + freshSlotData: slots, + forward: ForwardFromSlots, + computeLoss: ComputeLossFromSlots, + optimizerType: mfsOptType, + learningRate: mfsLr, + beta1: mfsB1, + beta2: mfsB2, + epsilon: mfsEps, + weightDecay: mfsWd, + out T _)) + { + continue; + } + } + + using var tape = new GradientTape(); + var predictedNoise = DiffusionMLPForwardOnTape(noisy, timeEmbed); + // ε-prediction MSE (Ho et al. 2020) on the VAE latent codes (Zhang et al. 2024 TabSyn). + var diff2 = Engine.TensorSubtract(predictedNoise, VectorToTensor(actualNoise)); + var loss = ReduceToScalar(Engine.TensorSquare(diff2)); + TapeStepOver(tape, loss, _diffMLPLayers); + } + } + finally + { + multiSlotStep?.Dispose(); } } diff --git a/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs b/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs index b189080b7c..57000481fa 100644 --- a/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs +++ b/src/NeuralNetworks/SyntheticData/TableGANGenerator.cs @@ -434,7 +434,40 @@ private void TrainDiscriminatorStepBatched(Tensor realBatch, Tensor noiseB var fakeBatch = GeneratorForwardBatched(noiseBatch); fakeBatch = ApplyOutputActivationsBatched(fakeBatch); - // GPU-RESIDENT WGAN-GP disc via Tensors PR #763. + // Preferred fused path: WganGpFusedStep (see ooples/AiDotNet#1845). + var discParams = TapeTrainingStep.CollectParameters(_discLayers.Cast>()); + if (discParams.Count > 0 + && NeuralNetworks.NeuralNetworkBase.TryMapToFusedOptimizerConfig( + _discriminatorOptimizer, + out var wganOptType, out var wganLr, out var wganB1, + out var wganB2, out var wganEps, out var wganWd, + out _, out _)) + { + using var wganStep = new AiDotNet.Training.WganGpFusedStep(); + Tensor DiscFwd(Tensor inp) => DiscriminatorForwardBatched(inp, isTraining: true); + Tensor EpsilonSampler(int bs) => + Engine.TensorRandomUniformRange(new[] { bs, 1 }, NumOps.Zero, NumOps.One); + if (wganStep.TryStep( + discParameters: discParams, + realBatch: realBatch, + fakeBatch: fakeBatch, + discForward: DiscFwd, + epsilonSampler: EpsilonSampler, + gradientPenaltyWeight: _options.GradientPenaltyWeight, + optimizerType: wganOptType, + learningRate: wganLr, + beta1: wganB1, + beta2: wganB2, + epsilon: wganEps, + weightDecay: wganWd, + out T _)) + { + return; + } + } + + // Secondary fused path: GpuResidentFusedStep with the loss composed via + // this class's ComputeGradientPenalty (createGraph=true GP fix, #1844). var trainableDiscLayers = _discLayers.OfType>().ToList(); if (trainableDiscLayers.Count > 0) { @@ -470,7 +503,7 @@ Tensor Loss(Tensor allScores, Tensor _) } using var tape = new GradientTape(); - var discParams = TapeTrainingStep.CollectParameters(_discLayers.Cast>()); + // discParams already collected above for the WganGpFusedStep attempt. var realScores = DiscriminatorForwardBatched(realBatch, isTraining: true); var fakeScores = DiscriminatorForwardBatched(fakeBatch, isTraining: true);