Efficient Polling-Based Learning Rate Optimization for Neural Networks

Get the accuracy of polling-based learning-rate selection at essentially the cost of plain SGD.

pip install efficient-polling-lr-scheduler

🇧🇷 Versão em português · 📄 Paper (LaTeX source)

This repository replicates the Polling Method of Tan et al. on CIFAR-10 and introduces Efficient Polling, a novel extension that recovers the same learning-rate schedule — and the same accuracy — while polling only 5.43% of batches, cutting optimizer steps by 79% and per-epoch wall-clock time from 8.05s to 2.81s, the cost of plain SGD. The method is benchmarked over five random seeds against eight comparison methods — Adam, three schedulers, SPS, Armijo backtracking, and the replicated base Polling method — plus two ablations that isolate the contribution of the adaptive polling trigger. A follow-up extension, Relative Polling, drops the fixed candidate grid: the user picks one rate and one multiplier, and a window that widens while blind, backed off without a cap, finds the same schedule from any start. All methods ship as a PyTorch package.

TL;DR

The learning rate is the single most influential hyperparameter in gradient-based training. Instead of picking it by hand or by a fixed schedule, polling tests several candidate learning rates at every batch and keeps the one that most improves batch accuracy. It works remarkably well, but it multiplies training time by the number of candidates.

Efficient Polling observes that the polled choice is highly redundant — within each training phase consecutive polls pick the same learning rate — and polls on demand instead: an exponential-backoff schedule doubles the gap between polls while the selection is stable, and a two-tier divergence guard protects the unpolled steps.

150 epochs, mean ± sample standard deviation over five seeds (42–46), NVIDIA RTX 5070. The full comparison — Adam, three schedulers, SPS, Armijo backtracking, two trigger-ablation variants and the Relative Polling extension, thirteen configurations in all — is in Results.

Efficient Polling matches the base method's accuracy (within 0.1 pp on test) while cutting optimizer steps by 79% and per-epoch time from 8.05s to 2.81s — a 2.9× speed-up over base Polling, indistinguishable from plain SGD's 2.86s.

Quickstart

pip install efficient-polling-lr-scheduler

Polling needs to re-evaluate the model to score a candidate step, so instead of the bare optimizer.step() you pass a closure that returns (loss, score) — the same contract as torch.optim.LBFGS, plus the score to maximize. make_closure builds it for you:

import torch
from efficient_polling_lr_scheduler import EfficientPollingSGD, make_closure

model = MyModel().to(device)
loss_fn = torch.nn.CrossEntropyLoss()

# Candidate LRs default to {1e-5, 1e-4, 1e-3, 1e-2, 1e-1} around lr.
optimizer = EfficientPollingSGD(model, lr=1e-3)

for inputs, targets in train_loader:
    inputs, targets = inputs.to(device), targets.to(device)
    info = optimizer.step(make_closure(model, loss_fn, inputs, targets))
    # info.lr, info.loss, info.polled, info.spike, info.rolled_back, ...

No learning-rate schedule, no warmup, no tuning: the learning rate is measured. Swap EfficientPollingSGD for PollingSGD to get the base method (polls every batch), or wrap any optimizer you like:

from efficient_polling_lr_scheduler import EfficientPollingOptimizer

optimizer = EfficientPollingOptimizer(
    torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9),
    candidate_lrs=(1e-5, 1e-4, 1e-3, 1e-2, 1e-1),
    module=model,          # so BatchNorm buffers are restored between trials
    max_poll_interval=64,  # backoff cap; 0 polls every batch
)

The package also ships two comparison baselines used in the paper — the stochastic Polyak step-size (SPS) and Armijo backtracking line search, both of which derive the step size from the current batch at no extra forward pass:

from efficient_polling_lr_scheduler import SPSSGD, ArmijoSGD

optimizer = SPSSGD(model, lr=1e-3, max_lr=0.1)
# or
optimizer = ArmijoSGD(model, lr=1e-3, lr_max=0.1)

The optional training helpers run a full comparison in a few lines, and accept a plain optimizer (and an optional torch.optim.lr_scheduler) too — so the baseline, Adam and the standard schedulers go through the same loop:

