https://github.com/Mu2e/Offline/pull/1960 sub PR #2 - #1963
Conversation
|
Hi @YongyiBWu,
which require these tests: build. @Mu2e/fnalbuild-users, @Mu2e/write have access to CI actions on main. ⌛ The following tests have been triggered for 3c9c2d0: build (Build queue - API unavailable) |
|
☔ The build is failing at 3c9c2d0.
N.B. These results were obtained from a build of this Pull Request at 3c9c2d0 after being merged into the base branch at 9bb1ef2. For more information, please check the job page here. |
FNALbuild
added
build failed
build pending
and removed
build running
build failed
labels
|
@FNALbuild run build test |
|
⌛ The following tests have been triggered for ecbdf8e: build (Build queue - API unavailable) |
|
☔ The build is failing at ecbdf8e.
N.B. These results were obtained from a build of this Pull Request at ecbdf8e after being merged into the base branch at f6e4350. For more information, please check the job page here. |
The new constructor parameters are appended after initializeRandomWeights rather than placed logically, and SBDMGeneratedSample carries a temporary conversion to vector<double>, so the VDResampler modules already in the tree keep compiling against this header. Both are undone by the later PRs that replace those callers.
ecbdf8e to
a32e69e
Compare
|
@FNALbuild run build test |
|
⌛ The following tests have been triggered for a32e69e: build (Build queue - API unavailable) |
|
☀️ The build tests passed at a32e69e.
N.B. These results were obtained from a build of this Pull Request at a32e69e after being merged into the base branch at f6e4350. For more information, please check the job page here. |
oksuzian
left a comment
oksuzian
left a comment
There was a problem hiding this comment.
PR Review Summary — #1963
Reviewed at head a32e69eb89f27580f7693b62bcca711b30fdef7c. First pass.
Decision
- 🔴 request changes
Scope understood
MachineLearningTools/ScoreBasedDiffusionModelonly (2 files, +3451/-354), split out of #1960. It changes the forward process fromx + sigma*epsto the variance-preservingsqrt(alphabar)*x + sigma*eps, adds a LOGSIG schedule, v-prediction, per-coordinate Fourier embeddings, an EMA weight copy, a per-dimension gradient-weight controller, peak importance sampling, a binary checkpoint format, and several diagnostics.- The STMMC callers are explicitly deferred to a follow-up PR, and the constructor's parameter order is documented as temporary so those callers keep binding. That intent is sound, but it does not hold for
generateSampleorsaveModel— see findings 1 and 2. - CI is green at this head, which is what makes findings 1 and 2 worth stating: both changes compile silently.
Findings
-
🔴 [S0] Two parameters inserted into the middle of
generateSamplesilently rebind the existing in-tree callers- Evidence: the signature goes from
(condition, bool useHeun = true, int diffusionSteps = -1)to(condition, bool useEMANetworkIfAvailable = true, bool useHeun = true, bool useSDE = true, int diffusionSteps = -1, double sdeToOdeSigmaThreshold = -1.0). Onmain,STMMC/src/VDResamplerGenerateFromModel_module.cc:263,274,284andSTMMC/src/VDResamplerGenerateMix_module.cc:376,387,397all callgenerateSample({}, useHeun_, diffusionSteps_)with abool useHeun_and anint diffusionSteps_. After this PR those arguments bind asuseEMANetworkIfAvailable = useHeun_,useHeun = (bool)diffusionSteps_(200 converts totrue),useSDE = truefrom the default, anddiffusionSteps = -1from the default. - Impact: three silent behaviour changes in production sampling code with no diagnostic. The EMA/base network choice becomes whatever
useHeun_happens to be; the caller's requested step count is discarded in favour of the model file's; and sampling switches from the deterministic reverse process to the SDE. Theint-to-boolconversion produces no warning under-Wall, so CI stays green and nothing announces the change. - Suggested fix: append the two new booleans at the end of the parameter list, the same way the constructor's new parameters were appended, so the existing positional calls keep their meaning until the caller PR updates them. If they genuinely must sit where they are, update the six call sites in this PR.
- Evidence: the signature goes from
-
🟠 [S1]
saveModelnow writes binary to a path the in-tree configuration documents as CSV, andloadModeldispatches on the extension- Evidence:
saveModelchanges from a CSV writer to a binary writer with default"DiffusionModel.bin"; the CSV writer is renamedsaveModelCsv. Onmain,VDResamplerTrain_module.cc:369,376,391callsaveModel(SBDMstage1ModelFile)etc., and those fhicl atoms are declared withComment("CSV filename for the trained stage-1 SBDM model parameters"). No fcl inSTMMC/fcl/sets a default, so the name comes from the operator. The newloadModelselects its parser from the extension and runs the CSV parser on anything that is not.dat/.bin. - Impact: an operator following the only in-tree documentation passes
foo.csv; the trainer writes binary bytes into it and the generator then parses those bytes as CSV. The train/generate round trip is broken for exactly the filename the configuration recommends. - Suggested fix: keep
saveModelwriting CSV and name the binary writersaveModelBinary, or havesaveModeldispatch on the extension symmetrically withloadModelso the format is a property of the filename in both directions. Either way, update those three fhiclCommentstrings in the same change.
- Evidence:
-
🟠 [S1]
SBDMGeneratedSample::operator const std::vector<double>&returns a reference into a temporary, and changes what the caller receives- Evidence:
ScoreBasedDiffusionModel.hh, the conversion operator onSBDMGeneratedSample.const std::vector<double>& v = model.generateSample(...);binds through a conversion function, which does not extend the temporary's lifetime, sovdangles at the end of the full expression. Separately, the conversion yieldsvalue(de-normalized), whereas the oldgenerateSamplereturned the reverse-diffusion state in normalized coordinates. The in-tree callers survive that today only becausenormalizeDatadoes not exist onmain, sodataMean_/dataStdev_are still the 0/1 defaults; once the follow-up trainer callsnormalizeData, those callers apply their own inverse transform to an already de-normalized vector. - Impact: undefined behaviour for one plausible spelling of the call, and a silent unit change for the other. Both are invisible at compile time, which is the property that makes a temporary shim dangerous rather than harmless.
- Suggested fix: return
std::vector<double>by value from the conversion — the dangling case disappears and the cost is a copy on a path that already copies. Please also state in the follow-up PR which member the callers should read, sincezscoreis the one that matches the old return value.
- Evidence:
-
🟠 [S1] LOGSIG accepts
logSigMax > 1for SCORE and EPS, which silently destroys the training signal- Evidence: the
logSigMax != 1.0coercion in the constructor is insideif (predictionTarget_ == PredictionTarget::V). For SCORE/EPS nothing checks it. WithlogSigMax > 1,sigma(t) > 1over a range oft, soalphabar()returnsstd::max(0.0, 1 - s*s) == 0andaddNoisecomputesxt[i] = sqrt(0)*x[i] + s*eps[i]— pure noise with the data removed. - Impact: a configuration that looks valid (the parameter is documented only as "Maximum sigma for LOGSIG schedule", default 1.0) trains on noise for part of the schedule and reports nothing. That is silent degradation from a detected-but-unhandled input; the fact that the v-prediction path already coerces the value shows the constraint is understood.
- Suggested fix: validate
0 < logSigMin < logSigMax <= 1in the constructor for every prediction target and throw otherwise, rather than coercing for one target and ignoring it for the others.
- Evidence: the
-
🟡 [S2] The CSV loader rejects every checkpoint written by the current
main- Evidence: the CSV path ends with
if ((int)dataMean.size() != (dim + conditionDim)) throw ... "Normalization parameter size mismatch".main'ssaveModelwrites no[DATA_NORMALIZATION]section, sodataMeanstays empty and this always fires. Every other key absent from a legacy file is handled with a default, and the log line even reportsformat CSV (unversioned). - Suggested fix: either treat an absent normalization section as the identity (
dataMean = 0,dataStdev = 1), matching the constructor defaults, or drop the legacy-key handling and reject an unversioned CSV up front with a message that says so. Half-supporting it costs the code without buying the compatibility.
- Evidence: the CSV path ends with
-
🟡 [S2] Gradient-clipping statistics are cumulative but reported as an epoch quantity
- Evidence:
clipCount_,totalClipChecks_andclipScaleAccum_are initialized in the constructor and never reset, yetClipRatioandAvgClipScaleare printed on the per-epochtrain()log line. They are also not serialized, so a resumed run restarts the average while the loss history continues. - Suggested fix: reset the three counters at the top of each epoch, or rename the printed fields to make clear they are run-to-date.
- Evidence:
-
🟡 [S2]
firstLayerBlockMagnitudesreports the norm of the summed gradient, not the gradient magnitude- Evidence:
blocks[b].gradL2 = std::sqrt(sg) * invUsed;wheresgaccumulates overnetwork_[0].gradWafternSamplesbackward passes have summed into it. That is||sum_k g_k|| / n, and the comment calls it "Gradient L2 per block (mean per sample)". - Impact: the function's documented use is "near-zero gradL2 on a block means that input feature is not being used". A feature whose per-sample gradients cancel across the batch — which is the normal state of a well-trained input — gives a near-zero norm-of-mean and reads as dead. The diagnostic can therefore report the opposite of the truth in the case it exists to detect.
- Suggested fix: accumulate the per-sample squared gradient (zero the buffer, backward, add
sum_c g_c^2, repeat) so the reported quantity is the RMS per-sample magnitude, or rename the field to say it is the mean gradient's norm.
- Evidence:
-
🟡 [S2] Comment contradicts the code it introduces
- Evidence:
// Heun's method (2nd order) Only ODE solver, no noise added, immediately above the block that buildsdwand adds it undereffectiveSDE. (// sahred noise vectoron the next line is a typo.) - Suggested fix: delete the stale half of the comment.
- Evidence:
Simplification and efficiency (§6 — these never gate the decision, and they are the bulk of what you asked about)
-
🟡 [S2] The binary format ships with compatibility code for eight versions that never existed.
main'sScoreBasedDiffusionModelhas a single CSVsaveModeland no binary path at all, so no committed code has ever written anSBDMfile. Everyversion >= Nbranch in the loader, theversion <= 3EMA-decay fix-up, the.bin"legacy spelling", theepsPredictionCSV key, the enum-value pinning that exists to makefalse/truemap ontoSCORE/EPS, and the thirty-line version-history comment are all compatibility with out-of-tree checkpoints. That is roughly 150 lines that the repo's dead-code rule says to delete — git has them. Ship the format as version 1 with one read path; you keep the truncation sentinel and the count bounds, which are the parts that earn their keep. -
🟡 [S2] The 28-argument positional constructor is the source of several other problems in this review. It already carries a comment explaining that its parameter order is deliberately illogical, and both
loadModelpaths spell out 28 positional arguments in that order. Astruct Configwith defaulted members, passed byconst&, would remove both call sites, make the ordering question moot, make future parameters non-breaking, and letlogSigMin/logSigMax/predictionTargetsit next to the schedule they belong to. Since the follow-up PR has to touch the callers anyway, doing it there costs little more than moving the parameters back. -
🟡 [S2]
train()multiplies gradients bydimWeights_whether or notuseDimWeightController_is set, and a large amount of machinery exists to manage the consequences:clampDimWeightswith its load-path invariant, the freeze-and-log branch inupdateUseDimWeightController, most ofresetDimWeightController's rationale, and two of the five snapshot fields. Gating the multiply onuseDimWeightController_would make all of that unnecessary. If the intent is that a phase can inherit the previous phase's weights, that is worth saying explicitly, because the three separate comment blocks currently read as documentation of a trap rather than of a feature. -
🟡 [S2] Peak sampling rebuilds its pools on every
train()call: an O(N·K) partition plus a full shuffle ofQ.getLastEpochLoss's own documentation says the curriculum planner callstrain(..., epochs=1, ...), so for a 5M-row dataset that partition runs once per epoch to produce a result that only changes when the window list or the data does. Caching it keyed on the window list would make it once per run. -
⚪ [S3] Four small ones, grouped: the
updateLossWeightPower/updateGradientClipThreshold/updateLearningRate/updateBatchSize/updateUseDimWeightControllersetters return the value they just stored and nothing in the diff reads it —voidwould be clearer;trainingSampleSize_is initialized to 0, never assigned, and now round-tripped through the binary format as a permanent zero;hasDataNormalization()detects "unset" by testing the constructor's-999/+999placeholders, where abool hasNormalization_set bynormalizeData()and the loader is one line and cannot be confused with data; andevaluateAverageLossshuffles all N indices to take the firstsubsetSize, which is the expensive way to draw a small sample from a large dataset.
Validation check
- Build/tests run: none by me. CI is green at this head (build 3344, prof build 9 min). FNALbuild's table reports
whitespace check :arrow_right: found whitespace errorsandtrigger :question: Return Code 1— the trigger return code was also non-clean on the two earlier runs, so it does not look specific to this diff, but it is worth a glance. - Config contract check: not applicable, no fcl in this PR. The mismatch between
saveModel's new format and the fhiclCommentstrings inSTMMCis finding 2. - Cross-repo consistency: needs follow-up. This PR is
MachineLearningTools-only, but three of its API changes reachSTMMCcallers that it does not update. Findings 1-3 are all instances of that. - Not checked: the numerical correctness of the v-prediction and LOGSIG derivations against the literature, the importance-sampling estimator's unbiasedness at
alpha = 1, and whether the trained model actually resolves the narrow pz feature — none of that is checkable from the diff, and it is the part your PR description points at.
Residual risk
- Findings 1-3 are silent at compile time and green in CI. If #1963 merges before its follow-up,
main's STMMC sampling changes behaviour with nothing to signal it. - The forward process changed from variance-exploding to variance-preserving. Any checkpoint produced before this PR is invalid under the new sampler, and nothing in the file format records which convention a checkpoint was trained under. Finding 5 makes the CSV ones throw, which is accidental protection rather than a design; a
.datwritten by an intermediate branch would load and sample quietly.
Author follow-ups
- Append
useEMANetworkIfAvailableanduseSDEto the end ofgenerateSample's parameter list (finding 1), or update the six STMMC call sites here. - Resolve the
saveModelformat/extension mismatch and the three fhiclCommentstrings (finding 2). - Return by value from the
SBDMGeneratedSampleconversion, and say which member the callers should read (finding 3). - Validate
logSigMax <= 1for all prediction targets (finding 4). - Decide whether unversioned CSVs are supported and make the loader consistent either way (finding 5).
- Consider collapsing the binary format to version 1 and moving the constructor to a
Configstruct (findings 9 and 10) — together they remove more code than the rest of this list combined. - Could you say what you ran to validate the v-prediction and LOGSIG paths, and roughly what the peak-window loss did across a curriculum? That is the part of this PR a reader cannot check from the diff.
Diffusion model
|
Comments: The following are not true. Need to refer to later PRs.
The following fixes are added. Since the sub-PRs of 1963-1968 are stacked, I only commit the change to the branch in #1968 to avoid rebasing every sub-PR. Check 13f7162 for the changes. The following edit suggestions are rejected. These are either stylish preferences or overcoding. |
|
📝 The HEAD of |
|
@oksuzian I've updated the codes in response to the comments. The changes are only appended to the last sub-PR of #1968, at 13f7162 to avoid rebasing all the PRs in the sequence. Compilation test should be fine for all sub-PRs but only the last PR's result is meaningful. Please direct the AI to look at the whole PR sequence before making change suggestions; it kept prompting compromises for splitting the big PR as errors. |
3 participants
ScoreBasedDiffusionModel: curriculum, EMA, prediction targets, format v9