Back to Data Structures

SGD / Mini-Batch SGD

October 17, 2026 Wasil Zafar 38 min read

Concepts, invariants, and Python, C++, and Java implementations of SGD / Mini-Batch SGD

Table of Contents

  1. Why SGD / Mini-Batch SGD Matters
  2. Core Concepts and Terminology
  3. How SGD / Mini-Batch SGD Works, Step by Step
  4. Worked Example with a Full Trace
  5. Implementation in Python, C++, and Java
  6. Complexity, Correctness, and Invariants
  7. Variants, Pitfalls, and Optimizations
  8. Practice Problems and Real-World Use

Why SGD / Mini-Batch SGD Matters

For an objective built from millions of training examples, computing the full gradient before every update can be expensive. Stochastic gradient descent, or SGD, uses one sampled example; mini-batch SGD averages gradients from a small sample and makes one update. The estimate can be much cheaper than a full pass, while introducing noise into the direction.

The central tradeoff is work versus uncertainty. A larger batch averages more gradient information at the same point, often improving hardware utilization and reducing variance under suitable sampling. It also costs more gradient evaluations per update. Smaller batches make more updates under the same example-work budget, but their trajectories can fluctuate and the full objective can increase on a valid step.

This lesson uses explicit sampling with replacement, an integer-seeded random stream, and a squared-distance objective with a known minimizer. The programs implement both SGD and mini-batch SGD through the batch-size parameter. Full-dataset loss monitoring is optional and counted separately, so the demonstration does not conceal the cost of printing a full objective after each update.

SGD and mini-batch SGD replace a full empirical-risk gradient with a sampled mean of component gradients, then update the parameter vector once. State the sampling rule and batch normalization explicitly: an unbiased ideal estimate supports expectation-based analysis, while individual updates need not decrease the full objective and monitoring costs must be counted separately.

Core Concepts and Terminology

Let the empirical objective be the mean of $n$ component losses:

$$F(\mathbf w)=\frac1n\sum_{i=0}^{n-1}\ell_i(\mathbf w).$$

For update $t$, sample $b$ indices $I_{t,1},\ldots,I_{t,b}$ independently and uniformly with replacement in the ideal random model. Repeats are allowed. Compute every component gradient at the same old parameter vector, form their mean, and update once:

$$\widehat{\mathbf g}_t=\frac1b\sum_{r=1}^b\nabla\ell_{I_{t,r}}(\mathbf w_t),\qquad\mathbf w_{t+1}=\mathbf w_t-\eta\widehat{\mathbf g}_t.$$

Batch size $b=1$ gives ordinary SGD. A full gradient instead evaluates every example once at the same point. Merely setting $b=n$ in a with-replacement sampler does not produce a full gradient: some examples can appear repeatedly while others are omitted. The program permits $b\gt n$ for the same reason.

An epoch usually denotes one traversal of a finite dataset, often with reshuffling. This with-replacement implementation has no epochs or shuffled permutations. Processing $n$ sampled gradient examples gives a work-equivalent pass, not a guarantee that every example has been visited.

The teaching loss is $\ell_i(\mathbf w)=\tfrac12\lVert\mathbf w-\mathbf c_i\rVert_2^2$, with component gradient $\mathbf w-\mathbf c_i$. If $\boldsymbol\mu=n^{-1}\sum_i\mathbf c_i$, then

$$F(\mathbf w)=\frac12\lVert\mathbf w-\boldsymbol\mu\rVert_2^2+\frac1{2n}\sum_i\lVert\mathbf c_i-\boldsymbol\mu\rVert_2^2.$$

Thus the unique minimizer is $\boldsymbol\mu$, but the minimum loss need not be zero. In the worked one-dimensional data, centers $[-2,0,2,4]$ have mean $1$ and variance $5$, so $F(w)=\tfrac12(w-1)^2+2.5$, its full gradient is $w-1$, and its minimum is $F(1)=2.5$.

Mini-batch trajectory on the full objective and an SGD loss increase The full objective has minimum two point five at parameter one. Mini-batch updates from five visit three point five, two point two five, zero point one two five, and zero point five six two five. A separate stochastic single-example update from zero point three seven five to two point one eight seven five increases the full loss. Sampled gradients create noisy motion on the full objective Centers: −2, 0, 2, 4. Full F(w) = 0.5(w − 1)² + 2.5. -1 1 3 5 2.5 6.5 10.5 14.5 F w 0 1 2 3 4 A noisy SGD step Before: w = 0.375 Sample center: 4 Sample gradient: −3.625 η = 0.5 → w = 2.1875 Full F: 2.6953 → 3.2051 The full loss increases. Optimum: w = 1, F = 2.5 Dashed links show the mini-batch run; the side example is the batch-size-1 run. A valid stochastic update need not reduce the full objective on every step.
TermWorked exampleMeaning
Dataset / componentCenter $c_2=2$ defines $\ell_2(w)=(w-2)^2/2$One contribution to mean empirical loss
Parameter$w_0=5$Shared value updated by optimization
Batch indicesFirst mini-batch is $(2,2)$With-replacement sampling allows duplicates
Component gradient$5-2=3$ for each occurrenceDerivative evaluated at the old parameter
Mean batch gradient$(3+3)/2=3$Divide by actual sample count, including repeats
Learning rate$\eta=0.5$Scales one mean-gradient update
Full gradient$w-1$Derivative of the complete empirical mean
Full-loss monitoringFour component losses per displayed full objectiveDiagnostic work, separate from sampled gradients
Gradient-example budgetEight evaluations for four batches of size $2$Counts repeated sampled examples separately
Gradient noiseAt $w=2$, sampling center $2$ gives $0$ while full gradient is $1$Explains why one sampled residual is not a convergence certificate

How SGD / Mini-Batch SGD Works, Step by Step