from efficient_polling_lr_scheduler import fit

history = fit(model, train_loader, val_loader, optimizer, loss_fn, epochs=150)
print(history.best_val_acc, sum(history.polls), sum(history.optimizer_steps))

API

Notes. Despite the distribution name, these are not torch.optim.lr_scheduler.LRScheduler subclasses: they wrap the optimizer and are driven entirely through optimizer.step(closure), so there is no separate scheduler.step() to call after it. Candidate learning rates are absolute and applied to every parameter group, overriding per-group learning rates. Pass module= (or the model itself as the first argument) whenever the forward pass mutates buffers, so trials cannot leak BatchNorm statistics. The closure must not call backward() or zero_grad() — the optimizer owns both.

How it works

Polling (replicated base method)

At each batch, after computing the gradient g from weights θ, every candidate learning rate is applied as a trial step and the winner is kept:

θ̂ₖ = θ − lrₖ · g               for each lrₖ ∈ C
k*  = argmax acc(θ̂ₖ, batch)     (ties favour the smallest lr)
θ   ← θ̂ₖ*

The candidate set is C = {1e-5, 1e-4, 1e-3, 1e-2, 1e-1}, spanning four orders of magnitude around the base LR. The trial updates are realized by snapshotting the model + optimizer state once, then reloading it before each candidate step, so every candidate departs from an identical pre-step condition. Because the trials overwrite the weights, the winning step is reapplied once from the snapshot, so a poll costs N + 1 = 6 optimizer steps rather than N = 5.

Efficient Polling (proposed extension)

The selection mechanism is untouched, but a batch is polled only when needed:

1. Adaptive polling schedule. Let K be the poll interval. After a poll, if the selection is unchanged, K ← min(2K, K_max) (geometric backoff, capped at K_max = 64); if it changed, K ← 1 (poll every batch until it stabilizes again). Between polls, a single blind SGD step uses the last selected LR. Backoff only engages once a poll has actually shown signal — i.e. the candidates' accuracies differ. At initialization every candidate ties, and treating a tie as a stable selection would stall training at the smallest candidate LR forever; the trigger ablation reproduces exactly this failure in a variant that lacks the rule.

2. Two-tier divergence guard. Blind steps have no per-step validation, so a high-LR step can diverge. The guard reuses quantities already computed:

   • Tier 2 — spike-triggered polls (prevention): if the batch loss exceeds γ · EMA(loss) (γ = 3, β = 0.9), poll immediately so the accuracy criterion can reject an explosive step.
   • Tier 1 — rollback checkpoints (recovery): each poll snapshot doubles as a known-good checkpoint; if the loss is non-finite or exceeds 2·ln(C) ≈ 4.61, restore the checkpoint and resume polling.

Across the five official runs, spike-triggered polls numbered 378 per run on average (6.6% of all polls) and the recovery tier fired only twice in the entire study, both in the same seed. Its necessity is real, though: an early unguarded run diverged to NaN at epoch 33 from a single blind step at lr = 1e-1 and never recovered — the same failure that killed most of the hand-designed schedulers in the comparison (see Results).

3. Alternative triggers (ablation). trigger="fixed" polls every K + 1 batches at a constant interval instead of backing off, and trigger="random" polls each batch independently with probability p. Both keep rule 2 (the guard) and the selection mechanism unchanged, and both are calibrated to the ~5% poll rate the adaptive backoff measures, so a comparison against them isolates the trigger rather than the budget. See Trigger ablation.

Relative Polling (proposed extension, per batch)

Both methods above choose from a fixed grid, and in the recorded runs the selected rate spends most of training pinned to that grid's edges: 1e-1 through the first phase, 1e-5 after the anneal. Relative Polling drops the grid. The user picks one learning rate and one multiplier m, and each poll tries three candidates around the rate in use; the winner becomes the new centre:

C_t = {X/m, X, X·m}        X ← argmax acc(θ̂, batch)

When a poll comes back blind — every candidate scoring the same, which batch accuracy does whenever one step moves no prediction — the next poll looks one multiplier farther out in both directions, and keeps widening until it sees a difference; a poll with signal narrows the window back. Three rules complete it:

