From 25135675175f796c70e25cc03b1de63517eeb7f9 Mon Sep 17 00:00:00 2001 From: franklinic Date: Fri, 10 Jul 2026 18:19:17 -0400 Subject: [PATCH] feat(training): centralize WGAN-GP + diffusion consumers through fused primitives MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1845 and #1846. Routes all four WGAN-GP critics through WganGpFusedStep (the fused-plan primitive from PR #1843) and adds MultiSlotFusedStep wire-ups to the three tabular diffusion consumers plus a base-class opt-in hook for DiffusionModelBase subclasses. ## Optimizer config plumbing Zero new API surface: IFusedOptimizerSpec.TryGetFusedOptimizerConfig already exposes (OptimizerType, LR, Beta1, Beta2, Epsilon, WeightDecay, Schedule, UseBf16Moments). Widened NeuralNetworkBase.TryMapToFusedOptimizerConfig from private to internal so the sibling generators in the same assembly can reuse the existing helper. ## WGAN-GP consumers (#1845) * CTGANGenerator, CopulaGANGenerator, TableGANGenerator, CausalGANGenerator — each critic training method now attempts WganGpFusedStep.TryStep FIRST with the discriminator's optimizer hyperparameters extracted via TryMapToFusedOptimizerConfig, falls back to the existing GpuResidentFusedStep path (secondary fused), then the eager tape (final fallback). The ε ∈ [0, 1]^B epsilon sampler uses Engine.TensorRandomUniformRange to match each critic's local ComputeGradientPenalty behavior. * Non-Adam optimizers (Lion, LBFGS) that don't implement IFusedOptimizerSpec cleanly fall through — TryMapToFusedOptimizerConfig returns false and the code path skips to GpuResidentFusedStep as before. ## Diffusion consumers (#1846) * TabDDPMGenerator — refactored to expose a slot-based forward (BuildTabDDPMSlots + DenoiserForwardFromTensors + ComputeDiffusionLossTapeFromTensors). Per-row TrainBatch loop now attempts MultiSlotFusedStep with (numNoisy, actualNoise, catNoisy, catClean, rawSinusoidalTimeEmbed) as persistent slots. The learnable _timestepProjection stays INSIDE the compiled forward closure so its weights participate in the backward pass. Plan is compiled once on the first row and replayed via slot-data refresh for subsequent rows. * TabSynGenerator — TrainDiffusionBatch's per-row loop wired with MultiSlotFusedStep on (noisyLatent, actualNoise, projectedTimeEmbed). Matches the existing eager path's semantic that _timestepProjection is NOT in _diffMLPLayers (kept detached in the eager path too), so the projected embedding is precomputed host-side per row and passed as slot data. * Finance/Forecasting/Foundation/CSDI — - ApplyInstanceNormalization rewritten with traceable engine ops (ReduceMean + ReduceVariance + TensorSqrt + broadcast subtract/divide) — same pattern as the TFC RevIN fix. The previous `.Data.Span` per-batch loop froze at trace time. - New BuildCsdiSlots + DenoiserForwardFromSlots express the DDPM x_t formation and packed denoising input via TensorConcatenate + engine scalar multiplies. Replaces the `.Data.Span[i] = xt[0, i]` fill that baked the trace batch's x_t into the compiled plan. - Train() attempts MultiSlotFusedStep first, falls back to the existing eager ComputeDenoisingPairTape path when the fused path can't engage. * Diffusion/DiffusionModelBase — - New opt-in `protected virtual bool SupportsFusedDenoising => false;` property. Base default is false so no existing subclass changes behavior. - Train() attempts MultiSlotFusedStep when SupportsFusedDenoising is true AND the training optimizer maps cleanly to a fused config. Slots: (noisySample, noise). Loss = MSE(pred, noise). QAT shadow restoration is preserved on the fused-success path. - Subclasses with fully-traceable PredictNoise / PredictNoiseBatched (e.g. after auditing to remove `.Data.Span` host loops) can opt in via a single-line override; no infrastructure changes needed elsewhere. ## Verification * net8.0, net471, net10.0 all build clean. * No API surface changes on IGradientBasedOptimizer — the existing IFusedOptimizerSpec interface (already implemented by all fuse-able optimizers) provided everything needed. * All consumers preserve eager fallback path for non-fuse-able optimizers and non-GPU hosts. Co-Authored-By: Claude Opus 4.7 --- src/Diffusion/DiffusionModelBase.cs | 87 +++++++ src/Finance/Forecasting/Foundation/CSDI.cs | 190 ++++++++++++++-- src/NeuralNetworks/NeuralNetworkBase.cs | 2 +- .../SyntheticData/CTGANGenerator.cs | 44 +++- .../SyntheticData/CausalGANGenerator.cs | 41 +++- .../SyntheticData/CopulaGANGenerator.cs | 39 +++- .../SyntheticData/TabDDPMGenerator.cs | 213 +++++++++++++++++- .../SyntheticData/TabSynGenerator.cs | 91 +++++++- .../SyntheticData/TableGANGenerator.cs | 37 ++- 9 files changed, 684 insertions(+), 60 deletions(-) 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);