Inputs are the initial parameter vector, a finite list of center vectors, seed, batch size $b$, constant learning rate $\eta$, update limit $K$, gradient-example budget $B$, and a full-monitoring flag.

  1. Validate and copy. Check shapes, numeric ranges, seed, and controls. Copy parameters and data. Initialize the seeded generator. Full monitoring, when enabled, evaluates the full objective at the starting point and counts $n$ example-loss evaluations. It uses no sampled gradients and consumes no randomness.
  2. Record the start. Step $0$ has an empty batch and empty gradient because no update has occurred. It always contains the starting point. If monitoring is off, the stored full-loss field is a dummy zero, and output explicitly prints OFF.
  3. Check limits before sampling. STEP_LIMIT takes precedence when $K$ updates have completed. Otherwise require room for all $b$ gradient examples; if unavailable, return GRAD_BUDGET before drawing indices. A leftover budget smaller than $b$ does not create a partial batch.
  4. Sample a full batch. Draw $b$ indices with replacement, in order. Every occurrence counts as a component-gradient evaluation, even when the same index repeats. The order also fixes floating-point accumulation.
  5. Accumulate at one frozen point. Initialize the gradient sum to zero. For each sampled index, add $\mathbf w-\mathbf c_i$ coordinate by coordinate without changing $\mathbf w$. Divide each accumulated coordinate by $b$ exactly once.
  6. Make one update. Multiply the mean gradient by $\eta$ and subtract it from the old point. Count one update and $b$ gradient examples. Zero-gradient batches and unchanged numerical proposals are counted, but they do not stop this stochastic run: the next sampled batch can have a different gradient.
  7. Monitor and record. If enabled, evaluate the full objective at the new point, charging another $n$ example-loss evaluations. Record the batch, mean gradient, new point, loss or monitoring flag, counters, and generator state, then repeat.

The algorithm does not use the monitored loss to accept or reject an update, change the learning rate, select a historical best, or stop. Consequently turning monitoring off preserves the training trajectory and random stream. There is no line search, epoch reshuffling, momentum, adaptive rate, gradient clipping, or convergence-tolerance exit in these programs.

A complete mini-batch averaging and update step At old parameter three point five, indices zero and three have centers minus two and four, producing gradients five point five and minus zero point five. Their mean is two point five. With rate one half the new parameter is two point two five and full objective is three point two eight one two five. One batch uses one old point and one mean-gradient update BEFORE MEAN GRADIENT AFTER UPDATE old w = 3.5 indices = (0, 3) centers = (−2, 4) full F = 5.625 gradients: 5.5, −0.5 sum = 5 mean = 5 / 2 = 2.5 η mean = 1.25 new w = 3.5 − 1.25 = 2.25 full F = 3.28125 two gradient examples Both component gradients are evaluated at w = 3.5; update only after averaging. Sampling is with replacement: repeated indices count as repeated gradient evaluations.
SGD / Mini-Batch SGD control flow
flowchart TD
    A["Validate, copy parameters and data, initialize RNG"] --> B["Optionally monitor full starting loss and record step zero"]
    B --> C{"Update limit reached?"}
    C -->|Yes| Y["Return STEP_LIMIT"]
    C -->|No| D{"Budget covers a complete batch?"}
    D -->|No| Z["Return GRAD_BUDGET"]
    D -->|Yes| E["Sample b indices with replacement"]
    E --> F["Sum gradients at the same old parameters"]
    F --> G["Divide by b and update parameters once"]
    G --> H["Count zero or unchanged updates without stopping"]
    H --> I["Optionally monitor full loss and record state"]
    I --> C

Worked Example with a Full Trace

Use centers $[-2,0,2,4]$, start $w_0=5$, seed $13$, batch size $b=2$, rate $\eta=0.5$, update limit $K=4$, gradient budget $B=100$, and full monitoring enabled. The initial full objective is $F(5)=10.5$.

The portable random generator follows

$$z_{k+1}=48271z_k\bmod2147483647,\qquad z_0=13.$$

Let $R=2147483646$. For a dataset of size $n$, map a raw state by $v=z-1$: reject and retry if $v\ge R-(R\bmod n)$, otherwise use index $v\bmod n$. Under a uniform raw-input model this balances residue classes. The seeded generator itself is deterministic and does not supply independent mathematical randomness. Its first state is $627523$, giving index $(627523-1)\bmod4=2$.

All four mini-batch updates are:

UpdateOld wIndices / centersComponent gradients at old wMean gradientNew wFull F(new w)
$1$$5$2,2 / 2,23,3$3$$3.5$$5.625000000$
$2$$3.5$0,3 / -2,45.5,-0.5$2.5$$2.25$$3.281250000$
$3$$2.25$0,0 / -2,-24.25,4.25$4.25$$0.125$$2.882812500$
$4$$0.125$2,1 / 2,0-1.875,0.125$-0.875$$0.5625$$2.595703125$

The repeated first batch $(2,2)$ evaluates the same component twice, each time at $w=5$. Both gradients equal $3$, so their mean is $3$ and the new point is $3.5$. The second batch uses centers $-2$ and $4$, giving gradients $5.5$ and $-0.5$ at the same old point $3.5$; their average is $2.5$, producing $2.25$.

At step $3$, both samples are center $-2$. Their gradients are both $4.25$, and their mean is $4.25$, not the sum $8.5$. One update moves from $2.25$ to $0.125$. At step $4$, sampled centers $2$ and $0$ give gradients $-1.875$ and $0.125$; their mean is $-0.875$, so subtracting the scaled negative gradient increases the parameter to $0.5625$.

Final status is STEP_LIMIT, with $4$ updates, $8$ gradient-example evaluations, $20$ monitoring loss-example evaluations, $8$ raw random draws, no retries, no zero-gradient batches, and final random state $373896330$. The exact final full loss is $2.595703125$, above the known minimum $2.5$.

The batch-end generator states and cumulative draw counts are:

UpdateBatch-end stateCumulative raw draws
$1$$226391675$$2$
$2$$1269048164$$4$
$3$$1000656585$$6$
$4$$373896330$$8$

For comparison, the batch-size-$1$ run uses the same seed and four updates, producing:

UpdateIndexOld wSample gradientNew wFull objective
$1$$2$$5$$3$$3.5$$5.625$
$2$$2$$3.5$$1.5$$2.75$$4.03125$
$3$$0$$2.75$$4.75$$0.375$$2.6953125$
$4$$3$$0.375$$-3.625$$2.1875$$3.205078125$

The fourth SGD step increases the full objective. This is a valid sampled update, not an implementation error. The comparison uses equal update counts but unequal gradient work: batch size $2$ spends eight gradient examples, while batch size $1$ spends four. It is an illustration of the mechanisms, not evidence that one wins under equal work.

