Machine Learning · Optimization

Gradient Noise, Batch Size, and Step Decay in Stochastic Gradient Descent

Why sampling a subset of the data makes the gradient unbiased but noisy, how that noise sets a floor under constant step sizes, and what step schedules can and cannot fix.

The reason stochastic gradient descent exists is arithmetic. With n examples, one full gradient costs n gradient evaluations, and near the end of training a single example already points in roughly the right direction. The reason it behaves differently from full-batch descent is statistics: the sampled gradient is a random variable, and every distinctive feature of the method, the noise floor, the rate schedules, the batch-size trade-offs, follows from the variance of that random variable rather than from its mean, which is the same quantity the deterministic analysis already covered in step sizes and the descent lemma.

Unbiased estimates and how the variance scales

Write the objective as an average over examples, f(x) = (1/n) Σ fᵢ(x), and draw a mini-batch B of size b uniformly without replacement. The estimator g = (1/b) Σ_{i∈B} ∇fᵢ(x) satisfies E[g] = ∇f(x), so the direction is right on average. Its covariance is what differs. For sampling with replacement from a population with per-example gradient covariance Σ(x), Cov(g) = Σ(x)/b: the standard deviation falls as 1/√b, so quadrupling the batch halves the noise. Without replacement the same scaling holds apart from a finite-population correction that matters only when b approaches n.

The per-example variance is not a fixed property of the dataset. It shrinks toward zero as the iterate approaches a point where all examples agree on the gradient, and it is largest where examples disagree most: around a decision boundary for a classifier, or in the early phase of training where individual examples pull in conflicting directions. This dependence on x is the reason the noise floor and the useful batch size both change during a run, and it is why the batch-size analysis is usually stated as a function of the current iterate.

The floor under a constant step size

Substituting a noisy gradient into the descent lemma gives a conditional expectation of the decrease that includes a negative term proportional to the variance. For a μ-strongly convex objective with gradient variance bounded by σ², the expected squared distance to the optimum under a constant step does not go to zero; it is bounded by a quantity on the order of ησ²/μ, and the expected function-value gap by a similar expression multiplied by L. Both shrink linearly as η → 0, and this is the definition of the noise floor: the iterate jitters inside a ball whose radius is set by the product of the step size and the noise-to-curvature ratio.

The practical consequence is uncomfortable. Reducing the step size to tighten the floor also slows the initial progress, which is proportional to η in the linear regime, so a schedule must spend its budget where the trade is favourable. It also explains an asymmetry between the two failure directions: a step that is too large produces visible divergence in the objective, whereas a step that is too small produces a plateau that is indistinguishable from convergence to a non-optimal point unless the gradient norm is inspected. Tracking ‖g‖² rather than the objective separates the two, since the noise floor holds the objective flat while the gradient norm stays bounded away from zero.

Robbins–Monro conditions and the shape of the rate

Diminishing steps resolve the floor by construction. The classical conditions on a schedule η_k are Σ η_k = ∞ and Σ η_k² < ∞: the first ensures the iterate can travel an arbitrary distance in a fixed direction, and the second ensures the accumulated variance stays finite instead of adding up to an unbounded random walk. The schedule η_k = c/k satisfies both. The schedule η_k = c/√k violates the second condition, and the resulting iterate has a persistent noise floor even though the step size goes to zero; it is nevertheless used, because it decays slowly enough to keep making progress on objectives whose gradients do not vanish near the solution.

The rates follow the schedule. For a convex, L-smooth objective, the average of the iterates under a 1/k schedule achieves a function-value gap of O(1/k), and reaching accuracy ε requires O(1/ε) iterations, the same exponent as full-batch gradient descent, but each iteration costs b example gradients instead of n. That comparison is where the win comes from, and it is also where it can evaporate: if b is comparable to n, the constant factor of the stochastic method is worse and the variance advantage has disappeared.

def sgd(params, grads_fn, batches, eta0=0.1, decay_every=20, decay=0.5):
    """Step decay: multiply the step size by decay every decay_every passes."""
    eta = eta0
    for epoch, batch in enumerate(batches):
        if epoch and epoch % decay_every == 0:
            eta *= decay
        for x, y in batch:
            g = grads_fn(params, x, y)
            params = [p - eta * gi for p, gi in zip(params, g)]
    return params

