Align ZSharp with the paper: algorithm, models, and datasets - #1
Merged
Conversation
The implementation diverged from arXiv:2505.02369 in four ways that changed training behavior. Each is corrected here. - Percentile default 70 -> 95, matching the paper's Q_p = 0.95. The old default kept the top 30% of components and sat outside the range the paper ablates (0.75-0.95). The mask also now uses strict `>`, and EPSILON is 1e-8 to match the paper's delta. - Gradient filtering now follows Eq. 9. Previously each layer fell back to keeping its top 20% when nothing passed the threshold, so no layer could ever be fully zeroed. The threshold is pooled across the network, so zeroing a layer is legitimate; the fallback to unfiltered gradients applies only when filtering zeroes every layer. - Gradient clipping removed from both the ZSharp and SGD paths. The paper specifies none, and a global clip shifts relative layer magnitudes and therefore the pooled percentile threshold. - ZSharp now builds on AdamW (lr 1e-3, weight decay 5e-5) with the paper's step decay of 0.75 every 10 epochs. `momentum` is no longer forwarded, since AdamW rejects it; it applies to the SGD baseline only. Configs, docs, the demo notebook, and the percentile ablation sweep are updated to match. Benchmark numbers in README and docs/algorithm.md were produced by the pre-fix code and are flagged as stale rather than regenerated. Models and datasets are unchanged: the paper's ResNet-56/110, VGG-16BN, small ViT variants, and Tiny-ImageNet remain unimplemented. Claude-Session: https://claude.ai/code/session_01AUnUVDcyMRHDib2K5YCuv5
Implements the architectures and datasets the ZSharp paper evaluates (arXiv:2505.02369), which were the last gap to its experimental setup. - CIFAR-style ResNet-56/110 (He et al., Sec. 4.2): three stages of basic blocks at 16/32/64 channels with parameter-free option-A shortcuts. These depths exist only in the CIFAR family, so torchvision does not ship them. Parameter counts match the published table (0.85M / 1.7M). - VGG-16BN with a CIFAR-adapted head: global average pooling into a single 512-unit classifier, as in the author's reference code, rather than torchvision's three 4096-wide ImageNet layers (~15M vs ~134M parameters on 32x32 inputs). - The paper's compact ViTs, `vit_7_8_8_384` and `vit_7_8_12_768`: 7 layers, 8 patches per side, embedding width 384, differing in head count (8 and 12) and MLP width. The trainer now passes the dataset's image size so patch size is derived per resolution. - Tiny-ImageNet-200, which torchvision does not provide. Downloads and extracts on first use; the validation split ships flat, so its labels are resolved from val_annotations.txt. Grayscale images are converted to RGB. Several details the paper leaves unstated, or states inconsistently, were resolved against the author's reference implementation at github.com/YUNBLAK/Sharpness-Aware-Minimization-with-Z-Score-Gradient-Filtering: the ResNet family, the VGG head, and the ViT dimensions. On the last, the paper reads `ViT-7/8/8-384` as layers/heads/patch-size/MLP, but that is not self-consistent — it makes both variants 8-headed with patch size 8, leaving the differing third field unexplained, and 12 patches per side does not divide a 32x32 input. The reference code's reading (fixed patches, varying heads) is used instead. All assumptions are documented in docs/algorithm.md, along with the paper's Tiny-ImageNet train-set size discrepancy. Existing models and configs are unchanged. Tests build a synthetic Tiny-ImageNet tree rather than downloading the real archive. Claude-Session: https://claude.ai/code/session_01AUnUVDcyMRHDib2K5YCuv5
Follow-ups to the model/dataset work, plus a script to check whether the percentile change actually helps. - configs/vit_zsharp.yaml used torchvision's ImageNet vit_b_16, which runs on CIFAR but is not one of the paper's architectures. It now uses vit_7_8_12_768 at the paper's batch size. - The percentile ablation in scripts/experiment.py ran on resnet18 and passed a momentum that AdamW ignores. It now uses ResNet-56, the model the paper's Table 2 ablates, and drops the unused argument. The paper configs it had commented out are enabled again, since the models they reference now exist. - scripts/ablate_percentile.py A/Bs the threshold with everything else held fixed, so the effect of the 70 -> 95 change can be measured without a full 200-epoch reproduction. Claude-Session: https://claude.ai/code/session_01AUnUVDcyMRHDib2K5YCuv5
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.
The implementation diverged from arXiv:2505.02369 in ways that changed training behavior, and it was missing the architectures and datasets the paper evaluates. This branch closes both gaps.
Algorithm fixes (
fix!)Four divergences, each of which altered results:
Q_p = 0.95clip_grad_norm_(1.0)on both pathsThe old percentile default sat outside the range the paper ablates (0.75–0.95). The old fallback guaranteed every layer kept gradients, which the pooled-threshold design explicitly permits zeroing. Clipping rescaled gradients before Z-scoring, shifting the global threshold — note this connects to the earlier revert in 27d5dd8; clipping is now gone entirely.
momentumis no longer forwarded to ZSharp (AdamW rejects it) and applies to the SGD baseline only.Models and datasets
val_annotations.txt. Verified against the real archive: 100,000 train / 10,000 val / 200 classes.Underspecified details
The paper omits normalization statistics, augmentation, the ResNet family, the VGG head, and the ViT embedding width. Each was resolved against the author's reference implementation and documented in
docs/algorithm.md.Worth reviewer attention: the paper reads
ViT-7/8/8-384as layers/heads/patch-size/MLP, but that isn't self-consistent — it makes both variants 8-headed with patch size 8, leaving the differing third field unexplained, and 12 patches per side doesn't divide a 32×32 input. The reference code's reading (fixed patches, varying heads) is used instead.Also documented: the paper describes Tiny-ImageNet as 90,000 training images; the canonical dataset has 100,000. The real dataset is used as distributed.
Benchmark numbers are stale
The README's +5.26% was produced by the pre-fix code. It's flagged as stale in the README and
docs/algorithm.mdrather than regenerated — the paper's protocol is 200 epochs at batch 256, which needs real hardware. This is the one open item.Verification
85 tests pass; ruff, mypy, and interrogate (100% docstrings) all clean. All 10 model × resolution combinations verified end-to-end through the real trainer. LR decay, the Eq. 9 fallback, and AdamW wiring were each confirmed empirically. Tests build a synthetic Tiny-ImageNet tree rather than downloading the archive.
https://claude.ai/code/session_01AUnUVDcyMRHDib2K5YCuv5