Budget $B=3$ allows one complete size-$2$ batch, leaving one gradient evaluation unused; no second batch is sampled. With zero gradient budget, the parameter is still the valid supplied start and full monitoring still evaluates its four losses. Turning monitoring off spends no full-loss evaluations and preserves every training update.

Full mini-batch trace and separate training and monitoring accounting The four mini-batches are two two, zero three, zero zero, and two one. The final loss is about two point five nine five seven using eight gradient examples. Four single-example SGD updates use only four gradient examples and finish near three point two zero five one; each monitored run separately evaluates twenty example losses. All mini-batch updates, with training and monitoring costs Step 0 w=5 batch start F=10.500000 g=- Step 1 w=3.5 batch 2,2 F=5.625000 g=3 Step 2 w=2.25 batch 0,3 F=3.281250 g=2.5 Step 3 w=0.125 batch 0,0 F=2.882812 g=4.25 Step 4 w=0.5625 batch 2,1 F=2.595703 g=-0.875 Same four updates, different gradient work: this is a trace comparison, not a fair-work race. MINI-BATCH SIZE 2 final F = 2.59570312 8 gradient examples SGD: BATCH SIZE 1 final F = 3.20507812 4 gradient examples Each monitored run also evaluates 20 example losses: 5 passes over 4 centers. Disabling full monitoring preserves every draw, batch, gradient, and parameter update.
Worked Example

Trace It by Hand

At $w=2.25$, both sampled indices are $0$, whose center is $-2$. With batch size $2$ and rate $0.5$, predict the two component gradients, mean gradient, new parameter, and full objective.

Reveal: each component gradient is $2.25-(-2)=4.25$. The mean is $(4.25+4.25)/2=4.25$, so $w_{\mathrm{next}}=2.25-0.5(4.25)=0.125$. The full objective is $0.5(0.125-1)^2+2.5=2.8828125$. Using the unnormalized sum would instead jump to $-2$, whose full objective is $7$.

Implementation in Python, C++, and Java

Each language panel is complete and independent. Use Python $3.10+$, C++17, or Java $17+$. Compile C++ with g++ -std=c++17 -O2 -ffp-contract=off main.cpp -o sgd and avoid fast-math transformations when comparing binary64 decisions. Save the Java panel as Main.java.

The solver accepts $1\le d\le16$ coordinates, $1\le n\le64$ center vectors, a seed in $1,\ldots,2147483646$, $1\le b\le64$, finite $0\lt\eta\le1$, $0\le K\le500$, and integer $0\le B\le10^6$. All coordinates must be finite and have magnitude at most $10^6$. Batch size may exceed dataset size because sampling is with replacement. The range for $\eta$ suits this common-Hessian quadratic; it is not a generic safe-learning-rate rule for arbitrary training losses.

Python rejects booleans as numeric data, requires actual integers for seed/batch/update/budget controls, and requires an actual Boolean monitoring flag. Inputs must be ordinary finite sequences; Java arrays and rows must be non-null. Invalid inputs raise an exception before training. The optimizer copies inputs and snapshots, preserving recorded points and gradients.

Result holds the final point, monitor flag, final monitored full loss, status, update and example-work counts, zero-gradient and unchanged-update counts, generator state/draw/retry counts, and full history. When monitoring is disabled, the numeric full-loss fields are dummy zeros, and output prints OFF; they must not be interpreted as a zero objective. Step $0$ has empty indices and gradient because no batch was sampled yet. Later history rows store the gradient computed at the preceding row's point and the new point reached by that update.

Training and monitoring budgets are intentionally separate. $B$ limits sampled component-gradient evaluations only. Monitoring at initialization and after every update, when enabled, still costs full-dataset loss evaluations, including with $B=0$. The program reports both costs explicitly. STEP_LIMIT takes precedence when both limits would stop; otherwise it stops before sampling a batch that cannot fit. It performs no partial batch and never silently drops repeated samples.

Arithmetic order is fixed: draw all batch indices, accumulate component gradients in sampled order at the unchanged old point, divide by actual batch size, multiply by rate, and subtract. Monitoring iterates data rows then coordinates, subtracts, squares, sums, multiplies by $1/2$, and divides by $n$. Java's BigDecimal helper only provides eight-place half-even display rounding; optimization uses double. The tested cross-language histories match at binary64 precision, while theoretical unbiasedness and variance claims refer to ideal sampling and exact arithmetic.

The common-Hessian loss is a transparent verification problem with an analytic answer, not a general model-training library. For other losses, provide correct per-example gradients, review their smoothness and numeric ranges, and define policies for non-finite values, constraints, and failures. The monitor flag is included so production-style training need not pay for a full objective after every update.

from dataclasses import dataclass, field
import math


MOD = 2147483647
RANGE = MOD - 1


class RandomSource:
    def __init__(self, seed):
        self.state, self.draws, self.rejections = seed, 0, 0

    def below(self, bound):
        limit = RANGE - RANGE % bound
        while True:
            self.state = 48271 * self.state % MOD
            self.draws += 1
            value = self.state - 1
            if value < limit:
                return value % bound
            self.rejections += 1


def full_objective(point, centers):
    total = 0.0
    for row in centers:
        for x, c in zip(point, row):
            delta = x - c
            square = delta * delta
            total = total + square
    return (0.5 * total) / len(centers)


@dataclass
class State:
    iteration: int
    point: list
    indices: list
    gradient: list
    full_loss: float
    gradient_examples: int
    loss_examples: int
    rng_state: int
    rng_draws: int


@dataclass
class Result:
    status: str = "STEP_LIMIT"
    point: list = field(default_factory=list)
    monitor_full: bool = True
    full_loss: float = 0.0
    updates: int = 0
    gradient_examples: int = 0
    loss_examples: int = 0
    zero_gradient_batches: int = 0
    unchanged_updates: int = 0
    rng_state: int = 0
    rng_draws: int = 0
    rng_rejections: int = 0
    history: list = field(default_factory=list)


