diff --git a/README.md b/README.md
index 633d2e2..3205f4d 100644
--- a/README.md
+++ b/README.md
@@ -32,11 +32,17 @@ 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.
+- **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
 
@@ -44,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
@@ -59,7 +65,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/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/configs/vit_zsharp.yaml b/configs/vit_zsharp.yaml
index ffc5370..005df08 100644
--- a/configs/vit_zsharp.yaml
+++ b/configs/vit_zsharp.yaml
@@ -1,15 +1,17 @@
+# 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
-  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
+  batch_size: 256
+  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..77ba63b 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
 
@@ -122,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
@@ -131,6 +171,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 +194,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..2f0e04f 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`
 
@@ -57,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:**
 
@@ -68,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,
@@ -75,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`
 
@@ -84,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.
@@ -114,12 +163,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 +184,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 +239,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/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 4deee2c..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 = {}
@@ -143,7 +142,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:
@@ -154,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": {
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/<wnid>/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_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..0016116 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
@@ -395,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)
@@ -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..8f87d62 100644
--- a/zsharp/constants.py
+++ b/zsharp/constants.py
@@ -20,32 +20,55 @@
 # Dataset names
 CIFAR10_DATASET = "cifar10"
 CIFAR100_DATASET = "cifar100"
+TINY_IMAGENET_DATASET = "tiny_imagenet"
 
-# Default batch and training parameters
-DEFAULT_BATCH_SIZE = 128
+# 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
+# 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
-
-# Numerical stability constants
-EPSILON = 1e-12
-EPSILON_STD = 1e-8
+DEFAULT_WEIGHT_DECAY = 5e-5
 
-# 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"
 
+# 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-<layers>/<patches>/<heads>-<mlp>`` 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/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..1d5a1f5 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()
 
@@ -267,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
 
@@ -281,12 +281,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 +304,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 +352,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 +361,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",