1. Uncapped backoff, bounded by failure. The poll interval k doubles when a scheduled poll with signal keeps the rate, a blind poll leaves it unchanged, and it has no K_max. Every poll checkpoints the best point seen so far — weights, optimizer state and rate, judged by a slow trend of the loss. When a blind stretch blows up, training goes back in time to that point, k drops to zero and the ceiling of the next slow start becomes half the interval that blew up; growth is exponential up to the ceiling and linear above it, as in TCP congestion control.
2. Ties keep the rate, except after a restart. A tie carries no information, so an ordinary poll keeps X. The poll right after a restart breaks ties one notch down, because a restart has a single cause — a rate too high for blind steps.
3. Adam-style trend. The loss trend is a bias-corrected exponential mean with a second moment of its deviations, so a spike is a loss more than z deviations above the trend rather than a fixed ratio to it.

from efficient_polling_lr_scheduler import RelativePollingSGD

optimizer = RelativePollingSGD(model, lr=1e-3, multiplier=10.0, lr_max=1e-1)

Results. Same protocol as the table below, five seeds, candidates capped at 1e-1:

Per batch, Relative Polling reaches the test accuracy of Armijo backtracking (84.09%, the best in the table below) at the cost of plain SGD, polling 7% of batches; a poll costs 4 optimizer steps instead of 6, so it takes fewer steps than Efficient Polling. It finds the same first phase at 1e-1, then settles at 1e-2 from epoch ~40 to the end instead of annealing to 1e-5: once batch accuracy saturates the criterion goes blind, ties keep the rate, and no blow-up forced a notch down (one restart in five runs). That plateau is what costs it test loss — training continues at 1e-2 on a training set it already fits, so its predictions grow overconfident. The poll fraction varies by seed (1.8% to 16.3%): a seed that keeps hunting between 1e-1 and 1e-2 resets its interval on every move.

The one rate the user picks does not have to be right. Started two decades below or above the default, on seed 42:

From either start the relative window is at 1e-1 within the first epoch and reproduces the same schedule.

Per epoch is a negative result. RelativeEpochPolling, driven by fit(..., epoch_polling=...), trains a whole epoch per candidate from one snapshot and keeps the one with the lowest mean training loss. It walks the rate down to 1e-5 within thirty epochs and stalls at 48%; selecting on end-of-epoch validation accuracy instead did the same, to 1e-8 and 45%. Any single-epoch score rewards the smoothness of a small step over the progress of a large one, and a blind epoch at 1e-1 has none of the per-step protection the batch method gets from its spike polls. The granularity that works is the batch.

Results

Table below reproduces the paper's Table I, eleven configurations, plus the two Relative Polling configurations added after the paper; each run on five seeds (42–46), 150 epochs, batch size 64 (704 batches/epoch, 105,600/run). Reported as mean ± sample standard deviation.

The three schedulers start at 1e-1 (the top of the candidate set) rather than at the baseline 1e-3, since a decay schedule needs somewhere to decay from; SPS and Armijo are capped at that same 1e-1, so no method may take a step the others were never allowed to consider. Despite that, cosine annealing, step decay and ReduceLROnPlateau diverge to NaN around epoch 28 on most of their seeds (5/5, 4/5 and 1/5 respectively) — the table still credits them with their best pre-divergence checkpoint, since every method is scored at its own best validation epoch. Both Polling variants hold that same 1e-1 for roughly thirty epochs across their 25 combined runs without a single failure: what breaks the schedulers is not the rate itself but the absence of a per-step check on it.