def mini_batch_sgd(initial, centers, seed, batch_size, learning_rate,
                   max_steps, gradient_budget, monitor_full=True):
    if any(type(x) is not int for x in [seed, batch_size, max_steps, gradient_budget]):
        raise ValueError("integer controls required")
    if not (1 <= seed < MOD and 1 <= batch_size <= 64 and
            0 <= max_steps <= 500 and 0 <= gradient_budget <= 1000000):
        raise ValueError("control out of range")
    if type(monitor_full) is not bool:
        raise ValueError("monitor flag must be boolean")

    def number(x):
        return type(x) is int or (type(x) is float and math.isfinite(x))

    if not number(learning_rate) or not 0 < learning_rate <= 1:
        raise ValueError("invalid learning rate")
    dimension, size = len(initial), len(centers)
    if not 1 <= dimension <= 16 or not 1 <= size <= 64:
        raise ValueError("invalid dimension or dataset size")
    if any(len(row) != dimension for row in centers):
        raise ValueError("row size mismatch")
    if any(not number(x) or abs(x) > 1000000 for row in [initial] + list(centers) for x in row):
        raise ValueError("invalid coordinate")
    point = list(map(float, initial))
    data = [list(map(float, row)) for row in centers]
    rate = float(learning_rate)
    rng = RandomSource(seed)
    result = Result(monitor_full=monitor_full, rng_state=seed)
    loss = full_objective(point, data) if monitor_full else 0.0
    result.loss_examples = size if monitor_full else 0

    def record(indices, gradient):
        result.history.append(State(result.updates, point[:], indices[:], gradient[:],
            loss, result.gradient_examples, result.loss_examples, rng.state, rng.draws))

    record([], [])
    while True:
        if result.updates == max_steps:
            result.status = "STEP_LIMIT"
            break
        if gradient_budget - result.gradient_examples < batch_size:
            result.status = "GRAD_BUDGET"
            break
        indices = [rng.below(size) for _ in range(batch_size)]
        gradient = [0.0] * dimension
        for index in indices:
            for j in range(dimension):
                contribution = point[j] - data[index][j]
                gradient[j] = gradient[j] + contribution
        for j in range(dimension):
            gradient[j] = gradient[j] / batch_size
        candidate = []
        for x, g in zip(point, gradient):
            displacement = rate * g
            candidate.append(x - displacement)
        result.zero_gradient_batches += int(all(g == 0.0 for g in gradient))
        result.unchanged_updates += int(candidate == point)
        point = candidate
        result.updates += 1
        result.gradient_examples += batch_size
        if monitor_full:
            loss = full_objective(point, data)
            result.loss_examples += size
        record(indices, gradient)
    result.point, result.full_loss = point[:], loss
    result.rng_state, result.rng_draws = rng.state, rng.draws
    result.rng_rejections = rng.rejections
    return result


def vector(xs):
    return "(" + ",".join(f"{x:.8f}" for x in xs) + ")"


def show(label, result):
    print(label)
    for row in result.history:
        indices = ",".join(map(str, row.indices)) or "-"
        gradient = vector(row.gradient) if row.gradient else "-"
        loss = f"{row.full_loss:.8f}" if result.monitor_full else "OFF"
        print(f"k={row.iteration} x={vector(row.point)} batch={indices} grad={gradient} "
              f"full={loss} grad_examples={row.gradient_examples} loss_examples={row.loss_examples}")
    print(f"status={result.status} updates={result.updates} grad_examples={result.gradient_examples} "
          f"loss_examples={result.loss_examples} zero_batches={result.zero_gradient_batches} "
          f"unchanged={result.unchanged_updates} draws={result.rng_draws} "
          f"rejections={result.rng_rejections} rng={result.rng_state}")


def main():
    data = [[-2], [0], [2], [4]]
    show("mini-batch", mini_batch_sgd([5], data, 13, 2, .5, 4, 100))
    show("sgd", mini_batch_sgd([5], data, 13, 1, .5, 4, 100))
    show("budget", mini_batch_sgd([5], data, 13, 2, .5, 4, 3))
    show("monitor-off", mini_batch_sgd([5], data, 13, 2, .5, 2, 100, False))
    show("zero-budget", mini_batch_sgd([5], data, 13, 2, .5, 4, 0))


if __name__ == "__main__":
    main()
