From 2f2cc34d5d329dd597cdc9d560e3db15549dafec Mon Sep 17 00:00:00 2001 From: Bangyen Pham Date: Wed, 26 Aug 2026 10:16:34 -0400 Subject: [PATCH 1/3] fix!: align ZSharp implementation with the paper 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 --- README.md | 9 ++++-- configs/cifar100_zsharp.yaml | 11 +++---- configs/vit_zsharp.yaml | 11 +++---- configs/zsharp_baseline.yaml | 16 ++++++---- configs/zsharp_quick.yaml | 9 +++--- docs/algorithm.md | 34 ++++++++++++++------- docs/api.md | 47 ++++++++++++++++------------- scripts/experiment.py | 3 +- tests/test_optimizer.py | 57 ++++++++++++++++++++++++++++++++++++ tests/test_train.py | 20 +++++++++---- zsharp/constants.py | 29 ++++++++++-------- zsharp/optimizer.py | 42 +++++++++++++++++++------- zsharp/trainer.py | 36 +++++++++++++---------- zsharp_demo.ipynb | 10 +++---- 14 files changed, 231 insertions(+), 103 deletions(-) diff --git a/README.md b/README.md index 633d2e2..b4d4b12 100644 --- a/README.md +++ b/README.md @@ -32,9 +32,14 @@ Or open in Colab: [Colab Notebook](https://colab.research.google.com/github/bang *\*Benchmark results from full training runs. Local results may vary based on configuration.* +> **Note**: this benchmark predates the alignment of the implementation to +> the paper (see [docs/algorithm.md](docs/algorithm.md)). It was produced +> with 70th-percentile filtering, an SGD base optimizer, and gradient +> clipping, and has not been regenerated under the current defaults. + ## Features -- **Z-Score Gradient Filtering** — Intelligent gradient filtering with a default 70th percentile threshold (configurable) for improved training stability. +- **Z-Score Gradient Filtering** — Layer-wise Z-score normalization with a global 95th percentile threshold (configurable), matching the paper's $Q_p = 0.95$. - **Apple Silicon Optimization** — Up to 4.39x speedup using MPS (Metal Performance Shaders) for faster training on Mac. - **Comprehensive Testing** — 95%+ test coverage with 62 unit tests ensuring reliability and reproducibility. @@ -59,7 +64,7 @@ zsharp/ ## References -- [Sharpness-Aware Minimization with Z-Score Gradient Filtering](https://arxiv.org/html/2505.02369v3) — Original research paper by Juyoung Yun. +- [Sharpness-Aware Minimization with Z-Score Gradient Filtering](https://arxiv.org/html/2505.02369v3) — Original research paper by Juyoung Yun. The optimizer and default hyperparameters follow this paper: $Q_p = 0.95$, $\rho = 0.05$, AdamW base optimizer (lr 1e-3, weight decay 5e-5), and an LR step decay of 0.75 every 10 epochs. - [Sharpness-Aware Minimization](https://arxiv.org/abs/2010.01412) — Foundation SAM algorithm research. ## License diff --git a/configs/cifar100_zsharp.yaml b/configs/cifar100_zsharp.yaml index 1dfc8b4..4ba4fb4 100644 --- a/configs/cifar100_zsharp.yaml +++ b/configs/cifar100_zsharp.yaml @@ -1,15 +1,16 @@ +# Paper hyperparameters (arXiv:2505.02369) on CIFAR-100. dataset: cifar100 model: resnet18 optimizer: type: zsharp rho: 0.05 - percentile: 70 - lr: 0.01 - momentum: 0.9 - weight_decay: 5e-4 + percentile: 95 + lr: 0.001 + momentum: 0.9 # unused by zsharp; AdamW is the base optimizer + weight_decay: 5e-5 train: batch_size: 256 - epochs: 10 + epochs: 200 device: auto num_workers: 4 pin_memory: false diff --git a/configs/vit_zsharp.yaml b/configs/vit_zsharp.yaml index ffc5370..d260292 100644 --- a/configs/vit_zsharp.yaml +++ b/configs/vit_zsharp.yaml @@ -1,15 +1,16 @@ +# Paper hyperparameters (arXiv:2505.02369) with a ViT backbone. dataset: cifar10 model: vit_b_16 optimizer: type: zsharp rho: 0.05 - percentile: 70 - lr: 0.01 - momentum: 0.9 - weight_decay: 5e-4 + percentile: 95 + lr: 0.001 + momentum: 0.9 # unused by zsharp; AdamW is the base optimizer + weight_decay: 5e-5 train: batch_size: 128 # Smaller batch size for ViT - epochs: 10 + epochs: 200 device: auto num_workers: 4 pin_memory: false diff --git a/configs/zsharp_baseline.yaml b/configs/zsharp_baseline.yaml index 728e4e9..a05fc14 100644 --- a/configs/zsharp_baseline.yaml +++ b/configs/zsharp_baseline.yaml @@ -1,15 +1,19 @@ +# Paper configuration (arXiv:2505.02369, "Experimental Settings"): +# AdamW base optimizer, lr 1e-3, weight decay 5e-5, batch size 256, +# 200 epochs, Q_p = 0.95. The LR step decay (x0.75 every 10 epochs) is +# applied by the trainer. dataset: cifar10 model: resnet18 optimizer: type: zsharp rho: 0.05 - percentile: 70 - lr: 0.01 - momentum: 0.9 - weight_decay: 5e-4 + percentile: 95 + lr: 0.001 + momentum: 0.9 # unused by zsharp; AdamW is the base optimizer + weight_decay: 5e-5 train: - batch_size: 128 - epochs: 20 + batch_size: 256 + epochs: 200 device: auto num_workers: 4 pin_memory: false diff --git a/configs/zsharp_quick.yaml b/configs/zsharp_quick.yaml index fddc441..bb15dce 100644 --- a/configs/zsharp_quick.yaml +++ b/configs/zsharp_quick.yaml @@ -1,12 +1,13 @@ +# Quick smoke test: paper hyperparameters, but only 2 epochs. dataset: cifar10 model: resnet18 optimizer: type: zsharp rho: 0.05 - percentile: 70 - lr: 0.01 - momentum: 0.9 - weight_decay: 5e-4 + percentile: 95 + lr: 0.001 + momentum: 0.9 # unused by zsharp; AdamW is the base optimizer + weight_decay: 5e-5 train: batch_size: 128 epochs: 2 diff --git a/docs/algorithm.md b/docs/algorithm.md index a50924d..5371132 100644 --- a/docs/algorithm.md +++ b/docs/algorithm.md @@ -43,14 +43,15 @@ For each layer $l$ with gradients $g_l$: 2. **Global filtering threshold**: $$t = \text{quantile}\left(\bigcup_l |z_l|,\; p\right)$$ - where $p$ is the percentile (default: 70). The threshold is computed over - the absolute Z-scores of **all layers concatenated**, not per layer. + where $p$ is the percentile (default: 95, i.e. $Q_p = 0.95$). The + threshold is computed over the absolute Z-scores of **all layers + concatenated**, not per layer. 3. **Masking**: - $$g_l^{filtered} = g_l \odot \mathbb{I}[|z_l| \geq t]$$ - If no component in a layer passes the threshold, the top - $\lceil 0.2 \cdot \text{numel}(g_l) \rceil$ components are kept so the - layer is never fully zeroed. + $$g_l^{filtered} = g_l \odot \mathbb{I}[|z_l| > t]$$ + Because the threshold is pooled across the network, a layer whose + Z-scores are all small may be zeroed entirely. If filtering zeroes the + gradient everywhere, the unfiltered gradient is used instead (Eq. 9). 4. **SAM perturbation**: $$\epsilon = \rho \frac{g^{filtered}}{\|g^{filtered}\|_2}$$ @@ -65,10 +66,15 @@ For each layer $l$ with gradients $g_l$: | Parameter | Default | Description | |-----------|---------|-------------| | `rho` | 0.05 | SAM perturbation radius | -| `percentile` | 70 | Global filtering threshold (%) | -| `lr` | 0.01 | Learning rate | -| `momentum` | 0.9 | Momentum coefficient | -| `weight_decay` | 5e-4 | Weight decay | +| `percentile` | 95 | Global filtering threshold (%) | +| `lr` | 0.001 | Learning rate | +| `momentum` | 0.9 | Momentum coefficient (SGD baseline only) | +| `weight_decay` | 5e-5 | Weight decay | + +ZSharp uses **AdamW** as its base optimizer, with the learning rate +multiplied by 0.75 every 10 epochs, matching the paper's experimental +settings. No gradient clipping is applied. The `momentum` field applies +only to the SGD baseline and is ignored when `type: zsharp`. ## Key Benefits @@ -131,6 +137,11 @@ parameters -= state["e"] # second_step (after re-backward) | Test Accuracy | 74.89% | 80.15% | +5.26% | | Training Time | Baseline | ~4.39x faster on MPS | Speedup | +> **Note**: these numbers were produced *before* the codebase was aligned +> to the paper (70th-percentile filtering, SGD base optimizer, gradient +> clipping). They have not been regenerated and no longer describe the +> current defaults. + ### Hyperparameter Sensitivity ZSharp is robust to hyperparameter variations: @@ -149,7 +160,8 @@ ZSharp is robust to hyperparameter variations: ## Best Practices 1. **Start with defaults**: Use default hyperparameters for initial experiments -2. **Adjust percentile**: Lower percentile (50-60%) for noisy datasets +2. **Adjust percentile**: The paper ablates $Q_p \in [0.75, 0.95]$ and + reports 0.95 as best; lower values retain more components 3. **Monitor convergence**: ZSharp typically converges in fewer epochs 4. **Use appropriate batch size**: 128 works well for most cases 5. **Enable MPS**: Use Apple Silicon GPU for up to 4.39x speedup diff --git a/docs/api.md b/docs/api.md index 4e7465e..746318f 100644 --- a/docs/api.md +++ b/docs/api.md @@ -14,7 +14,7 @@ Optimizer implementations for SAM and ZSharp. - `SAM(base_optimizer, rho=0.05, **kwargs)` — Sharpness-Aware Minimization. Subclass of `torch.optim.Optimizer`. -- `ZSharp(base_optimizer, rho=0.05, percentile=70, **kwargs)` — SAM with +- `ZSharp(base_optimizer, rho=0.05, percentile=95, **kwargs)` — SAM with Z-Score gradient filtering. Subclass of `SAM`. Both take `params` (an iterable of parameters) as their first positional @@ -30,12 +30,15 @@ argument. **Parameters:** - `params`: Model parameters -- `base_optimizer`: Base optimizer class (e.g., `torch.optim.SGD`) +- `base_optimizer`: Base optimizer class. The trainer uses + `torch.optim.AdamW`, matching the paper. - `rho`: SAM perturbation radius (default: 0.05) -- `percentile`: Global filtering threshold in percent (default: 70) -- `lr`: Learning rate (default: 0.01) -- `momentum`: Momentum coefficient (default: 0.9) -- `weight_decay`: Weight decay (default: 5e-4) +- `percentile`: Global filtering threshold in percent (default: 95) +- `lr`: Learning rate (default: 0.001) +- `weight_decay`: Weight decay (default: 5e-5) + +Any further keyword arguments are forwarded to `base_optimizer`, so +`momentum` is accepted only when the base optimizer is SGD. ### `zsharp.trainer` @@ -114,12 +117,14 @@ Configuration models and default values. **Key Constants:** - `DEFAULT_SEED`: 42 -- `DEFAULT_LEARNING_RATE`: 0.01 -- `DEFAULT_MOMENTUM`: 0.9 +- `DEFAULT_LEARNING_RATE`: 0.001 +- `DEFAULT_MOMENTUM`: 0.9 (SGD baseline only) - `DEFAULT_RHO`: 0.05 -- `DEFAULT_PERCENTILE`: 70 -- `DEFAULT_WEIGHT_DECAY`: 5e-4 -- `DEFAULT_BATCH_SIZE`: 128 +- `DEFAULT_PERCENTILE`: 95 +- `DEFAULT_WEIGHT_DECAY`: 5e-5 +- `DEFAULT_BATCH_SIZE`: 256 +- `DEFAULT_LR_STEP_SIZE`: 10 +- `DEFAULT_LR_GAMMA`: 0.75 - `RESULTS_DIR`: "results" ## Configuration @@ -133,13 +138,13 @@ model: resnet18 optimizer: type: zsharp rho: 0.05 - percentile: 70 - lr: 0.01 - momentum: 0.9 - weight_decay: 5e-4 + percentile: 95 + lr: 0.001 + momentum: 0.9 # unused by zsharp; AdamW is the base optimizer + weight_decay: 5e-5 train: - batch_size: 128 - epochs: 20 + batch_size: 256 + epochs: 200 device: auto num_workers: 4 pin_memory: false @@ -188,11 +193,11 @@ import torch # Create ZSharp optimizer optimizer = ZSharp( list(model.parameters()), - base_optimizer=torch.optim.SGD, + base_optimizer=torch.optim.AdamW, rho=0.05, - percentile=70, - lr=0.01, - momentum=0.9, + percentile=95, + lr=0.001, + weight_decay=5e-5, ) # Training loop diff --git a/scripts/experiment.py b/scripts/experiment.py index 4deee2c..b435c9b 100644 --- a/scripts/experiment.py +++ b/scripts/experiment.py @@ -143,7 +143,8 @@ def run_comparison_experiments(fast_mode: bool = False) -> dict: def run_hyperparameter_study() -> dict: """Run hyperparameter study for percentile threshold as mentioned in the paper""" - percentiles = [50, 60, 70, 80, 90] + # The paper's Table 2 ablates Q_p over {0.75, 0.80, 0.85, 0.90, 0.95}. + percentiles = [75, 80, 85, 90, 95] results = {} for percentile in percentiles: diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 243fff2..af6972e 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -698,6 +698,63 @@ def test_sam_vs_zsharp_behavior(self): assert diverged, "SAM and ZSharp should behave differently" + def test_zsharp_falls_back_to_unfiltered_gradient(self): + """Test Eq. 9 fallback when filtering removes every component. + + With percentile=100 the threshold is the maximum absolute Z-score, + and the strict ``>`` comparison retains nothing. The paper then + specifies using the unfiltered gradient for the ascent step. + """ + model = SimpleModel() + zsharp = ZSharp( + list(model.parameters()), + optim.SGD, + rho=0.05, + percentile=100, + lr=0.01, + ) + + x = torch.randn(4, 10) + y = torch.randint(0, 2, (4,)) + criterion = nn.CrossEntropyLoss() + criterion(model(x), y).backward() + + before = [p.grad.detach().clone() for p in model.parameters()] + zsharp.first_step() + + # Gradients are restored unfiltered, and the ascent step still runs. + for original, p in zip(before, model.parameters(), strict=True): + assert torch.equal(original, p.grad) + assert all("e" in zsharp.state[p] for p in model.parameters()) + + def test_zsharp_may_zero_an_entire_layer(self): + """Test that filtering can zero a whole layer. + + The threshold is pooled across the network, so a layer whose + Z-scores are all small contributes nothing to the ascent direction. + The paper permits this; only an all-zero gradient triggers Eq. 9. + """ + model = SimpleModel() + zsharp = ZSharp( + list(model.parameters()), + optim.SGD, + rho=0.05, + percentile=50, + lr=0.01, + ) + + params = list(model.parameters()) + # One layer with a wide Z-score spread, one perfectly uniform (all + # Z-scores zero, so nothing in it can exceed a positive threshold). + params[0].grad = torch.randn_like(params[0]) * 10 + for p in params[1:]: + p.grad = torch.full_like(p, 0.5) + + zsharp.first_step() + + assert torch.count_nonzero(params[0].grad) > 0 + assert torch.count_nonzero(params[1].grad) == 0 + def test_sam_first_step_with_existing_state(self): """Test SAM first step when parameter already has state in optimizer""" model = SimpleModel() diff --git a/tests/test_train.py b/tests/test_train.py index 0f240b6..3797f56 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -239,8 +239,12 @@ def test_train_zsharp_optimizer( mock_model = SimpleTestModel() mock_get_model.return_value = mock_model - # Mock ZSharp optimizer + # Mock ZSharp optimizer. ``base_optimizer`` must be a real + # optimizer because the trainer attaches the LR scheduler to it. mock_optimizer = MagicMock() + mock_optimizer.base_optimizer = torch.optim.AdamW( + mock_model.parameters(), lr=0.001 + ) mock_zsharp.return_value = mock_optimizer # Mock data @@ -519,8 +523,14 @@ def test_train_results_saving(self, mock_get_model, mock_get_dataset): @patch("zsharp.trainer.get_dataset") @patch("zsharp.trainer.get_model") - def test_train_gradient_clipping(self, mock_get_model, mock_get_dataset): - """Test that gradient clipping is applied""" + def test_train_no_gradient_clipping( + self, mock_get_model, mock_get_dataset + ): + """Test that no gradient clipping is applied. + + The paper (arXiv:2505.02369) specifies no gradient clipping, so the + trainer must not rescale gradients before the optimizer step. + """ # Mock dataset mock_trainloader = MagicMock() mock_testloader = MagicMock() @@ -566,8 +576,8 @@ def test_train_gradient_clipping(self, mock_get_model, mock_get_dataset): ): train(config) - # Check that gradient clipping was called - mock_clip.assert_called() + # Check that gradient clipping was not called + mock_clip.assert_not_called() @patch("zsharp.trainer.get_dataset") @patch("zsharp.trainer.get_model") diff --git a/zsharp/constants.py b/zsharp/constants.py index df4aeef..a3e29dc 100644 --- a/zsharp/constants.py +++ b/zsharp/constants.py @@ -21,27 +21,30 @@ CIFAR10_DATASET = "cifar10" CIFAR100_DATASET = "cifar100" -# Default batch and training parameters -DEFAULT_BATCH_SIZE = 128 +# Default batch and training parameters. Batch size matches the paper; the +# epoch default below stays low deliberately, since the paper's 200 epochs +# is a poor default for an unattended run. The shipped configs set it. +DEFAULT_BATCH_SIZE = 256 DEFAULT_NUM_WORKERS = 2 DEFAULT_PIN_MEMORY = False # Optimizer constants +# Paper defaults (arXiv:2505.02369, "Experimental Settings"): AdamW with +# lr 1e-3 and weight decay 5e-5, and Q_p = 0.95, which keeps the top 5% of +# gradient components by absolute Z-score. DEFAULT_RHO = 0.05 -DEFAULT_PERCENTILE = 70 -DEFAULT_LEARNING_RATE = 0.01 +DEFAULT_PERCENTILE = 95 +DEFAULT_LEARNING_RATE = 1e-3 DEFAULT_MOMENTUM = 0.9 -DEFAULT_WEIGHT_DECAY = 5e-4 +DEFAULT_WEIGHT_DECAY = 5e-5 -# Numerical stability constants -EPSILON = 1e-12 -EPSILON_STD = 1e-8 - -# Gradient clipping -MAX_GRADIENT_NORM = 1.0 +# Learning rate schedule: multiplied by 0.75 every 10 epochs. +DEFAULT_LR_STEP_SIZE = 10 +DEFAULT_LR_GAMMA = 0.75 -# Z-score filtering constants -DEFAULT_TOP_K_RATIO = 0.2 # Keep top 20% if no gradients pass threshold +# Numerical stability constant (delta in the paper). +EPSILON = 1e-8 +EPSILON_STD = 1e-8 # Model architecture constants RESNET18_NAME = "resnet18" diff --git a/zsharp/optimizer.py b/zsharp/optimizer.py index aaaaf78..ddfa585 100644 --- a/zsharp/optimizer.py +++ b/zsharp/optimizer.py @@ -19,7 +19,6 @@ from zsharp.constants import ( DEFAULT_PERCENTILE, DEFAULT_RHO, - DEFAULT_TOP_K_RATIO, EPSILON, EPSILON_STD, MAX_QUANTILE_NUMEL, @@ -146,6 +145,12 @@ class ZSharp(SAM): This helps focus on the most important gradients and improves training stability. + Following the paper (arXiv:2505.02369), Z-scores are normalized within + each layer, the threshold is the ``percentile``-th quantile of the + absolute Z-scores pooled across all layers, and filtering is applied to + the ascent step only. If the filtered gradient vanishes everywhere, the + unfiltered gradient is used instead (Eq. 9). + Args: params: Parameters to optimize base_optimizer: Base optimizer class (e.g., torch.optim.SGD) @@ -283,20 +288,37 @@ def _apply_gradient_filtering( ) -> None: """Apply filtering mask to gradients based on threshold. + Components whose absolute Z-score does not exceed the threshold are + zeroed. Whole layers may be zeroed out, which the paper permits: the + threshold is pooled across the network, so a layer with uniformly + small Z-scores contributes nothing to the ascent direction. Only if + *every* layer is zeroed does the filtering back off, restoring the + unfiltered gradients per Eq. 9. + Args: layer_grad_info: Metadata to map back to parameters. zscores_list: Precomputed Z-scores. threshold: Absolute Z-score threshold. """ + retained = 0 for i, (p, original_grad) in enumerate(layer_grad_info): - layer_zscores = zscores_list[i] - mask = layer_zscores.abs() >= threshold + mask = (zscores_list[i].abs() > threshold).view_as(original_grad) + retained += int(mask.any()) + p.grad = cast("torch.Tensor", p.grad) * mask - if not mask.any(): - top_k = max(1, int(DEFAULT_TOP_K_RATIO * mask.numel())) - _, indices = torch.topk(layer_zscores.abs(), top_k) - mask = torch.zeros_like(mask) - mask[indices] = True + if not retained: + self._restore_unfiltered_gradients(layer_grad_info) - mask = mask.view_as(original_grad) - p.grad = cast("torch.Tensor", p.grad) * mask + @staticmethod + def _restore_unfiltered_gradients( + layer_grad_info: list[tuple[torch.nn.Parameter, torch.Tensor]], + ) -> None: + """Undo filtering when it zeroed the gradient everywhere (Eq. 9). + + Args: + layer_grad_info: Parameters paired with their pre-filter + gradients. Masking rebinds ``p.grad`` rather than mutating + it, so the originals are still intact. + """ + for p, original_grad in layer_grad_info: + p.grad = original_grad diff --git a/zsharp/trainer.py b/zsharp/trainer.py index 438cb02..e3c8b89 100644 --- a/zsharp/trainer.py +++ b/zsharp/trainer.py @@ -29,8 +29,9 @@ AUTO_DEVICE, CPU_DEVICE, CUDA_DEVICE, + DEFAULT_LR_GAMMA, + DEFAULT_LR_STEP_SIZE, DEFAULT_SEED, - MAX_GRADIENT_NORM, MPS_DEVICE, RESULTS_DIR, SGD_OPTIMIZER, @@ -143,13 +144,13 @@ def _setup_optimizer( ) return optimizer, SGD_OPTIMIZER - # ZSharp optimizer + # ZSharp optimizer. The paper trains with AdamW as the base optimizer, + # so ``momentum`` is an SGD-only setting and is not forwarded here. optimizer = ZSharp( params, - base_optimizer=optim.SGD, + base_optimizer=optim.AdamW, rho=float(opt_config.rho), lr=lr, - momentum=momentum, weight_decay=wd, percentile=int(opt_config.percentile), ) @@ -174,9 +175,6 @@ def _run_train_step( outputs = ctx.model(x) loss = ctx.criterion(outputs, y) loss.backward() - torch.nn.utils.clip_grad_norm_( - ctx.model.parameters(), MAX_GRADIENT_NORM - ) zsharp_opt = cast("ZSharp", ctx.optimizer) zsharp_opt.first_step() # Zero before the second backward: the update must use the gradient @@ -190,9 +188,6 @@ def _run_train_step( outputs = ctx.model(x) loss = ctx.criterion(outputs, y) loss.backward() - torch.nn.utils.clip_grad_norm_( - ctx.model.parameters(), MAX_GRADIENT_NORM - ) ctx.optimizer.step() return float(loss.item()), outputs.detach() @@ -281,12 +276,13 @@ def _prepare_training( tuple[DataLoader[torch.Tensor], DataLoader[torch.Tensor]], int, str, + optim.lr_scheduler.StepLR, ]: """Prepare training context and loaders. Returns: - tuple: Training context, data loaders, epoch count, and the - resolved optimizer type. + tuple: Training context, data loaders, epoch count, the resolved + optimizer type, and the learning rate scheduler. """ cfg = config.train m, opt, opt_type = _init_components(config, device) @@ -303,7 +299,16 @@ def _prepare_training( num_workers=int(cfg.num_workers), pin_memory=cfg.pin_memory, ) - return ctx, ldrs, int(cfg.epochs), opt_type + # Step decay from the paper: lr is multiplied by 0.75 every 10 epochs. + # For ZSharp the schedule is attached to the base optimizer, which is + # what actually applies the update. + scheduled = getattr(opt, "base_optimizer", opt) + scheduler = optim.lr_scheduler.StepLR( + scheduled, + step_size=DEFAULT_LR_STEP_SIZE, + gamma=DEFAULT_LR_GAMMA, + ) + return ctx, ldrs, int(cfg.epochs), opt_type, scheduler @dataclass(frozen=True) @@ -342,8 +347,8 @@ def train(config: TrainingConfig) -> Optional[ExperimentResults]: """Train a model using the provided configuration.""" set_seed(DEFAULT_SEED) device = get_device(config) - ctx, (train_ldr, test_ldr), epochs, opt_type = _prepare_training( - config, device + ctx, (train_ldr, test_ldr), epochs, opt_type, scheduler = ( + _prepare_training(config, device) ) start_time = time.time() l_list, t_list, v_list = [], [], [] @@ -351,6 +356,7 @@ def train(config: TrainingConfig) -> Optional[ExperimentResults]: try: for epoch in range(epochs): e_loss, a = _run_epoch(ctx, epoch, train_ldr) + scheduler.step() va, _ = _validate(ctx, test_ldr) l_list.append(e_loss) t_list.append(a) diff --git a/zsharp_demo.ipynb b/zsharp_demo.ipynb index a1598c1..c33341a 100644 --- a/zsharp_demo.ipynb +++ b/zsharp_demo.ipynb @@ -165,14 +165,14 @@ " optimizer_sgd = optim.SGD(\n", " model_sgd.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4\n", " )\n", + " # ZSharp follows the paper: AdamW base optimizer, Q_p = 0.95.\n", " optimizer_zsharp = ZSharp(\n", " list(model_zsharp.parameters()),\n", - " base_optimizer=optim.SGD,\n", + " base_optimizer=optim.AdamW,\n", " rho=0.05,\n", - " percentile=70,\n", - " lr=0.01,\n", - " momentum=0.9,\n", - " weight_decay=1e-4,\n", + " percentile=95,\n", + " lr=0.001,\n", + " weight_decay=5e-5,\n", " )\n", "\n", " criterion = nn.CrossEntropyLoss()\n", From 2a81e377b4dea36af1a5e50a5483a8fdf1219816 Mon Sep 17 00:00:00 2001 From: Bangyen Pham Date: Wed, 26 Aug 2026 11:00:44 -0400 Subject: [PATCH 2/3] feat: add the paper's models and Tiny-ImageNet MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- README.md | 5 +- configs/resnet56_zsharp.yaml | 17 ++ configs/tiny_imagenet_zsharp.yaml | 18 ++ configs/vit_paper_zsharp.yaml | 18 ++ docs/algorithm.md | 34 +++ docs/api.md | 56 ++++- tests/test_data.py | 97 ++++++++- tests/test_models.py | 83 ++++++++ tests/test_train.py | 2 +- zsharp/constants.py | 20 ++ zsharp/data.py | 186 ++++++++++++++++- zsharp/models.py | 331 +++++++++++++++++++++++++++++- zsharp/trainer.py | 7 +- 13 files changed, 855 insertions(+), 19 deletions(-) create mode 100644 configs/resnet56_zsharp.yaml create mode 100644 configs/tiny_imagenet_zsharp.yaml create mode 100644 configs/vit_paper_zsharp.yaml diff --git a/README.md b/README.md index b4d4b12..3205f4d 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,8 @@ Or open in Colab: [Colab Notebook](https://colab.research.google.com/github/bang - **Z-Score Gradient Filtering** — Layer-wise Z-score normalization with a global 95th percentile threshold (configurable), matching the paper's $Q_p = 0.95$. - **Apple Silicon Optimization** — Up to 4.39x speedup using MPS (Metal Performance Shaders) for faster training on Mac. -- **Comprehensive Testing** — 95%+ test coverage with 62 unit tests ensuring reliability and reproducibility. +- **Paper Architectures** — CIFAR-style ResNet-56/110, VGG-16BN, and the paper's compact ViTs, on CIFAR-10/100 and Tiny-ImageNet. +- **Comprehensive Testing** — 95%+ test coverage with 85 unit tests ensuring reliability and reproducibility. ## Repo Structure @@ -49,7 +50,7 @@ Or open in Colab: [Colab Notebook](https://colab.research.google.com/github/bang zsharp/ ├── zsharp_demo.ipynb # Colab notebook demo ├── scripts/ # Training and experiment scripts -├── tests/ # Unit/integration tests (62 tests) +├── tests/ # Unit/integration tests (85 tests) ├── docs/ # Documentation and training curves ├── configs/ # Configuration files ├── results/ # Experimental results diff --git a/configs/resnet56_zsharp.yaml b/configs/resnet56_zsharp.yaml new file mode 100644 index 0000000..bd5de01 --- /dev/null +++ b/configs/resnet56_zsharp.yaml @@ -0,0 +1,17 @@ +# Paper configuration (arXiv:2505.02369) with the CIFAR-style ResNet-56. +dataset: cifar10 +model: resnet56 +optimizer: + type: zsharp + rho: 0.05 + percentile: 95 + lr: 0.001 + momentum: 0.9 # unused by zsharp; AdamW is the base optimizer + weight_decay: 5e-5 +train: + batch_size: 256 + epochs: 200 + device: auto + num_workers: 4 + pin_memory: false + use_mixed_precision: false diff --git a/configs/tiny_imagenet_zsharp.yaml b/configs/tiny_imagenet_zsharp.yaml new file mode 100644 index 0000000..0d8444c --- /dev/null +++ b/configs/tiny_imagenet_zsharp.yaml @@ -0,0 +1,18 @@ +# Paper configuration (arXiv:2505.02369) on Tiny-ImageNet (200 classes, +# 64x64). The dataset is downloaded on first use. +dataset: tiny_imagenet +model: resnet56 +optimizer: + type: zsharp + rho: 0.05 + percentile: 95 + lr: 0.001 + momentum: 0.9 # unused by zsharp; AdamW is the base optimizer + weight_decay: 5e-5 +train: + batch_size: 256 + epochs: 200 + device: auto + num_workers: 4 + pin_memory: false + use_mixed_precision: false diff --git a/configs/vit_paper_zsharp.yaml b/configs/vit_paper_zsharp.yaml new file mode 100644 index 0000000..5b35669 --- /dev/null +++ b/configs/vit_paper_zsharp.yaml @@ -0,0 +1,18 @@ +# Paper configuration (arXiv:2505.02369) with ViT-7/8/8-384: 7 layers, +# 8 heads, 8 patches per side, MLP dimension 384. +dataset: cifar10 +model: vit_7_8_8_384 +optimizer: + type: zsharp + rho: 0.05 + percentile: 95 + lr: 0.001 + momentum: 0.9 # unused by zsharp; AdamW is the base optimizer + weight_decay: 5e-5 +train: + batch_size: 256 + epochs: 200 + device: auto + num_workers: 4 + pin_memory: false + use_mixed_precision: false diff --git a/docs/algorithm.md b/docs/algorithm.md index 5371132..77ba63b 100644 --- a/docs/algorithm.md +++ b/docs/algorithm.md @@ -128,6 +128,40 @@ parameters += parameters.grad * scale # first_step parameters -= state["e"] # second_step (after re-backward) ``` +## Architectures and Datasets + +The paper evaluates ResNet-56/110, VGG-16BN, and compact ViTs on CIFAR-10, +CIFAR-100, and Tiny-ImageNet; all are implemented here. Two details the +paper leaves unstated were taken from the author's reference +implementation ([YUNBLAK/Sharpness-Aware-Minimization-with-Z-Score-Gradient-Filtering](https://github.com/YUNBLAK/Sharpness-Aware-Minimization-with-Z-Score-Gradient-Filtering)): + +- **ResNet style.** The paper cites He et al. but does not say which + family. Depths 56 and 110 exist only as CIFAR-style ResNets (6n+2 + layers, 16/32/64 channels, option-A shortcuts), and the reference code + confirms this. Parameter counts match the published table: 0.85M for + ResNet-56, 1.7M for ResNet-110. +- **ViT dimensions.** The paper reads `ViT-7/8/8-384` as layers / heads / + patch size / MLP dimension, but that is not self-consistent: it makes + both variants 8-headed with patch size 8, leaving the differing third + field unexplained, and the `12` of `ViT-7/8/12-768` does not divide a + 32x32 input as a patch count. The reference code fixes patches at 8 per + side and varies heads (8 and 12) at an embedding width of 384 that the + paper never states, which is the reading implemented here. + +- **VGG-16BN.** The reference implementation uses a CIFAR-adapted head — + global average pooling into a single 512-unit linear classifier — rather + than torchvision's three 4096-wide ImageNet layers, which carry roughly + nine times the parameters on 32x32 inputs. + +Normalization statistics and augmentation are also unspecified in the +paper. CIFAR uses conventional per-dataset statistics and Tiny-ImageNet +its own commonly cited values, with random crop, horizontal flip, and +normalization throughout, matching the reference implementation. + +Note that the paper describes Tiny-ImageNet as "90,000 training and +10,000 test images", while the canonical dataset has 100,000 training +images. The real dataset is used as distributed. + ## Experimental Results ### Performance Metrics diff --git a/docs/api.md b/docs/api.md index 746318f..2f0e04f 100644 --- a/docs/api.md +++ b/docs/api.md @@ -60,7 +60,8 @@ Training utilities and the main training loop. ### `zsharp.data` -Data loading and preprocessing utilities for CIFAR-10 and CIFAR-100. +Data loading and preprocessing utilities for CIFAR-10, CIFAR-100, and +Tiny-ImageNet. **Functions:** @@ -71,6 +72,14 @@ Data loading and preprocessing utilities for CIFAR-10 and CIFAR-100. - `get_cifar100(batch_size=128, num_workers=2, *, pin_memory=False)`: CIFAR-100 data loaders +**Classes:** + +- `TinyImageNet(root, *, train=True, download=True, transform=None)`: + Tiny-ImageNet-200 dataset. Not distributed through torchvision, so the + archive is downloaded and extracted on first use. The validation split + ships as a flat directory, and its labels are resolved from + `val_annotations.txt`. + **Data:** - `DATASET_METADATA`: Registry of normalization statistics, class counts, @@ -78,8 +87,15 @@ Data loading and preprocessing utilities for CIFAR-10 and CIFAR-100. **Supported Datasets:** -- `cifar10`: CIFAR-10 dataset -- `cifar100`: CIFAR-100 dataset +- `cifar10`: CIFAR-10 dataset (10 classes, 32x32) +- `cifar100`: CIFAR-100 dataset (100 classes, 32x32) +- `tiny_imagenet`: Tiny-ImageNet-200 (200 classes, 64x64) + +The paper does not state normalization statistics or augmentation for any +dataset. CIFAR uses the conventional per-dataset statistics; Tiny-ImageNet +uses the commonly cited Tiny-ImageNet values. All three apply the same +augmentation as the author's reference implementation: random crop with +padding, horizontal flip, then normalization. ### `zsharp.models` @@ -87,15 +103,45 @@ Model loading utilities. **Functions:** -- `get_model(model_name="resnet18", num_classes=10) -> nn.Module`: Get a - PyTorch model by name +- `get_model(model_name="resnet18", num_classes=10, image_size=32) -> nn.Module`: + Get a PyTorch model by name. `image_size` is used by the paper's ViT + variants to derive their patch size. + +**Classes:** + +- `CifarResNet(blocks_per_stage, num_classes=10)`: CIFAR-style ResNet with + `6n+2` layers (He et al., Sec. 4.2) +- `PaperViT(num_classes=10, image_size=32, *, ...)`: the paper's compact + Vision Transformer **Supported Models:** +Architectures used in the ZSharp paper: + +- `resnet56`, `resnet110`: CIFAR-style ResNets — three stages of basic + blocks at 16/32/64 channels with parameter-free (option A) shortcuts. + These depths exist only in the CIFAR family; torchvision does not ship + them, so they are implemented here. +- `vgg16_bn`: VGG-16 with batch normalization +- `vit_7_8_8_384`, `vit_7_8_12_768`: compact ViTs with 7 layers, 8 heads, + and an embedding width of 384 + +Also available: + - `resnet18`: ResNet-18 architecture - `vgg11`: VGG-11 architecture - `vit_b_16`: Vision Transformer B-16 +> **On the ViT naming**: the paper writes `ViT-7/8/8-384` and +> `ViT-7/8/12-768`, and its text reads the fields as layers / heads / +> patch size / MLP dimension. That reading 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 author's reference implementation fixes patches at 8 per side +> and varies heads (8 and 12) at a constant embedding width of 384, which +> is what is implemented here. The second field is patches *per side*, not +> pixels: on a 32x32 input, 8 per side gives 4x4 pixel patches. + ### `zsharp.constants` Configuration models and default values. diff --git a/tests/test_data.py b/tests/test_data.py index 9f3f013..d6f45ea 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -6,7 +6,12 @@ import pytest import torch -from zsharp.data import get_cifar10, get_cifar100, get_dataset +from zsharp.data import ( + TinyImageNet, + get_cifar10, + get_cifar100, + get_dataset, +) # Patch the dataset classes so tests exercise the loader wiring without # downloading the real CIFAR datasets (which are ~163MB each and would @@ -175,3 +180,93 @@ def test_only_train_loader_shuffles(self): assert isinstance(test.sampler, torch.utils.data.SequentialSampler), ( "test loader must not shuffle" ) + + +def _build_fake_tiny_imagenet(root, wnids=("n01443537", "n01629819")): + """Write a minimal Tiny-ImageNet tree matching the real archive layout. + + Args: + root: Directory to hold the ``tiny-imagenet-200`` folder. + wnids: Class IDs to synthesize. + + Returns: + pathlib.Path: The dataset directory that was created. + """ + from PIL import Image + + base = root / "tiny-imagenet-200" + base.mkdir(parents=True) + (base / "wnids.txt").write_text("\n".join(wnids) + "\n") + + for index, wnid in enumerate(wnids): + images = base / "train" / wnid / "images" + images.mkdir(parents=True) + for j in range(2): + Image.new("RGB", (64, 64), (index * 40, j * 40, 0)).save( + images / f"{wnid}_{j}.JPEG" + ) + + val_images = base / "val" / "images" + val_images.mkdir(parents=True) + lines = [] + for k in range(4): + filename = f"val_{k}.JPEG" + Image.new("RGB", (64, 64), (k * 30, 0, 0)).save(val_images / filename) + lines.append(f"{filename}\t{wnids[k % len(wnids)]}\t0\t0\t63\t63") + (base / "val" / "val_annotations.txt").write_text("\n".join(lines) + "\n") + return base + + +class TestTinyImageNet: + """Test cases for the Tiny-ImageNet dataset.""" + + def test_train_split_indexes_per_class_directories(self, tmp_path): + """Train images live under train//images/.""" + _build_fake_tiny_imagenet(tmp_path) + dataset = TinyImageNet(root=str(tmp_path), train=True, download=False) + + assert len(dataset) == 4 + assert dataset.class_to_idx == {"n01443537": 0, "n01629819": 1} + image, label = dataset[0] + assert image.shape == (3, 64, 64) + assert label in (0, 1) + + def test_val_split_resolves_labels_from_annotations(self, tmp_path): + """Val ships flat, so labels come from val_annotations.txt.""" + _build_fake_tiny_imagenet(tmp_path) + dataset = TinyImageNet(root=str(tmp_path), train=False, download=False) + + assert len(dataset) == 4 + # Classes alternate across the four synthetic val images. + assert sorted(dataset[i][1] for i in range(4)) == [0, 0, 1, 1] + + def test_grayscale_images_become_rgb(self, tmp_path): + """Tiny-ImageNet contains grayscale images that must be converted.""" + from PIL import Image + + base = _build_fake_tiny_imagenet(tmp_path) + Image.new("L", (64, 64), 128).save( + base / "train" / "n01443537" / "images" / "gray.JPEG" + ) + dataset = TinyImageNet(root=str(tmp_path), train=True, download=False) + + assert all(dataset[i][0].shape[0] == 3 for i in range(len(dataset))) + + def test_missing_dataset_without_download_raises(self, tmp_path): + """Never silently download when the caller opted out.""" + with pytest.raises(RuntimeError, match="not found"): + TinyImageNet(root=str(tmp_path / "absent"), download=False) + + def test_get_dataset_routes_to_tiny_imagenet(self, tmp_path): + """get_dataset must wire tiny_imagenet through to the loaders.""" + _build_fake_tiny_imagenet(tmp_path) + with patch("zsharp.data.DATA_ROOT", str(tmp_path)): + train, test = get_dataset( + "tiny_imagenet", batch_size=2, num_workers=0 + ) + + assert isinstance(train.sampler, torch.utils.data.RandomSampler) + assert isinstance(test.sampler, torch.utils.data.SequentialSampler) + images, labels = next(iter(train)) + assert images.shape == (2, 3, 64, 64) + assert labels.shape == (2,) diff --git a/tests/test_models.py b/tests/test_models.py index c11c7da..8e96222 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -52,6 +52,89 @@ def test_get_model_vit_b_16(self): output = model(x) assert output.shape == (1, 10) # Updated expected shape + def test_get_model_resnet56(self): + """Test get_model with the CIFAR-style resnet56""" + model = get_model("resnet56", num_classes=10) + + assert isinstance(model, nn.Module) + + # CIFAR ResNets take 32x32 inputs, not the ImageNet 224x224. + output = model(torch.randn(1, 3, 32, 32)) + assert output.shape == (1, 10) + + # ResNet-56 is 6n+2 with n=9, and He et al. report 0.85M params. + conv_and_fc = sum( + 1 for m in model.modules() if isinstance(m, (nn.Conv2d, nn.Linear)) + ) + assert conv_and_fc == 56 + params = sum(p.numel() for p in model.parameters()) + assert 0.8e6 < params < 0.9e6 + + def test_get_model_resnet110(self): + """Test get_model with the CIFAR-style resnet110""" + model = get_model("resnet110", num_classes=100) + + assert isinstance(model, nn.Module) + output = model(torch.randn(1, 3, 32, 32)) + assert output.shape == (1, 100) + + # 6n+2 with n=18, reported at 1.7M params. + conv_and_fc = sum( + 1 for m in model.modules() if isinstance(m, (nn.Conv2d, nn.Linear)) + ) + assert conv_and_fc == 110 + params = sum(p.numel() for p in model.parameters()) + assert 1.6e6 < params < 1.8e6 + + def test_get_model_vgg16_bn(self): + """Test get_model with the CIFAR-adapted vgg16_bn""" + model = get_model("vgg16_bn", num_classes=10) + + assert isinstance(model, nn.Module) + + # Runs natively on 32x32, unlike the torchvision ImageNet VGG. + output = model(torch.randn(1, 3, 32, 32)) + assert output.shape == (1, 10) + + # A single 512-unit classifier, not torchvision's 4096-wide stack, + # so the model is ~15M parameters rather than ~134M. + assert isinstance(model.classifier, nn.Linear) + assert model.classifier.in_features == 512 + params = sum(p.numel() for p in model.parameters()) + assert params < 20e6 + + @pytest.mark.parametrize( + ("name", "num_heads", "mlp_hidden"), + [("vit_7_8_8_384", 8, 384), ("vit_7_8_12_768", 12, 768)], + ) + def test_get_model_paper_vit(self, name, num_heads, mlp_hidden): + """Test the paper's compact ViT variants. + + Both variants use 8 patches per side and differ in head count and + MLP width, so both must run at the native CIFAR resolution. + """ + model = get_model(name, num_classes=10, image_size=32) + + assert isinstance(model, nn.Module) + assert model.patches_per_side == 8 + assert len(model.encoder.layers) == 7 + + layer = model.encoder.layers[0] + assert layer.self_attn.num_heads == num_heads + assert layer.linear1.out_features == mlp_hidden + + output = model(torch.randn(2, 3, 32, 32)) + assert output.shape == (2, 10) + + def test_paper_vit_patch_size_follows_image_size(self): + """Test that patch size is derived from the input resolution. + + The paper's third field is patches per side, so a 64x64 + Tiny-ImageNet input yields 8x8 pixel patches rather than 4x4. + """ + assert get_model("vit_7_8_8_384", image_size=32).patch_size == 4 + assert get_model("vit_7_8_8_384", image_size=64).patch_size == 8 + def test_get_model_unknown_model(self): """Test get_model with unknown model raises error""" with pytest.raises(ValueError, match="Unknown model"): diff --git a/tests/test_train.py b/tests/test_train.py index 3797f56..0016116 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -399,7 +399,7 @@ def test_train_cifar100(self, mock_get_model, mock_get_dataset): # Check that model was created with correct num_classes mock_get_model.assert_called_with( - model_name="resnet18", num_classes=100 + model_name="resnet18", num_classes=100, image_size=32 ) assert isinstance(results, ExperimentResults) diff --git a/zsharp/constants.py b/zsharp/constants.py index a3e29dc..8f87d62 100644 --- a/zsharp/constants.py +++ b/zsharp/constants.py @@ -20,6 +20,12 @@ # Dataset names CIFAR10_DATASET = "cifar10" CIFAR100_DATASET = "cifar100" +TINY_IMAGENET_DATASET = "tiny_imagenet" + +# Tiny-ImageNet is not distributed through torchvision; it is downloaded +# from the canonical Stanford CS231n mirror and extracted under DATA_ROOT. +TINY_IMAGENET_URL = "http://cs231n.stanford.edu/tiny-imagenet-200.zip" +TINY_IMAGENET_DIRNAME = "tiny-imagenet-200" # Default batch and training parameters. Batch size matches the paper; the # epoch default below stays low deliberately, since the paper's 200 epochs @@ -49,6 +55,20 @@ # Model architecture constants RESNET18_NAME = "resnet18" +# CIFAR-style ResNets (He et al., Sec. 4.2): three stages of basic blocks +# starting at 16 channels and doubling, for a depth of 6n + 2. +CIFAR_RESNET_BASE_WIDTH = 16 +CIFAR_RESNET_STAGES = 3 + +# The paper's compact ViTs: 7 layers, 8 heads, an embedding width of 384, +# and 8 patches per side. Taken from the author's reference implementation +# (github.com/YUNBLAK/Sharpness-Aware-Minimization-with-Z-Score-Gradient-Filtering), +# since the paper itself does not state the embedding dimension. +VIT_PAPER_LAYERS = 7 +VIT_PAPER_HEADS = 8 +VIT_PAPER_HIDDEN = 384 +VIT_PAPER_PATCHES_PER_SIDE = 8 + # Optimizer types SGD_OPTIMIZER = "sgd" ZSHARP_OPTIMIZER = "zsharp" diff --git a/zsharp/data.py b/zsharp/data.py index 30c696d..8f0b4f6 100644 --- a/zsharp/data.py +++ b/zsharp/data.py @@ -1,22 +1,29 @@ # Copyright (c) 2025 Bangyen Pham -"""Data loading utilities for CIFAR-10 and CIFAR-100 datasets. +"""Data loading utilities for the datasets used in the ZSharp paper. -This module provides functions to load and preprocess CIFAR-10 and CIFAR-100 -datasets with appropriate data augmentation and normalization. +This module provides functions to load and preprocess CIFAR-10, CIFAR-100 +and Tiny-ImageNet with appropriate data augmentation and normalization. """ -from typing import Union +from collections.abc import Callable +from pathlib import Path +from typing import Optional, Union, cast import torch import torch.utils.data import torchvision import torchvision.transforms as T +from PIL import Image +from torchvision.datasets.utils import download_and_extract_archive from zsharp.constants import ( DATA_ROOT, DEFAULT_BATCH_SIZE, DEFAULT_NUM_WORKERS, DEFAULT_PIN_MEMORY, + TINY_IMAGENET_DATASET, + TINY_IMAGENET_DIRNAME, + TINY_IMAGENET_URL, ) # Dataset metadata registry @@ -36,6 +43,16 @@ "image_size": 32, "crop_padding": 4, }, + # Tiny-ImageNet: 200 classes of 64x64 images. The paper does not state + # normalization statistics, so the commonly cited Tiny-ImageNet values + # are applied here; augmentation mirrors the CIFAR recipe above. + TINY_IMAGENET_DATASET: { + "mean": (0.4802, 0.4481, 0.3975), + "std": (0.2770, 0.2691, 0.2821), + "num_classes": 200, + "image_size": 64, + "crop_padding": 8, + }, } _DATASET_CLASSES: dict[str, type[torchvision.datasets.VisionDataset]] = { @@ -115,6 +132,158 @@ def _get_cifar( return trainloader, testloader +class TinyImageNet(torch.utils.data.Dataset[tuple[torch.Tensor, int]]): + """Tiny-ImageNet-200: 200 classes of 64x64 images. + + Not available through torchvision, so this downloads and extracts the + canonical archive on first use. The train split is laid out per class, + but the validation split ships as a flat directory whose labels live in + ``val_annotations.txt``, so both are resolved to a common index here. + """ + + def __init__( + self, + root: str = DATA_ROOT, + *, + train: bool = True, + download: bool = True, + transform: Optional[Callable[[Image.Image], torch.Tensor]] = None, + ) -> None: + """Build the sample index for one split. + + Args: + root: Directory holding (or to receive) the extracted dataset. + train: Load the training split rather than the validation split. + download: Fetch the archive if it is not already present. + transform: Optional transform applied to each PIL image. + + Raises: + RuntimeError: If the dataset is missing and ``download`` is + False. + """ + self.transform = transform + self.root = Path(root) / TINY_IMAGENET_DIRNAME + + if not self.root.exists(): + if not download: + msg = ( + f"Tiny-ImageNet not found at {self.root}. " + "Pass download=True to fetch it." + ) + raise RuntimeError(msg) + download_and_extract_archive(TINY_IMAGENET_URL, download_root=root) + + wnids = sorted((self.root / "wnids.txt").read_text().split()) + self.class_to_idx = {wnid: i for i, wnid in enumerate(wnids)} + self.samples = self._train_samples() if train else self._val_samples() + + def _train_samples(self) -> list[tuple[Path, int]]: + """Index the per-class training directories.""" + samples: list[tuple[Path, int]] = [] + for wnid, label in self.class_to_idx.items(): + image_dir = self.root / "train" / wnid / "images" + samples.extend( + (path, label) for path in sorted(image_dir.glob("*.JPEG")) + ) + return samples + + def _val_samples(self) -> list[tuple[Path, int]]: + """Index the flat validation directory via its annotations file.""" + annotations = self.root / "val" / "val_annotations.txt" + samples = [] + for line in annotations.read_text().splitlines(): + if not line.strip(): + continue + filename, wnid = line.split("\t")[:2] + path = self.root / "val" / "images" / filename + samples.append((path, self.class_to_idx[wnid])) + return samples + + def __len__(self) -> int: + """Return the number of samples in the split. + + Returns: + int: Sample count. + """ + return len(self.samples) + + def __getitem__(self, index: int) -> tuple[torch.Tensor, int]: + """Load one sample. + + Args: + index: Position of the sample within the split. + + Returns: + tuple: The transformed image and its class index. + """ + path, label = self.samples[index] + with Image.open(path) as image: + # Tiny-ImageNet contains a few grayscale images. + converted = image.convert("RGB") + if self.transform is not None: + return self.transform(converted), label + return T.functional.to_tensor(converted), label + + +def _get_tiny_imagenet( + batch_size: int = DEFAULT_BATCH_SIZE, + num_workers: int = DEFAULT_NUM_WORKERS, + *, + pin_memory: bool = DEFAULT_PIN_MEMORY, +) -> tuple[ + torch.utils.data.DataLoader[torch.Tensor], + torch.utils.data.DataLoader[torch.Tensor], +]: + """Load Tiny-ImageNet with train and test data loaders. + + Args: + batch_size: Batch size for data loaders + num_workers: Number of worker processes for data loading + pin_memory: Whether to pin memory for faster GPU transfer + + Returns: + tuple: (train_loader, test_loader) for Tiny-ImageNet + + """ + meta = DATASET_METADATA[TINY_IMAGENET_DATASET] + transform_train = T.Compose( + [ + T.RandomCrop(meta["image_size"], padding=meta["crop_padding"]), + T.RandomHorizontalFlip(), + T.ToTensor(), + T.Normalize(meta["mean"], meta["std"]), + ], + ) + transform_test = T.Compose( + [ + T.ToTensor(), + T.Normalize(meta["mean"], meta["std"]), + ], + ) + + trainset = TinyImageNet(train=True, transform=transform_train) + testset = TinyImageNet(train=False, transform=transform_test) + + trainloader = torch.utils.data.DataLoader( + trainset, + batch_size=batch_size, + shuffle=True, + num_workers=num_workers, + pin_memory=pin_memory, + ) + testloader = torch.utils.data.DataLoader( + testset, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + pin_memory=pin_memory, + ) + return ( + cast("torch.utils.data.DataLoader[torch.Tensor]", trainloader), + cast("torch.utils.data.DataLoader[torch.Tensor]", testloader), + ) + + def get_cifar10( batch_size: int = DEFAULT_BATCH_SIZE, num_workers: int = DEFAULT_NUM_WORKERS, @@ -184,7 +353,8 @@ def get_dataset( """Get dataset by name with train and test data loaders. Args: - dataset_name: Name of the dataset ('cifar10' or 'cifar100') + dataset_name: Name of the dataset ('cifar10', 'cifar100' or + 'tiny_imagenet') batch_size: Batch size for data loaders num_workers: Number of worker processes for data loading pin_memory: Whether to pin memory for faster GPU transfer @@ -203,5 +373,11 @@ def get_dataset( num_workers=num_workers, pin_memory=pin_memory, ) + if dataset_name == TINY_IMAGENET_DATASET: + return _get_tiny_imagenet( + batch_size=batch_size, + num_workers=num_workers, + pin_memory=pin_memory, + ) error_msg = f"Unknown dataset: {dataset_name}" raise ValueError(error_msg) diff --git a/zsharp/models.py b/zsharp/models.py index 1e8eedc..c119efb 100644 --- a/zsharp/models.py +++ b/zsharp/models.py @@ -3,26 +3,334 @@ This module provides functions to load and configure different PyTorch models including ResNet, VGG, and Vision Transformer variants. + +Alongside the torchvision ImageNet models, this module implements the +architectures used in the ZSharp paper (arXiv:2505.02369): CIFAR-style +ResNets (He et al., Sec. 4.2) and the paper's small Vision Transformers. """ -from typing import cast +from typing import TYPE_CHECKING, cast +import torch from torch import nn from torchvision import models from torchvision.models import vit_b_16 -from zsharp.constants import RESNET18_NAME +if TYPE_CHECKING: + from collections.abc import Callable + +from zsharp.constants import ( + CIFAR_RESNET_BASE_WIDTH, + CIFAR_RESNET_STAGES, + RESNET18_NAME, + VIT_PAPER_HEADS, + VIT_PAPER_HIDDEN, + VIT_PAPER_LAYERS, + VIT_PAPER_PATCHES_PER_SIDE, +) + + +class _CifarBasicBlock(nn.Module): + """Two-convolution residual block for CIFAR-style ResNets.""" + + def __init__(self, in_planes: int, planes: int, stride: int = 1) -> None: + """Initialize the block. + + Args: + in_planes: Number of input channels. + planes: Number of output channels. + stride: Stride of the first convolution. + """ + super().__init__() + self.conv1 = nn.Conv2d( + in_planes, + planes, + kernel_size=3, + stride=stride, + padding=1, + bias=False, + ) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d( + planes, + planes, + kernel_size=3, + stride=1, + padding=1, + bias=False, + ) + self.bn2 = nn.BatchNorm2d(planes) + + # Option A shortcut from He et al. Sec. 4.2: downsample by striding + # and zero-pad the extra channels, keeping the block parameter-free. + self.pad_channels = 0 + self.stride = stride + if stride != 1 or in_planes != planes: + self.pad_channels = (planes - in_planes) // 2 + + def _shortcut(self, x: torch.Tensor) -> torch.Tensor: + """Apply the parameter-free option-A shortcut.""" + if self.pad_channels == 0 and self.stride == 1: + return x + subsampled = x[:, :, :: self.stride, :: self.stride] + return nn.functional.pad( + subsampled, + (0, 0, 0, 0, self.pad_channels, self.pad_channels), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run the residual block. + + Args: + x: Input feature map. + + Returns: + torch.Tensor: Output feature map. + """ + out = nn.functional.relu(self.bn1(self.conv1(x))) + out = self.bn2(self.conv2(out)) + out = out + self._shortcut(x) + return nn.functional.relu(out) + + +class CifarResNet(nn.Module): + """CIFAR-style ResNet with 6n+2 layers (He et al., Sec. 4.2). + + Three stages of ``n`` basic blocks operate at 16, 32 and 64 channels on + 32x32 inputs. This differs from the torchvision ImageNet ResNets, which + use a 7x7 stem and max-pooling; depths such as 56 and 110 exist only in + this CIFAR family. + """ + + def __init__(self, blocks_per_stage: int, num_classes: int = 10) -> None: + """Initialize the network. + + Args: + blocks_per_stage: Number of blocks ``n`` in each of the three + stages, giving a depth of ``6n + 2``. + num_classes: Number of output classes. + """ + super().__init__() + width = CIFAR_RESNET_BASE_WIDTH + self.conv1 = nn.Conv2d( + 3, + width, + kernel_size=3, + stride=1, + padding=1, + bias=False, + ) + self.bn1 = nn.BatchNorm2d(width) + + stages, final_width = self._build_stages(blocks_per_stage, width) + self.layers = nn.Sequential(*stages) + self.fc = nn.Linear(final_width, num_classes) + + for module in self.modules(): + if isinstance(module, (nn.Conv2d, nn.Linear)): + nn.init.kaiming_normal_(module.weight) + + @staticmethod + def _build_stages( + blocks_per_stage: int, width: int + ) -> tuple[list[nn.Module], int]: + """Build the three residual stages. + + Args: + blocks_per_stage: Number of blocks in each stage. + width: Channel count of the first stage. + + Returns: + tuple: The blocks and the final channel count. + """ + blocks: list[nn.Module] = [] + in_planes = width + for stage in range(CIFAR_RESNET_STAGES): + planes = width * (2**stage) + for block in range(blocks_per_stage): + # Each stage after the first halves the spatial resolution. + stride = 2 if stage > 0 and block == 0 else 1 + blocks.append(_CifarBasicBlock(in_planes, planes, stride)) + in_planes = planes + return blocks, in_planes + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Classify a batch of images. + + Args: + x: Input images of shape ``(batch, 3, height, width)``. + + Returns: + torch.Tensor: Class logits of shape ``(batch, num_classes)``. + """ + out = nn.functional.relu(self.bn1(self.conv1(x))) + out = self.layers(out) + out = nn.functional.adaptive_avg_pool2d(out, 1).flatten(1) + return cast("torch.Tensor", self.fc(out)) + + +class CifarVGG16BN(nn.Module): + """VGG-16 with batch normalization, adapted for small images. + + Matches the author's reference implementation: the standard VGG-16 + convolutional stack followed by global average pooling and a single + 512-unit linear classifier. The torchvision ImageNet VGG instead ends + in three 4096-wide layers, which on a 32x32 input would upsample a 1x1 + feature map to 7x7 and carry roughly nine times the parameters. + """ + + # Standard VGG-16 ("D") convolutional configuration. + # fmt: off + CONFIG = ( + 64, 64, "M", 128, 128, "M", 256, 256, 256, "M", + 512, 512, 512, "M", 512, 512, 512, + ) + # fmt: on + + def __init__(self, num_classes: int = 10) -> None: + """Initialize the network. + + Args: + num_classes: Number of output classes. + """ + super().__init__() + layers: list[nn.Module] = [] + in_channels = 3 + for entry in self.CONFIG: + if entry == "M": + layers.append(nn.MaxPool2d(kernel_size=2, stride=2)) + continue + channels = int(entry) + layers.extend( + ( + nn.Conv2d(in_channels, channels, kernel_size=3, padding=1), + nn.BatchNorm2d(channels), + nn.ReLU(inplace=True), + ) + ) + in_channels = channels + + self.features = nn.Sequential(*layers) + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.classifier = nn.Linear(in_channels, num_classes) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Classify a batch of images. + + Args: + x: Input images of shape ``(batch, 3, height, width)``. + + Returns: + torch.Tensor: Class logits of shape ``(batch, num_classes)``. + """ + out = self.avgpool(self.features(x)).flatten(1) + return cast("torch.Tensor", self.classifier(out)) + + +class PaperViT(nn.Module): + """The compact Vision Transformer used in the ZSharp paper. + + Named ``ViT-//-`` in the paper, e.g. + ``ViT-7/8/8-384`` and ``ViT-7/8/12-768``. + + The paper's prose describes the fields as layers / heads / patch size / + MLP dimension, but that reading is not self-consistent: it would make + 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 author's reference implementation fixes patches at 8 per + side and varies heads (8 and 12) at a constant embedding width of 384, + which is the reading used here. + + Note that the second field is the number of patches *per side*, not the + patch size in pixels: on a 32x32 input, 8 patches per side means each + patch is 4x4 pixels. Patches are embedded by flattening pixels and + applying a single linear projection. + """ + + def __init__( # noqa: PLR0913 + self, + num_classes: int = 10, + image_size: int = 32, + *, + patches_per_side: int = VIT_PAPER_PATCHES_PER_SIDE, + num_layers: int = VIT_PAPER_LAYERS, + num_heads: int = VIT_PAPER_HEADS, + hidden: int = VIT_PAPER_HIDDEN, + mlp_hidden: int = VIT_PAPER_HIDDEN, + ) -> None: + """Initialize the transformer. + + Args: + num_classes: Number of output classes. + image_size: Height and width of the input images. + patches_per_side: Number of patches along each spatial axis. + num_layers: Number of transformer encoder layers. + num_heads: Number of self-attention heads. + hidden: Embedding dimension. + mlp_hidden: Feed-forward dimension, the paper's trailing field. + """ + super().__init__() + self.patches_per_side = patches_per_side + self.patch_size = image_size // patches_per_side + num_patches = patches_per_side**2 + patch_dim = (self.patch_size**2) * 3 + + self.embed = nn.Linear(patch_dim, hidden) + self.cls_token = nn.Parameter(torch.randn(1, 1, hidden)) + self.pos_embed = nn.Parameter(torch.randn(1, num_patches + 1, hidden)) + + encoder_layer = nn.TransformerEncoderLayer( + d_model=hidden, + nhead=num_heads, + dim_feedforward=mlp_hidden, + activation="gelu", + batch_first=True, + norm_first=True, + ) + # enable_nested_tensor is incompatible with norm_first and only + # emits a warning; disable it explicitly. + self.encoder = nn.TransformerEncoder( + encoder_layer, num_layers, enable_nested_tensor=False + ) + self.head = nn.Sequential( + nn.LayerNorm(hidden), nn.Linear(hidden, num_classes) + ) + + def _to_patches(self, x: torch.Tensor) -> torch.Tensor: + """Split images into flattened patch vectors.""" + size = self.patch_size + out = x.unfold(2, size, size).unfold(3, size, size) + out = out.permute(0, 2, 3, 4, 5, 1) + return out.reshape(x.size(0), self.patches_per_side**2, -1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Classify a batch of images. + + Args: + x: Input images of shape ``(batch, 3, height, width)``. + + Returns: + torch.Tensor: Class logits of shape ``(batch, num_classes)``. + """ + out = self.embed(self._to_patches(x)) + cls = self.cls_token.repeat(out.size(0), 1, 1) + out = torch.cat([cls, out], dim=1) + self.pos_embed + out = self.encoder(out) + return cast("torch.Tensor", self.head(out[:, 0])) def get_model( model_name: str = RESNET18_NAME, num_classes: int = 10, + image_size: int = 32, ) -> nn.Module: """Get a PyTorch model by name. Args: model_name: Name of the model to load num_classes: Number of output classes + image_size: Input resolution, used by the paper's ViT variants to + derive the patch size. Returns: torch.nn.Module: PyTorch model @@ -31,14 +339,29 @@ def get_model( ValueError: If model name is not supported """ - model_map = { + model_map: dict[str, Callable[[], nn.Module]] = { "resnet18": lambda: models.resnet18(num_classes=num_classes), "vgg11": lambda: models.vgg11(num_classes=num_classes), "vit_b_16": lambda: vit_b_16(num_classes=num_classes), + # Architectures from the ZSharp paper. + "resnet56": lambda: CifarResNet(9, num_classes), + "resnet110": lambda: CifarResNet(18, num_classes), + "vgg16_bn": lambda: CifarVGG16BN(num_classes=num_classes), + "vit_7_8_8_384": lambda: PaperViT( + num_classes=num_classes, + image_size=image_size, + mlp_hidden=384, + ), + "vit_7_8_12_768": lambda: PaperViT( + num_classes=num_classes, + image_size=image_size, + num_heads=12, + mlp_hidden=768, + ), } if model_name not in model_map: error_msg = f"Unknown model {model_name}" raise ValueError(error_msg) - return cast("nn.Module", model_map[model_name]()) + return model_map[model_name]() diff --git a/zsharp/trainer.py b/zsharp/trainer.py index e3c8b89..1d5a1f5 100644 --- a/zsharp/trainer.py +++ b/zsharp/trainer.py @@ -262,8 +262,13 @@ def _init_components( raise ValueError(error_msg) classes = cast("int", DATASET_METADATA[ds_name]["num_classes"]) + image_size = cast("int", DATASET_METADATA[ds_name]["image_size"]) model_name = config.model - model = get_model(model_name=model_name, num_classes=classes).to(device) + model = get_model( + model_name=model_name, + num_classes=classes, + image_size=image_size, + ).to(device) optimizer, opt_type = _setup_optimizer(config, model) return model, optimizer, opt_type From 0705091a38088f4683c9db481934c2bad64e2d81 Mon Sep 17 00:00:00 2001 From: Bangyen Pham Date: Wed, 26 Aug 2026 11:22:23 -0400 Subject: [PATCH 3/3] chore: point experiment tooling at the paper's architectures 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 --- configs/vit_zsharp.yaml | 7 +- scripts/ablate_percentile.py | 138 +++++++++++++++++++++++++++++++++++ scripts/experiment.py | 12 +-- 3 files changed, 148 insertions(+), 9 deletions(-) create mode 100644 scripts/ablate_percentile.py diff --git a/configs/vit_zsharp.yaml b/configs/vit_zsharp.yaml index d260292..005df08 100644 --- a/configs/vit_zsharp.yaml +++ b/configs/vit_zsharp.yaml @@ -1,6 +1,7 @@ -# Paper hyperparameters (arXiv:2505.02369) with a ViT backbone. +# Paper hyperparameters (arXiv:2505.02369) with the larger of the paper's +# two ViT variants: 7 layers, 8 patches per side, 12 heads, MLP 768. dataset: cifar10 -model: vit_b_16 +model: vit_7_8_12_768 optimizer: type: zsharp rho: 0.05 @@ -9,7 +10,7 @@ optimizer: momentum: 0.9 # unused by zsharp; AdamW is the base optimizer weight_decay: 5e-5 train: - batch_size: 128 # Smaller batch size for ViT + batch_size: 256 epochs: 200 device: auto num_workers: 4 diff --git a/scripts/ablate_percentile.py b/scripts/ablate_percentile.py new file mode 100644 index 0000000..ce7a1d8 --- /dev/null +++ b/scripts/ablate_percentile.py @@ -0,0 +1,138 @@ +# Copyright (c) 2025 Bangyen Pham +"""A/B the Z-score percentile threshold with everything else held fixed. + +Isolates the hyperparameter that moved most when the implementation was +aligned to the paper (70 -> 95), so the effect can be read without the +cost of a full 200-epoch reproduction. Each arm re-seeds identically, so +the only difference between runs is the threshold. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import sys +import time +from pathlib import Path + +from zsharp.constants import RESULTS_DIR, TrainingConfig +from zsharp.trainer import train + +logging.basicConfig( + level=logging.INFO, + format="%(message)s", + handlers=[logging.StreamHandler(sys.stdout)], +) +logger = logging.getLogger(__name__) + + +def _build_config( + percentile: int, epochs: int, model: str, dataset: str +) -> TrainingConfig: + """Build a ZSharp config that varies only in the percentile. + + Args: + percentile: Z-score filtering threshold. + epochs: Number of training epochs. + model: Model name. + dataset: Dataset name. + + Returns: + TrainingConfig: The validated configuration. + """ + return TrainingConfig.model_validate( + { + "dataset": dataset, + "model": model, + "optimizer": {"type": "zsharp", "percentile": percentile}, + "train": { + "epochs": epochs, + "batch_size": 256, + "device": "auto", + "num_workers": 4, + }, + } + ) + + +def main() -> None: + """Run one arm per percentile and report the comparison.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--percentiles", + type=int, + nargs="+", + default=[70, 95], + help="Thresholds to compare (default: the old and new defaults)", + ) + parser.add_argument("--epochs", type=int, default=10) + parser.add_argument("--model", default="resnet56") + parser.add_argument("--dataset", default="cifar10") + args = parser.parse_args() + + results = {} + for percentile in args.percentiles: + logger.info("=" * 60) + logger.info( + "percentile=%d model=%s dataset=%s epochs=%d", + percentile, + args.model, + args.dataset, + args.epochs, + ) + logger.info("=" * 60) + + start = time.time() + output = train( + _build_config(percentile, args.epochs, args.model, args.dataset) + ) + if output is None: + logger.warning("percentile=%d interrupted", percentile) + continue + + results[percentile] = { + "final_test_accuracy": output.final_test_accuracy, + "final_test_loss": output.final_test_loss, + "test_accuracies": output.test_accuracies, + "runtime": time.time() - start, + } + + logger.info("=" * 60) + logger.info("PERCENTILE ABLATION") + logger.info("=" * 60) + for percentile, result in sorted(results.items()): + logger.info( + "percentile %3d: %.2f%% (%.0fs)", + percentile, + result["final_test_accuracy"], + result["runtime"], + ) + if len(results) > 1: + best = max(results, key=lambda k: results[k]["final_test_accuracy"]) + worst = min(results, key=lambda k: results[k]["final_test_accuracy"]) + delta = ( + results[best]["final_test_accuracy"] + - results[worst]["final_test_accuracy"] + ) + logger.info("best: %d (+%.2f%% over %d)", best, delta, worst) + + out_dir = Path(RESULTS_DIR) + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "percentile_ablation.json" + with out_path.open("w") as f: + json.dump( + { + "model": args.model, + "dataset": args.dataset, + "epochs": args.epochs, + "results": results, + }, + f, + indent=2, + ) + logger.info("saved to %s", out_path) + + +if __name__ == "__main__": + main() diff --git a/scripts/experiment.py b/scripts/experiment.py index b435c9b..be0a1f7 100644 --- a/scripts/experiment.py +++ b/scripts/experiment.py @@ -17,7 +17,6 @@ from zsharp.constants import ( DEFAULT_BATCH_SIZE, DEFAULT_LEARNING_RATE, - DEFAULT_MOMENTUM, DEFAULT_NUM_WORKERS, DEFAULT_RHO, DEFAULT_WEIGHT_DECAY, @@ -73,9 +72,9 @@ def run_comparison_experiments(fast_mode: bool = False) -> dict: experiments = [ ("configs/sgd_baseline.yaml", "SGD Baseline"), ("configs/zsharp_baseline.yaml", "ZSharp"), - # Temporarily disabled for testing: - # ("configs/cifar100_zsharp.yaml", "ZSharp CIFAR-100"), - # ("configs/vit_zsharp.yaml", "ZSharp ViT"), + ("configs/resnet56_zsharp.yaml", "ZSharp ResNet-56"), + ("configs/cifar100_zsharp.yaml", "ZSharp CIFAR-100"), + ("configs/vit_paper_zsharp.yaml", "ZSharp ViT-7/8/8-384"), ] results = {} @@ -155,13 +154,14 @@ def run_hyperparameter_study() -> dict: # Create temporary config config = { "dataset": "cifar10", - "model": "resnet18", + # Table 2 ablates the percentile on ResNet-56. + "model": "resnet56", "optimizer": { + # ZSharp builds on AdamW, which takes no momentum. "type": "zsharp", "rho": DEFAULT_RHO, "percentile": percentile, "lr": DEFAULT_LEARNING_RATE, - "momentum": DEFAULT_MOMENTUM, "weight_decay": DEFAULT_WEIGHT_DECAY, }, "train": {