Both polling methods autonomously discover the same two-phase schedule entirely from batch-level feedback: the mean selected LR converges to ≈1e-1 within the first epoch, holds there for ~30 epochs, then collapses toward ≈1e-5 for fine refinement near convergence — base Polling completes the anneal between epochs 42–45, Efficient Polling more gradually, between epochs 49–78 (the interval hasn't reset to one on every batch yet).

Training & validation loss, all thirteen configurations.	Mean selected LR per epoch, symlog, with a cross marking divergence.

Polls concentrate exactly where the schedule changes: the steady-state floor is 704 / (K_max + 1) ≈ 11 polls/epoch, the median over all epochs is 16.4, and the count peaks at 231 in epoch 32 — the exact epoch the selected LR begins collapsing from 1e-1 to 1e-5, where disagreeing polls keep resetting the interval to one. This is the mechanism that lets ~5% of the polls recover the full two-phase schedule.

Cost model

A poll costs N + 1 = 6 optimizer steps (one trial per candidate, plus reapplying the winner). With P polls out of B = 105,600 batches:

S_eff = P·(N+1) + (B − P) = B + N·P

With the P = 5,732.2 polls measured on average over the five runs, this gives 134,261 steps — a 79% reduction versus base Polling's 633,600, exactly matching Table above. Only 26% of those steps come from polls; the remaining 74% are ordinary SGD steps, which is why the per-epoch wall-clock time sits at the level of plain SGD.

Trigger ablation

Two control variants isolate the contribution of the adaptive backoff by replacing only the trigger, keeping the candidate set, selection rule and two-tier guard unchanged: fixed polls at a constant interval (K = 19) and random polls each batch independently with probability p = 0.05, both calibrated to the ~5% rate the backoff measures.

On final accuracy the random trigger is competitive — 84.02% test vs. 83.76% for the backoff, well inside the seed-to-seed spread. This is an honest negative result for the strong reading of the claim: at this budget, distributing polls uniformly at random is enough to track the schedule, provided the guard absorbs the cost of arriving late. The supported claim is the weaker one — the backoff reaches the same quality while spending its polls where they carry information (16.4 polls in a median epoch vs. a peak of 231 at the transition, a 14× ratio, against a flat ~40/epoch for both controls), needing ~2.4× fewer spike-triggered guard interventions (378 vs. 919 and 888) and slightly fewer optimizer steps.

The fixed-interval control exposes a real failure mode: on 4/5 seeds it matches the other variants (84.17% ± 0.96% val), but on the remaining seed it never leaves the initialization plateau, ending at ~10% (random-guess) accuracy — because at initialization every candidate ties, the poll keeps returning the smallest candidate LR by the tie-break rule, and at 1e-5 the weights move too little to ever break the tie. The adaptive backoff is immune by construction, since it refuses to back off until a poll has actually discriminated between candidates (rule 1 above).

Repository structure

.
├── src/efficient_polling_lr_scheduler/     # the installable package
│   ├── polling.py             # base method (Tan et al.)
│   ├── efficient.py           # Efficient Polling (ours), incl. the fixed/random triggers
│   ├── relative.py            # Relative Polling (ours): widening window, TCP-style backoff, best-point restarts
│   ├── baselines.py           # SPS and Armijo backtracking comparison optimizers
│   ├── _snapshot.py           # exact state save/restore for trial steps
│   ├── closures.py            # batch closures and selection criteria
│   └── training.py            # optional fit/train_epoch/evaluate helpers
├── tests/                     # pytest suite for the algorithms
├── examples/
│   ├── cifar10.py             # reproduces all thirteen configurations from the CLI
│   └── plot_results.py        # redraws the figures from the recorded runs
├── notebooks/
│   └── cifar10.ipynb          # original experiments: data, model, all 13 methods, plots
├── docs/
│   └── main.tex                # the paper (IEEE format)
├── images/                    # figures used in the paper and this README
├── results/cifar10/           # the 69 recorded runs: 13 configurations × 5 seeds, plus 4 initial-rate runs
├── models/                    # best checkpoints per method (.pt, gitignored)
├── pyproject.toml
├── CHANGELOG.md
├── README.md
└── README(pt-br).md

Setup

To use the methods, all you need is the package (Python 3.10+, PyTorch 2.0+):

pip install efficient-polling-lr-scheduler

To reproduce the experiments, clone the repository and install with the extras. A CUDA-capable GPU is recommended (CPU works but is slow):

git clone https://github.com/luiz-linkezio/Efficient-Polling-Based-Learning-Rate-Optimization-for-Neural-Networks.git
cd Efficient-Polling-Based-Learning-Rate-Optimization-for-Neural-Networks
python -m venv venv
source venv/bin/activate
pip install -e ".[dev,examples]" jupyter

Dataset

The experiments load the CIFAR-10 Python version from a local directory (the pickled data_batch_* / test_batch files). Download it from the official site:

curl -O https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz
tar -xzf cifar-10-python.tar.gz

Running

The example script runs every configuration over one or more seeds and prints the comparison table, resuming a sweep from any results already on disk:

python examples/cifar10.py --data-dir /path/to/cifar-10-batches-py --seeds 42 43 44 45 46
# one method, one seed, shorter run:
python examples/cifar10.py --data-dir ... --methods efficient --seeds 42 --epochs 20

Results are written to --results-dir (default results/cifar10/) as one JSON file per (method, seed). examples/plot_results.py redraws the figures from those files, so a plot can never disagree with the table:

python examples/plot_results.py --results-dir results/cifar10 --out-dir images

Run the test suite with pytest.

Alternatively, open the original notebook and run the cells top to bottom, pointing DATA_DIR (in the Constants cell) at the extracted cifar-10-batches-py directory:

jupyter notebook notebooks/cifar10.ipynb

The notebook is organized as: Imports → Constants → Configs (seeds 42–46, device) → Data (dataset, normalization stats, 90/10 train/val split) → Model (SimpleCIFAR10CNN, ~0.56M params) → Train (all thirteen configurations, one shared loop) → Animations & plots → Test. Best checkpoints are written to models/.

Reproducibility. Five seeds (42–46) each fix weight init, data shuffling and the train/val split, so on a given seed every method starts from the same weights and sees the same batch order. All numbers above are the mean ± sample standard deviation over the five runs.

Experimental setup

• Dataset: CIFAR-10 — 45,000 train / 5,000 val / 10,000 test, normalized per channel with training statistics.
• Model: SimpleCIFAR10CNN, a 5-layer CNN (64→64→128→128→256 conv channels, 3×3 kernels, ReLU, MaxPool, AdaptiveAvgPool, Linear head), 557,898 parameters, no batch norm or dropout so the optimizer is the only source of adaptation.
• Optimizer: vanilla SGD (no momentum, no weight decay) for the proposed and replicated methods, batch size 64, base LR 1e-3, 150 epochs (704 batches/epoch, 105,600 total).
• Comparison methods: Adam, cosine annealing, step decay, ReduceLROnPlateau, SPS (Polyak step-size), Armijo backtracking line search, and the replicated base Polling method — eight in total, plus two trigger-ablation variants of the proposed method.
• Seeds: five (42–46) per configuration, thirteen configurations, 65 runs total, plus four initial-rate robustness runs on seed 42.
• Hardware: single NVIDIA GeForce RTX 5070 (12 GB).

Citation

If you use this work, please cite the paper:

@misc{souzasilva_efficient_polling_lr_scheduler,
  title  = {Efficient Polling-Based Learning Rate Optimization for Neural Networks},
  author = {de Souza Silva, Jos{\'e} R. and B. A. da Silva, Luiz Henrique and B. de Souza, Caio B. and Balieiro, Andson M.},
  year   = {2026},
  note   = {Centro de Inform{\'a}tica (CIn), Universidade Federal de Pernambuco (UFPE)},
  url    = {https://github.com/luiz-linkezio/Efficient-Polling-Based-Learning-Rate-Optimization-for-Neural-Networks}
}

The base Polling method is from Tan et al. (see docs/base_paper.pdf). The full paper, with related work, the cost model derivation and the trigger ablation, is at docs/main.tex.

To cite the software specifically, add note = {Python package \texttt{efficient-polling-lr-scheduler}} or reference the PyPI project.

🧑‍💻 Authors

Luiz Henrique Developer LinkedIn Portfolio	José Ronaldo Developer LinkedIn Portfolio

Universidade Federal de Pernambuco, Recife, Brazil. The paper additionally credits Caio B. B. de Souza (UPE) and Andson M. Balieiro (CIn/UFPE) — see Citation.

License

MIT — see the LICENSE file in this repository.