#include <cmath>
#include <cstdint>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>
using namespace std;
using I64 = int64_t;
using Rows = vector<vector<double>>;
constexpr I64 MOD = 2147483647, RANGE = MOD - 1;
struct RandomSource {
    I64 state; int draws = 0, rejections = 0;
    explicit RandomSource(I64 seed) : state(seed) {}
    int below(int bound) {
        I64 limit = RANGE - RANGE % bound;
        while (true) {
            state = 48271 * state % MOD; ++draws; I64 value = state - 1;
            if (value < limit) return int(value % bound);
            ++rejections;
        }
    }
};
double full_objective(const vector<double>& point, const Rows& centers) {
    double total = 0.0;
    for (const auto& row : centers) for (size_t j = 0; j < point.size(); ++j) {
        double delta = point[j] - row[j]; double square = delta * delta;
        total = total + square;
    }
    return (0.5 * total) / double(centers.size());
}
struct State {
    int iteration; vector<double> point; vector<int> indices; vector<double> gradient;
    double full_loss; int gradient_examples, loss_examples; I64 rng_state; int rng_draws;
};
struct Result {
    string status = "STEP_LIMIT"; vector<double> point; bool monitor_full = true;
    double full_loss = 0.0; int updates = 0, gradient_examples = 0, loss_examples = 0;
    int zero_gradient_batches = 0, unchanged_updates = 0; I64 rng_state = 0;
    int rng_draws = 0, rng_rejections = 0; vector<State> history;
};
Result mini_batch_sgd(const vector<double>& initial, const Rows& centers, I64 seed,
        int batch_size, double learning_rate, int max_steps, int gradient_budget, bool monitor_full = true) {
    if (seed < 1 || seed >= MOD || batch_size < 1 || batch_size > 64 || max_steps < 0 || max_steps > 500 ||
        gradient_budget < 0 || gradient_budget > 1000000) throw invalid_argument("control out of range");
    if (!isfinite(learning_rate) || learning_rate <= 0 || learning_rate > 1)
        throw invalid_argument("invalid learning rate");
    int dimension = int(initial.size()), size = int(centers.size());
    if (dimension < 1 || dimension > 16 || size < 1 || size > 64)
        throw invalid_argument("invalid dimension or dataset size");
    for (double x : initial) if (!isfinite(x) || abs(x) > 1000000) throw invalid_argument("invalid coordinate");
    for (const auto& row : centers) {
        if (row.size() != initial.size()) throw invalid_argument("row size mismatch");
        for (double x : row) if (!isfinite(x) || abs(x) > 1000000) throw invalid_argument("invalid coordinate");
    }
    auto point = initial; auto data = centers; RandomSource rng(seed); Result result;
    result.monitor_full = monitor_full; result.rng_state = seed;
    double loss = monitor_full ? full_objective(point, data) : 0.0;
    result.loss_examples = monitor_full ? size : 0;
    auto record = [&](const vector<int>& indices, const vector<double>& gradient) {
        result.history.push_back(State{result.updates, point, indices, gradient, loss,
            result.gradient_examples, result.loss_examples, rng.state, rng.draws});
    };
    record(vector<int>{}, vector<double>{});
    while (true) {
        if (result.updates == max_steps) { result.status = "STEP_LIMIT"; break; }
        if (gradient_budget - result.gradient_examples < batch_size) { result.status = "GRAD_BUDGET"; break; }
        vector<int> indices; for (int k = 0; k < batch_size; ++k) indices.push_back(rng.below(size));
        vector<double> gradient(dimension, 0.0);
        for (int index : indices) for (int j = 0; j < dimension; ++j) {
            double contribution = point[j] - data[index][j]; gradient[j] = gradient[j] + contribution;
        }
        for (double& g : gradient) g = g / batch_size;
        vector<double> candidate; bool zero = true;
        for (int j = 0; j < dimension; ++j) {
            double displacement = learning_rate * gradient[j]; candidate.push_back(point[j] - displacement);
            if (gradient[j] != 0.0) zero = false;
        }
        result.zero_gradient_batches += zero; result.unchanged_updates += candidate == point;
        point = candidate; ++result.updates; result.gradient_examples += batch_size;
        if (monitor_full) { loss = full_objective(point, data); result.loss_examples += size; }
        record(indices, gradient);
    }
    result.point = point; result.full_loss = loss; result.rng_state = rng.state;
    result.rng_draws = rng.draws; result.rng_rejections = rng.rejections; return result;
}
string vector_text(const vector<double>& xs) {
    ostringstream out; out << fixed << setprecision(8) << '(';
    for (size_t j = 0; j < xs.size(); ++j) { if (j) out << ','; out << xs[j]; }
    out << ')'; return out.str();
}
void show(const string& label, const Result& result) {
    cout << fixed << setprecision(8) << label << '\n';
    for (const auto& row : result.history) {
        cout << "k=" << row.iteration << " x=" << vector_text(row.point) << " batch=";
        if (row.indices.empty()) cout << '-';
        for (size_t k = 0; k < row.indices.size(); ++k) { if (k) cout << ','; cout << row.indices[k]; }
        cout << " grad=" << (row.gradient.empty() ? "-" : vector_text(row.gradient)) << " full=";
        if (result.monitor_full) cout << row.full_loss; else cout << "OFF";
        cout << " grad_examples=" << row.gradient_examples << " loss_examples=" << row.loss_examples << '\n';
    }
    cout << "status=" << result.status << " updates=" << result.updates << " grad_examples=" << result.gradient_examples
         << " loss_examples=" << result.loss_examples << " zero_batches=" << result.zero_gradient_batches
         << " unchanged=" << result.unchanged_updates << " draws=" << result.rng_draws
         << " rejections=" << result.rng_rejections << " rng=" << result.rng_state << '\n';
}
int main() {
    Rows data; for (double c : vector<double>{-2, 0, 2, 4}) data.push_back(vector<double>{c});
    vector<double> x{5};
    show("mini-batch", mini_batch_sgd(x, data, 13, 2, .5, 4, 100));
    show("sgd", mini_batch_sgd(x, data, 13, 1, .5, 4, 100));
    show("budget", mini_batch_sgd(x, data, 13, 2, .5, 4, 3));
    show("monitor-off", mini_batch_sgd(x, data, 13, 2, .5, 2, 100, false));
    show("zero-budget", mini_batch_sgd(x, data, 13, 2, .5, 4, 0));
    return 0;
}
import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.ArrayList;
import java.util.List;

