feat(kwyk): import published TF kwyk weights with a verified mapping - #384
Open
dhritimandas wants to merge 1 commit into
Open
feat(kwyk): import published TF kwyk weights with a verified mapping#384dhritimandas wants to merge 1 commit into
dhritimandas wants to merge 1 commit into
Conversation
The kwyk architecture (McClure et al. 2019) was reimplemented in PyTorch, but users could not obtain the trusted published weights -- the three SavedModels ship only inside the neuronets/kwyk Docker container, and no conversion path existed. - nobrainer/datasets/convert_kwyk.py: TF SavedModel/.npz -> KWYKMeshNet state_dict converter. Strict load, offline parity check over all 8 conv layers (including the classifier), and a .provenance.json sidecar recording the source path + SHA-256 -- the MAP and BD checkpoints are structurally indistinguishable, so the content hash is the only durable record of which model a .pth came from. - KWYKMeshNet: the output classifier is now an FFGConv3d (was plain nn.Conv3d), matching the TF logits/ layer, which is itself a full VWN conv. All five logits variables map 1:1 with no information loss; the previous mean-path collapse made MC uncertainty under-dispersed at the output (TF contributes ~11.5 of run-to-run logit variance there; the plain conv contributed 0). Also adds a bias flag threaded through the hidden layers (default False, unchanged). - ConcreteDropout3d clamp widened [0.05, 0.95] -> [0.01, 0.99] in lockstep with the converter: the published SSD checkpoint stores p up to 0.954055, so the old ceiling silently clipped 89% of trained values (598/672). A unit test pins the two clamp ranges together. - Biases are imported by default (--drop-bias is an explicit, logged opt-out): every published checkpoint carries bias_m/bias_a for every conv layer, and dropping them shifts output logits by max ~20. Every mapping claim is verified against the real container and validated numerically against the live TF graph: the converted MAP model reproduces the original's logits to max|diff| = 9.5e-05 (relative error ~4e-6) on identical input; converted SSD preserves all 672 dropout probabilities to 6e-08 and reproduces output-layer MC variance. Full verification report, including the six discrepancies found and fixed and a second independent re-verification: docs/kwyk_mapping_verification.md.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
What: We add a converter for the published kwyk TensorFlow weights. It loads them into the PyTorch
KWYKMeshNet. We verify the mapping against the real container and the live TF graph.Why: The trusted weights exist only inside the
neuronets/kwykDocker container. Retraining from scratch does not reproduce them. Verification also found three model defects. We fix them here.Converter (
nobrainer/datasets/convert_kwyk.py)--npz(pre-extracted, no TF needed) or--tf-path(requires tensorflow)load_state_dict(strict=True)as a correctness gate<out>.provenance.jsonsidecar: source path + SHA-256 + arch args (MAP and BD checkpoints are structurally identical — the hash is the only durable record of which model a.pthcame from)--drop-biasis an explicit, logged opt-out (dropping shifts logits by max ~20 vs ~1e-4)Model fixes exposed by verification
classifieris nowFFGConv3d, notnn.Conv3d— the TFlogits/layer is a full VWN conv; the plain conv discarded its sigmas and contributed 0 MC variance (TF original: ~11.5). All fivelogits/conv3d/*vars now map 1:1kl_divergence()and samples undermc_vwn=True(more faithful to the paper, but a training-objective change)ConcreteDropout3dclamp widened[0.05, 0.95] -> [0.01, 0.99]in model + converter together; old ceiling silently clipped 598/672 (89%) of the published SSDpvalues (max 0.954055). Unit test pins the two ranges to each otherbiasflag threaded through hidden layers (defaultFalse, unchanged)Verification
logits/conv3d/*namespace (notlayer_8/),layer_1..71-based, biases 24/24max|diff| = 9.5e-05(rel ~4e-6), 100% argmax agreementppreserved to 6e-08; deterministic path bit-reproducible; MC variance ~10docs/kwyk_mapping_verification.mdTest plan
test_croissantfailure, fixed by feat(data): DataSpec validation, Croissant fix, MPS guard #376)main(), MC variance)