Where mini-batch acceleration stops paying

Doubling the batch halves the gradient variance, but it also doubles the work per step, so the wall-clock benefit is not automatic. The two effects cancel in the regime where the noise dominates the update: raising b by a factor of k with the step size held fixed reduces the noise by √k while costing k times the arithmetic. Keeping the amount of noise per unit of progress constant instead requires raising the step size in proportion to b, and that only works while the step is small enough to satisfy the descent lemma’s stability bound.

There is a point past which additional samples buy almost nothing. It occurs when the gradient noise stops being the limiting factor and the direction of the true gradient becomes the bottleneck; increasing the batch beyond that point wastes arithmetic on a better estimate of a direction that is already known accurately enough. The crossover scales with the noise-to-signal ratio of the gradient and moves during training, and it is the conceptual reason behind the observation that a large batch with a large step can reproduce the trajectory of a small batch with a decaying step. Batch size and step size therefore act as two handles on the same trade-off between per-step progress and per-step noise, and momentum and Nesterov acceleration change the picture further by filtering the gradient sequence over time rather than averaging it within a step.

Epoch schedules and the order of the samples

Almost all implementations iterate over a shuffled pass of the data rather than sampling independently with replacement, and the difference is not only bookkeeping. Sampling without replacement within an epoch makes the successive gradients negatively correlated conditional on the permutation: an example used early is unavailable later, so the average gradient over a full epoch equals the exact full gradient. For finite sums of smooth convex functions this random reshuffling attains a faster rate in the number of epochs than with-replacement sampling, which is why the epoch-shuffling loop is the default despite the messier analysis. On the other hand, the guarantee depends on the shuffle being uniform; a deterministic order reintroduces a systematic bias that behaves like a rotation around the optimum.

Schedules in common use are step decay by a factor of roughly 0.1 or 0.5 at fixed epoch boundaries, linear or cosine decay to zero over a fixed horizon, and inverse-time schedules of the form η_k = η₀/(1 + k/K). The practical differences are smaller than the differences between a well-chosen constant step and a badly chosen one, because all of them reduce the noise floor eventually.

Why the first steps are often the smallest

A schedule that only decays assumes the early phase is the safe one, and the opposite is frequently true. The curvature near a random initialization can exceed the curvature near the solution, whether because the parameters sit in a steeper region or because the effective scale of the model output is larger at the start. A step size that satisfies the stability bound at the end of training can therefore violate it at the start, and the observable symptom is a loss that rises for the first few hundred iterations before it begins to fall. The usual mitigation is a warmup: begin with a small step, raise it over a fixed number of iterations that is a small fraction of the total budget, and then hand over to the decay schedule. Warmup interacts with batch size in the way the descent analysis predicts, since a larger batch has smaller gradient variance and tolerates a larger step, and it interacts with the accumulator state of adaptive step sizes, whose bias correction is largest in exactly those early iterations and whose estimates of gradient magnitude are therefore least reliable where the schedule is most aggressive.

Where the picture breaks down

Three assumptions are routinely violated. Mini-batches drawn from a stream of correlated samples, such as successive frames of a video or adjacent tokens of a document, are not independent, so the variance of the batch mean does not fall as 1/b and a batch built from neighbouring examples behaves much like a single example. Objectives whose curvature changes as the iterate moves, which includes every deep network, invalidate the fixed L and μ that the guarantees assume, and a step size that is stable at initialization can become unstable later. And when the data distribution itself shifts over time, the justified schedule converges to a moving target, so a decreasing step can lock in parameters that matched an early phase of the stream. The failure is not detectably a divergence: the loss curve looks like healthy convergence to a slightly worse value.

References

  1. A Stochastic Approximation Method — Robbins & Monro, Annals of Mathematical Statistics, 1951
  2. Optimization Methods for Large-Scale Machine Learning — Bottou, Curtis & Nocedal, SIAM Review, 2018
  3. An Empirical Model of Large-Batch Training — McCandlish, Kaplan, Amodei, OpenAI, 2018
  4. Don't Decay the Learning Rate, Increase the Batch Size — Smith, Kindermans, Ying & Le, 2017