public class Main {
    static final long MOD=2147483647L,RANGE=MOD-1;
    static class RandomSource {
        long state;int draws,rejections;
        RandomSource(long seed) { state=seed; }
        int below(int bound) {
            long limit=RANGE-RANGE%bound;
            while (true) {
                state=48271*state%MOD;++draws;long value=state-1;
                if (value<limit) return (int)(value%bound);++rejections;
            }
        }
    }
    static double[][] copy(double[][] rows) {
        double[][] out=new double[rows.length][];
        for (int i=0;i<rows.length;++i) out[i]=rows[i].clone();return out;
    }
    static double full_objective(double[] point,double[][] centers) {
        double total=0.0;
        for (double[] row : centers) for (int j=0;j<point.length;++j) {
            double delta=point[j]-row[j];double square=delta*delta;total=total+square;
        }
        return (0.5*total)/centers.length;
    }
    static class State {
        int iteration;double[] point;int[] indices;double[] gradient;double full_loss;
        int gradient_examples,loss_examples;long rng_state;int rng_draws;
        State(int k,double[] x,int[] b,double[] g,double loss,int ge,int le,long rs,int rd) {
            iteration=k;point=x.clone();indices=b.clone();gradient=g.clone();full_loss=loss;
            gradient_examples=ge;loss_examples=le;rng_state=rs;rng_draws=rd;
        }
    }
    static class Result {
        String status="STEP_LIMIT";double[] point=new double[0];boolean monitor_full=true;double full_loss;
        int updates,gradient_examples,loss_examples,zero_gradient_batches,unchanged_updates;
        long rng_state;int rng_draws,rng_rejections;List<State> history=new ArrayList<>();
    }
    static Result mini_batch_sgd(double[] initial,double[][] centers,long seed,int batch_size,
            double learning_rate,int max_steps,int gradient_budget,boolean monitor_full) {
        if (seed<1 || seed>=MOD || batch_size<1 || batch_size>64 || max_steps<0 || max_steps>500 ||
            gradient_budget<0 || gradient_budget>1000000) throw new IllegalArgumentException("control out of range");
        if (!Double.isFinite(learning_rate) || learning_rate<=0 || learning_rate>1)
            throw new IllegalArgumentException("invalid learning rate");
        int dimension=initial.length,size=centers.length;
        if (dimension<1 || dimension>16 || size<1 || size>64)
            throw new IllegalArgumentException("invalid dimension or dataset size");
        for (double x : initial)
            if (!Double.isFinite(x) || Math.abs(x)>1000000) throw new IllegalArgumentException("invalid coordinate");
        for (double[] row : centers) {
            if (row.length!=dimension) throw new IllegalArgumentException("row size mismatch");
            for (double x : row)
                if (!Double.isFinite(x) || Math.abs(x)>1000000) throw new IllegalArgumentException("invalid coordinate");
        }
        double[] point=initial.clone();double[][] data=copy(centers);RandomSource rng=new RandomSource(seed);
        Result result=new Result();result.monitor_full=monitor_full;result.rng_state=seed;
        double loss=monitor_full ? full_objective(point,data) : 0.0;result.loss_examples=monitor_full ? size : 0;
        result.history.add(new State(0,point,new int[0],new double[0],loss,0,result.loss_examples,rng.state,0));
        while (true) {
            if (result.updates==max_steps) { result.status="STEP_LIMIT";break; }
            if (gradient_budget-result.gradient_examples<batch_size) { result.status="GRAD_BUDGET";break; }
            int[] indices=new int[batch_size];for (int k=0;k<batch_size;++k) indices[k]=rng.below(size);
            double[] gradient=new double[dimension];
            for (int index : indices) for (int j=0;j<dimension;++j) {
                double contribution=point[j]-data[index][j];gradient[j]=gradient[j]+contribution;
            }
            for (int j=0;j<dimension;++j) gradient[j]=gradient[j]/batch_size;
            double[] candidate=new double[dimension];boolean zero=true,changed=false;
            for (int j=0;j<dimension;++j) {
                double displacement=learning_rate*gradient[j];candidate[j]=point[j]-displacement;
                if (gradient[j]!=0.0) zero=false;if (candidate[j]!=point[j]) changed=true;
            }
            if (zero) ++result.zero_gradient_batches;if (!changed) ++result.unchanged_updates;
            point=candidate;++result.updates;result.gradient_examples+=batch_size;
            if (monitor_full) { loss=full_objective(point,data);result.loss_examples+=size; }
            result.history.add(new State(result.updates,point,indices,gradient,loss,
                result.gradient_examples,result.loss_examples,rng.state,rng.draws));
        }
        result.point=point.clone();result.full_loss=loss;result.rng_state=rng.state;
        result.rng_draws=rng.draws;result.rng_rejections=rng.rejections;return result;
    }
    static String fixed8(double x) {
        String s=new BigDecimal(x).setScale(8,RoundingMode.HALF_EVEN).toPlainString();
        return Double.doubleToRawLongBits(x)<0 && s.charAt(0)!='-' ? "-"+s : s;
    }
    static String vector(double[] xs) {
        StringBuilder s=new StringBuilder("(");
        for (int j=0;j<xs.length;++j) { if (j>0) s.append(',');s.append(fixed8(xs[j])); }
        return s.append(')').toString();
    }
    static void show(String label,Result result) {
        System.out.println(label);
        for (State row : result.history) {
            StringBuilder batch=new StringBuilder();
            for (int k=0;k<row.indices.length;++k) { if (k>0) batch.append(',');batch.append(row.indices[k]); }
            System.out.println("k="+row.iteration+" x="+vector(row.point)+" batch="+(batch.length()==0?"-":batch)+
                " grad="+(row.gradient.length==0?"-":vector(row.gradient))+" full="+(result.monitor_full?fixed8(row.full_loss):"OFF")+
                " grad_examples="+row.gradient_examples+" loss_examples="+row.loss_examples);
        }
        System.out.println("status="+result.status+" updates="+result.updates+" grad_examples="+result.gradient_examples+
            " loss_examples="+result.loss_examples+" zero_batches="+result.zero_gradient_batches+
            " unchanged="+result.unchanged_updates+" draws="+result.rng_draws+
            " rejections="+result.rng_rejections+" rng="+result.rng_state);
    }
    public static void main(String[] args) {
        double[][] data=new double[4][];data[0]=new double[]{-2};data[1]=new double[]{0};
        data[2]=new double[]{2};data[3]=new double[]{4};double[] x={5};
        show("mini-batch",mini_batch_sgd(x,data,13,2,.5,4,100,true));
        show("sgd",mini_batch_sgd(x,data,13,1,.5,4,100,true));
        show("budget",mini_batch_sgd(x,data,13,2,.5,4,3,true));
        show("monitor-off",mini_batch_sgd(x,data,13,2,.5,2,100,false));
        show("zero-budget",mini_batch_sgd(x,data,13,2,.5,4,0,true));
    }
}

Expected standard output, identical in all three languages:

  • mini-batch
  • k=0 x=(5.00000000) batch=- grad=- full=10.50000000 grad_examples=0 loss_examples=4
  • k=1 x=(3.50000000) batch=2,2 grad=(3.00000000) full=5.62500000 grad_examples=2 loss_examples=8
  • k=2 x=(2.25000000) batch=0,3 grad=(2.50000000) full=3.28125000 grad_examples=4 loss_examples=12
  • k=3 x=(0.12500000) batch=0,0 grad=(4.25000000) full=2.88281250 grad_examples=6 loss_examples=16
  • k=4 x=(0.56250000) batch=2,1 grad=(-0.87500000) full=2.59570312 grad_examples=8 loss_examples=20
  • status=STEP_LIMIT updates=4 grad_examples=8 loss_examples=20 zero_batches=0 unchanged=0 draws=8 rejections=0 rng=373896330
  • sgd
  • k=0 x=(5.00000000) batch=- grad=- full=10.50000000 grad_examples=0 loss_examples=4
  • k=1 x=(3.50000000) batch=2 grad=(3.00000000) full=5.62500000 grad_examples=1 loss_examples=8
  • k=2 x=(2.75000000) batch=2 grad=(1.50000000) full=4.03125000 grad_examples=2 loss_examples=12
  • k=3 x=(0.37500000) batch=0 grad=(4.75000000) full=2.69531250 grad_examples=3 loss_examples=16
  • k=4 x=(2.18750000) batch=3 grad=(-3.62500000) full=3.20507812 grad_examples=4 loss_examples=20
  • status=STEP_LIMIT updates=4 grad_examples=4 loss_examples=20 zero_batches=0 unchanged=0 draws=4 rejections=0 rng=1269048164
  • budget
  • k=0 x=(5.00000000) batch=- grad=- full=10.50000000 grad_examples=0 loss_examples=4
  • k=1 x=(3.50000000) batch=2,2 grad=(3.00000000) full=5.62500000 grad_examples=2 loss_examples=8
  • status=GRAD_BUDGET updates=1 grad_examples=2 loss_examples=8 zero_batches=0 unchanged=0 draws=2 rejections=0 rng=226391675
  • monitor-off
  • k=0 x=(5.00000000) batch=- grad=- full=OFF grad_examples=0 loss_examples=0
  • k=1 x=(3.50000000) batch=2,2 grad=(3.00000000) full=OFF grad_examples=2 loss_examples=0
  • k=2 x=(2.25000000) batch=0,3 grad=(2.50000000) full=OFF grad_examples=4 loss_examples=0
  • status=STEP_LIMIT updates=2 grad_examples=4 loss_examples=0 zero_batches=0 unchanged=0 draws=4 rejections=0 rng=1269048164
  • zero-budget
  • k=0 x=(5.00000000) batch=- grad=- full=10.50000000 grad_examples=0 loss_examples=4
  • status=GRAD_BUDGET updates=0 grad_examples=0 loss_examples=4 zero_batches=0 unchanged=0 draws=0 rejections=0 rng=13

Complexity, Correctness, and Invariants

Let $T$ be completed updates. This program has no early convergence exit, so

$$T=\min\left(K,\left\lfloor\frac{B}{b}\right\rfloor\right),\qquad G=bT.$$

Here $G$ counts component-gradient examples, including repeated indices. Full monitoring separately costs $H=n(T+1)$ component-loss examples; with monitoring off, $H=0$. If $J$ raw draws are rejected, random draw count is $Q=bT+J$. Checks occur before sampling, so an unaffordable batch consumes no random draws and no gradient work.

For this generator and bounds $n\le64$, a rejected state is followed by an accepted one: it is $2147483647-a$ for $1\le a\le63$, whose successor $2147483647-48271a$ lies outside every relevant rejection interval. Thus mapping has bounded constant cost here. Input validation and copying cost $\Theta(nd)$. Training costs $\Theta(Tbd)$, and optional full monitoring adds $\Theta((T+1)nd)$. Including these components, actual runtime is

$$\Theta\bigl(nd+Tbd+\mathbf1_{\mathrm{monitor}}(T+1)nd\bigr).$$

Best, average, and worst objective landscapes do the same amount of work for fixed valid sizes and controls; attained quality can differ. Best work for zero updates still includes input validation and any requested initial full-loss pass. For a general loss, substitute actual per-example gradient and loss costs. Wall-clock speed also depends on vectorization, memory traffic, and hardware utilization, so increasing batch size does not automatically improve quality or time to target.

Live training state needs $O(b+d)$ auxiliary space beyond data, and this implementation copies the $O(nd)$ dataset. It retains history, so total auxiliary storage is $O(nd+(T+1)(b+d))$, including sampled indices, parameter vectors, and mean gradients. Removing or thinning history reduces that extra trace storage. Full monitoring does not use a training gradient or alter randomness.

The implementation invariant is that each committed update uses the arithmetic mean of exactly $b$ sampled component gradients evaluated at the same preceding point. Its counters satisfy $G=bT$, histories contain $T+1$ coherent states, and a monitored score belongs to the point stored in its row. A history gradient belongs to the previous point, so its provenance is explicit. These are consistency guarantees, not monotone-loss or optimality guarantees.

Under independent uniform sampling with replacement and exact arithmetic,

$$\mathbb E[\widehat{\mathbf g}_t\mid\mathbf w_t]=\nabla F(\mathbf w_t).$$

If the conditional single-example squared noise is bounded by $\sigma^2$, averaging $b$ independent samples gives conditional squared noise at most $\sigma^2/b$. This needs the sampling assumptions: random reshuffling within an epoch has a different conditional distribution, and nonuniform sampling generally needs appropriate weighting to target the same empirical objective.

For the teaching family these facts can be proved directly. Let $\sigma_c^2=n^{-1}\sum_i\lVert\mathbf c_i-\boldsymbol\mu\rVert_2^2$. A sampled mean center $\overline{\mathbf c}_t$ satisfies $\widehat{\mathbf g}_t=\mathbf w_t-\overline{\mathbf c}_t$, so error $\mathbf e_t=\mathbf w_t-\boldsymbol\mu$ obeys

$$\mathbf e_{t+1}=(1-\eta)\mathbf e_t+\eta(\overline{\mathbf c}_t-\boldsymbol\mu).$$

The sampled-center noise has mean zero and squared norm expectation $\sigma_c^2/b$, giving the exact conditional recurrence

$$\mathbb E[\lVert\mathbf e_{t+1}\rVert_2^2\mid\mathbf w_t]=(1-\eta)^2\lVert\mathbf e_t\rVert_2^2+\frac{\eta^2\sigma_c^2}{b}.$$

For constant $0\lt\eta\le1$ and nonzero variance, the limiting expected squared parameter error is $\eta\sigma_c^2/[b(2-\eta)]$. Expected objective gap is half that value. For the worked data, $\sigma_c^2=5$, $\eta=1/2$, and $b=2$, the ideal limiting expected objective gap is $5/12$. A lucky short run can have a smaller gap; the limit describes an expectation over random runs, not a lower bound on every iterate. Fixed-rate stochastic iterates generally fluctuate rather than converge exactly to the minimizer.

This also explains why $b=n$ with replacement retains variance $\sigma_c^2/n$, whereas averaging every example once at a fixed point has no sampling variance. Decreasing learning rates, increasing batch sizes, or iterate averaging can improve asymptotic behavior under appropriate assumptions, but none is implemented silently here. The finite deterministic generator and floating-point arithmetic support reproducibility, not a proof that the ideal probability model holds for its sequence.

Method / policyTypical gradient work per updateSampling and diagnostic implications
Full gradient descent$n$ component gradientsEvery example once at one point; no sampling noise
SGDOne component gradientFresh with-replacement draw in this lesson; noisy individual updates
Mini-batch SGD$b$ component gradientsMean of a sampled batch; repeats allowed; ideal variance shrinks by $b$
Random reshufflingUsually $b$ gradients per batch, one shuffled pass per epochNo repeats within the pass; conditional sampling assumptions differ
Gradient accumulationSum of several micro-batches before one updateMatches one large mean gradient only with correct scaling, frozen parameters, and compatible model behavior
Momentum / adaptive optimizerBatch-gradient cost plus parameter-state workAdditional update state; not identical to plain SGD
This full monitor$n$ component-loss evaluations, separate from gradientsOptional at initialization and every update; can dominate small-batch training

Variants, Pitfalls, and Optimizations

Mean versus sum matters. Dividing by batch size keeps the update scale consistent with a mean empirical objective. Using a sum silently multiplies the effective rate by $b$. A final smaller batch in an epoch-based implementation must use its actual size. This reference implementation instead requires a complete fixed-size batch and stops before an unaffordable one; it does not divide a short batch by an unrelated nominal size.

Freeze parameters throughout a batch. Accumulating several gradients and then updating once is different from taking an SGD step after each example. Gradient accumulation across micro-batches can reproduce a larger mean gradient when parameters remain unchanged and normalization is correct. Stateful operations such as batch normalization, independently sampled dropout, or updates to other model state can prevent exact equivalence even when the algebraic scaling looks right.

Do not stop on one zero sampled gradient. At $w=2$, the component centered at $2$ has zero gradient while the full objective's gradient is $1$. With the two-center dataset $[2,0]$, seed $1$ first samples center $2$ and then center $0$: the first update is unchanged, but the second reaches the true mean $1$ at rate $0.5$. A small or unchanged stochastic step can also reflect rounding, a tiny rate, or sampling noise. Use a defined diagnostic and stopping policy rather than calling that one batch converged.

Define sampling, epochs, and unbiasedness together. Sampling with replacement can repeat an example even inside one batch. Shuffling a finite dataset once per epoch is a different schedule; the next draw conditioned on past draws is no longer uniform over all examples. Importance sampling and class balancing change probabilities and may require weights to preserve the desired objective. Duplicate sampling is not inherently a bug, but an unintended sampling rule is.

Do not require monotone full loss. An unbiased direction is an expectation statement, not a guarantee about a particular draw. Constant step sizes can maintain a noise neighborhood even on strongly convex quadratics. Decay schedules or averaging require their own choices and assumptions. Comparing only a smoothed training curve can hide transient increases, and comparing only batch losses can conflate changing data with parameter progress.

Count the expensive work. Computing the full training loss after every small-batch update can erase much of its computational advantage. Use periodic full diagnostics or other stated metrics in larger systems and report their costs. Compare methods at equal sampled-example budgets, wall-clock budgets, or a clearly chosen target criterion; equal update counts are not equal work when batch sizes differ. Larger batches can improve hardware utilization without implying a universal linear learning-rate scaling rule.

Keep objectives and regularization consistent. Sum versus mean reduction, example weights, and penalty scaling affect gradients. A regularizer included identically in every component remains once after averaging; adding it incorrectly per example outside that averaging can overcount it. Check model gradients with suitable finite differences or an independent reference on small cases, and separate training optimization from validation performance and data leakage.

Use the right comparison. Momentum, Adam, variance-reduced methods, distributed sampling, and data-parallel gradient averaging add policies beyond plain SGD. Synchronization, averaging weights for unequal batch sizes, and random streams affect reproducibility. Cornell's primary mini-batch analysis and SGD notes explain how variance and learning-rate assumptions enter convergence reasoning; match those assumptions before carrying a theorem into an implementation.

A zero sampled gradient and the danger of stopping on one mini-batch At parameter two, the component centered at two has gradient zero but the full gradient is one. Continuing with a component centered at zero produces gradient two and rate one half reaches the true minimizer one. Stopping at the first zero estimate would incorrectly label a suboptimal point converged. One zero sampled gradient does not certify convergence Original four centers: at w = 2, sample gradient 0 can hide full gradient 1. CONTINUE THE STOCHASTIC RUN STOP ON THIS ONE BATCH First update leaves w = 2. Next sampled center is 0. Gradient 2 → new w = 1. Full F reaches its minimum 2.5. Returns w = 2 as “converged”. Full F is still 3. The optimum has F = 2.5. A zero estimate hid nonzero slope. Use a stated stopping policy based on appropriate diagnostics, not one noisy batch. The reference programs count zero-gradient batches and keep going until a budget limit.

Average using the actual batch size, evaluate all batch gradients at the same old parameters, and state whether sampling uses replacement or reshuffling. Separate monitored full loss from sampled training work, and never infer convergence from one zero or unchanged stochastic update; noise can increase full loss even when the implementation is correct.

Practice Problems and Real-World Use

Start with the center-fitting family, whose minimizer is the center mean. Exhaustively enumerate all possible small batches to verify expected gradients and variances, then compare repeated seeds at matched gradient-example budgets. Turn monitoring on and off to verify that diagnostic work does not alter the training path.

For a realistic learning task, derive the per-example loss and gradient before choosing the sampling and batch policy. Record reduction, regularization, initialization, seed streams, update schedule, partial-batch handling, monitoring cadence, training work, and validation protocol. Linear and logistic models make useful next steps before introducing model state, autodifferentiation, distributed execution, or adaptive optimizers.

Problem or applicationExerciseMain idea testedDifficulty
Mean-center fittingReproduce both traces and compare with the analytic meanBatch averaging, duplicates, and noiseEasy
Zero-gradient trapSample a component whose gradient vanishes away from the full optimumStopping-policy correctnessEasy
Exact small-batch enumerationEnumerate with-replacement tuples and compute momentsUnbiasedness and variance reductionMedium
Mini-batch linear regressionDerive residual gradients and compare with a full-gradient referenceMatrix batches, reduction, and gradient workMedium
Regularized logistic regressionCheck stable per-example loss and penalty scalingCorrect objectives and validationHard
Reshuffling laboratoryCompare sampling schedules under equal gradient workEpoch coverage and conditional distributionsMedium
Micro-batch accumulationMatch a larger frozen-parameter batch on a simple modelNormalization and update timingHard
Monitoring and parallel trainingSeparate diagnostics and define worker averagingCost accounting and reproducibilityAdvanced