Variance-Weighted Batch Distribution

• Variance-weighted batch distribution is a technique that reweights mini-batch contributions based on variance in gradient space to produce unbiased and lower-variance estimators.
• It partitions data into clusters where each representative is weighted by cluster size, penalizing large, heterogeneous clusters to ensure tighter variance control.
• The approach proves most effective in settings with high gradient redundancy, enhancing optimization in duplicate-rich or clearly clusterable datasets.

Variance-weighted batch distribution denotes a class of batch-construction or batch-reweighting schemes in which the contribution of a sampled set is shaped by variance-related structure rather than by uniform empirical averaging. In one explicit stochastic-gradient formulation, the dataset is partitioned in gradient space, one example is sampled from each partition cell, and the selected examples are weighted by cluster size, yielding an unbiased mini-batch gradient estimator whose variance is minimized by a weighted clustering objective (Faghri et al., 2020). Across adjacent literatures, the same general idea appears in nonuniform sampling over batches, inverse-variance weighting of losses within a minibatch, variance-sensitive batch-size schedules, and explicit variance-tilted joint laws over a whole generated batch (Needell et al., 2016, Mai et al., 2021, Azangulov et al., 20 Jun 2026).

1. Canonical formulation in stochastic gradient descent

The most direct formalization begins with empirical risk minimization

$L(\theta)=\frac1N\sum_{i=1}^N \ell(x_i;\theta),$

with per-example gradients

$$g_i=\nabla_\theta \ell(x_i;\theta)\in\mathbb{R}^d, \qquad g=\nabla_\theta L(\theta)=\frac1N\sum_{i=1}^N g_i.$$

The dataset is partitioned into $K$ subsets in gradient space, with assignment $a_i\in\{1,\dots,K\}$ and cluster sizes

$N_k=\sum_{i=1}^N \mathbf{1}(a_i=k).$

If cluster $k$ contains gradients $\{g_{j,k}\}_{j=1}^{N_k}$, the stratified estimator is

$$\hat g(\theta)=\frac1N\sum_{k=1}^K N_k\, g_{j,k}, \qquad j\sim \mathrm{Unif}\{1,\dots,N_k\}. \tag{1}$$

This estimator is unbiased for any partition,

$\mathbb E[\hat g]=g,$

so the design problem is purely variational: choose the partition so that the variance of the unbiased estimator is as small as possible (Faghri et al., 2020).

Variance is defined as the trace of the covariance matrix of the vector estimator. Under the paper’s i.i.d. assumption, cross-cluster covariance terms vanish, yielding

$$\mathrm{Var}[\hat g] = \frac{1}{N^2}\sum_{k=1}^K N_k^2\,\mathrm{Var}[g_{j,k}]. \tag{2}$$

The dependence on $$g_i=\nabla_\theta \ell(x_i;\theta)\in\mathbb{R}^d, \qquad g=\nabla_\theta L(\theta)=\frac1N\sum_{i=1}^N g_i.$$0 is decisive. Large, heterogeneous clusters are costly, because their within-cluster dispersion is amplified quadratically by cluster size. This gives the precise sense in which the batch distribution is variance-weighted: a selected representative stands in for the entire cluster, and the variance contribution of that cluster is weighted by its represented mass.

For fixed $$g_i=\nabla_\theta \ell(x_i;\theta)\in\mathbb{R}^d, \qquad g=\nabla_\theta L(\theta)=\frac1N\sum_{i=1}^N g_i.$$1, minimizing $$g_i=\nabla_\theta \ell(x_i;\theta)\in\mathbb{R}^d, \qquad g=\nabla_\theta L(\theta)=\frac1N\sum_{i=1}^N g_i.$$2 is equivalent to minimizing

$$g_i=\nabla_\theta \ell(x_i;\theta)\in\mathbb{R}^d, \qquad g=\nabla_\theta L(\theta)=\frac1N\sum_{i=1}^N g_i.$$3

and, with cluster centers

$$g_i=\nabla_\theta \ell(x_i;\theta)\in\mathbb{R}^d, \qquad g=\nabla_\theta L(\theta)=\frac1N\sum_{i=1}^N g_i.$$4

this becomes

$$g_i=\nabla_\theta \ell(x_i;\theta)\in\mathbb{R}^d, \qquad g=\nabla_\theta L(\theta)=\frac1N\sum_{i=1}^N g_i.$$5

The extra factor $$g_i=\nabla_\theta \ell(x_i;\theta)\in\mathbb{R}^d, \qquad g=\nabla_\theta L(\theta)=\frac1N\sum_{i=1}^N g_i.$$6 distinguishes the objective from ordinary $$g_i=\nabla_\theta \ell(x_i;\theta)\in\mathbb{R}^d, \qquad g=\nabla_\theta L(\theta)=\frac1N\sum_{i=1}^N g_i.$$7-means. Without that factor, the criterion would reduce to standard within-cluster sum of squares. With it, variance in large clusters is penalized more heavily, so large clusters are driven to be tight while scattered or unusual gradients are pushed into smaller groups (Faghri et al., 2020).

2. Weighted clustering and Gradient Clustering

The optimization problem above is realized in the method called Gradient Clustering (GC). If all per-example gradients are available, GC alternates between an assignment step and a center update. For fixed centers and fixed current cluster sizes, assignment is

$$g_i=\nabla_\theta \ell(x_i;\theta)\in\mathbb{R}^d, \qquad g=\nabla_\theta L(\theta)=\frac1N\sum_{i=1}^N g_i.$$8

and the center update is

$$g_i=\nabla_\theta \ell(x_i;\theta)\in\mathbb{R}^d, \qquad g=\nabla_\theta L(\theta)=\frac1N\sum_{i=1}^N g_i.$$9

Because $K$0 appears inside the assignment cost, the usual monotonic-decrease guarantees of Lloyd’s algorithm do not directly apply; assignments are therefore solved for fixed cluster sizes and then $K$1 is updated (Faghri et al., 2020).

Exact clustering in full gradient space is prohibitive for deep networks, so the method uses low-rank layerwise approximations. For a fully connected layer with weights $K$2, each per-example gradient has outer-product form

$K$3

with input activations $K$4 and backpropagated output gradients $K$5. Cluster centers are approximated by rank-1 matrices

$K$6

which allows the squared Frobenius distance to be evaluated through low-dimensional inner products: $K$7 Under a K-FAC-like independence assumption, the factorized center updates are

$K$8

Analogous formulas are given for convolutional layers by exploiting the sum-of-outer-products structure across spatial locations (Faghri et al., 2020).

Operationally, the number of clusters is set equal to the mini-batch size, clustering is updated only every few epochs, and center updates can be online or mini-batch incremental. The paper argues that the overall overhead can be made roughly at most $K$9 a normal backprop step, while memory overhead is like storing $a_i\in\{1,\dots,K\}$0 extra samples because only rank-1 cluster-center approximations are stored. The method assumes i.i.d. training examples and uses zero cross-covariance of independently sampled gradients in the variance derivation. A further practical limitation is that GC may fail to help when the gradient distribution is not clusterable in the relevant sense; on ImageNet, the authors report substantial overlap with ordinary SGD, suggesting little exploitable cluster structure under their approximation (Faghri et al., 2020).

3. Gradient variance during training and the role of normalized variance

The same study places the batch-distribution construction in a broader empirical analysis of gradient noise. Contrary to common intuition, raw gradient variance often increases during training on CIFAR-10 and ImageNet, especially after learning-rate drops. On MNIST, by contrast, variance decreases steadily and becomes numerically near zero, a behavior connected in the paper to the Strong Growth Condition. The authors also report that smaller learning rates coincide with higher variance in deep models, whereas in random-features models larger learning rates can produce smaller variance (Faghri et al., 2020).

Because raw variance alone is not comparable across tasks and training phases, the paper defines normalized gradient variance as

$a_i\in\{1,\dots,K\}$1

Its inverse is a signal-to-noise ratio. Empirically, this statistic correlates better with optimization speed than raw variance. On ImageNet, normalized variance quickly exceeds $a_i\in\{1,\dots,K\}$2, meaning noise power exceeds signal power; on MNIST and CIFAR-10 it often stays below $a_i\in\{1,\dots,K\}$3. Since the denominator is shared across estimators on a fixed SGD trajectory, normalized variance preserves relative ordering while improving comparability across problems (Faghri et al., 2020).

These measurements also clarify where variance-weighted batch construction helps. On MNIST, GC’s variance is consistently below uniform SGD with double batch size. On CIFAR-10, GC improves on standard SGD but usually not on the doubled-batch baseline except during fluctuation phases. On ImageNet, GC often overlaps with SGD. The clearest validation comes from duplicate-data experiments: when many duplicates are inserted into random-feature models or image datasets, GC strongly outperforms ordinary random mini-batching, because one weighted representative is sufficient for a large low-variance cluster. This suggests that variance-weighted batch distributions are most effective when gradient redundancy or cluster structure is pronounced (Faghri et al., 2020).

4. Other direct forms of variance-based batch weighting

Variance-related batch design also appears in several other direct forms, but the weighted object differs across settings. In batched SGD with weighted sampling, the stochastic atom is itself a batch $a_i\in\{1,\dots,K\}$4, and one samples $a_i\in\{1,\dots,K\}$5 nonuniformly with importance correction. For smooth objectives, the recommended batch distribution is

$a_i\in\{1,\dots,K\}$6

where $a_i\in\{1,\dots,K\}$7 is the batch Lipschitz constant of $a_i\in\{1,\dots,K\}$8. In least squares this specializes to

$a_i\in\{1,\dots,K\}$9

so batch spectral norms determine the sampling law (Needell et al., 2016).

Deep heteroscedastic regression uses a different construction. Batch Inverse-Variance replaces the uniform minibatch average by a normalized inverse-variance weighted loss

$N_k=\sum_{i=1}^N \mathbf{1}(a_i=k).$0

Sampling is unchanged; the batch empirical distribution is reweighted internally by per-sample label-noise variances $N_k=\sum_{i=1}^N \mathbf{1}(a_i=k).$1. The normalization prevents near-zero-variance labels from causing an uncontrolled increase in the effective learning rate (Mai et al., 2021).

At the opposite extreme, variance weighting can be imposed on an entire joint batch law. In variance-tilted diffusion sampling, the target distribution over a batch $N_k=\sum_{i=1}^N \mathbf{1}(a_i=k).$2 is

$N_k=\sum_{i=1}^N \mathbf{1}(a_i=k).$3

where

$N_k=\sum_{i=1}^N \mathbf{1}(a_i=k).$4

Here the variance weighting is not a sampling probability for examples within a minibatch; it is a multiplicative tilt of an i.i.d. product law toward batches with larger empirical feature spread (Azangulov et al., 20 Jun 2026).

Setting	Weighted object	Representative formulation
Gradient clustering	Cluster representatives	$N_k=\sum_{i=1}^N \mathbf{1}(a_i=k).$5
Batched SGD sampling	Batch probabilities	$N_k=\sum_{i=1}^N \mathbf{1}(a_i=k).$6
Deep heteroscedastic regression	Within-batch losses	$N_k=\sum_{i=1}^N \mathbf{1}(a_i=k).$7
Diffusion diverse sampling	Joint batch law	$N_k=\sum_{i=1}^N \mathbf{1}(a_i=k).$8

The common structure is nonuniform contribution of represented mass, but the mechanism can lie in sample selection, loss weighting, or the definition of the batch target itself.

5. Indirect variance-sensitive batching

A separate line of work adjusts the total batch size rather than the within-batch distribution. In adaptive-batch SGD, the theoretical rule is

$N_k=\sum_{i=1}^N \mathbf{1}(a_i=k).$9

for PL, strongly convex, and convex settings, and

$k$0

for smooth nonconvex problems. The key variance relation is

$k$1

This is variance-sensitive batch sizing, but it is not a variance-weighted distribution over examples or clusters (Sievert et al., 2019).

Variance-reduced finite-sum methods employ a related idea at the epoch level. History-gradient adaptation sets

$k$2

The outer batch is therefore small when historical stochastic gradients are large and large when they are small. Again, the policy reallocates total sampling effort across epochs, not probability mass across examples inside one batch (Ji et al., 2019).

Optimizer-state design provides a third indirect mechanism. Batch size invariant Adam changes the second-moment update from “average first, then square” to “square first, then average”: $k$3 This preserves

$k$4

across merged micro-batches, so the expected $k$5-statistic does not change with batch aggregation. The method is variance-preserving in optimizer statistics, but it does not alter the batch sampling law (Wang et al., 2024).

6. Boundaries of the concept and adjacent methods

The literature makes a clear distinction between variance-weighted batch distributions and methods that merely reduce variance after the batch has already been sampled. Reduced-variance random batch methods for nonlocal PDEs keep the standard uniform random subsampling of particles and add a control-variate correction; the novelty lies in the coefficient

$k$6

not in a nonuniform batch distribution (Pareschi et al., 2023). Variance-reduced random batch Langevin dynamics likewise samples shell neighbors uniformly and then subtracts the estimated force covariance from the thermostat noise,

$k$7

so variance enters as post-sampling correction rather than as a sampling law (Xu et al., 2024).

A broader probabilistic usage appears in batch Bayesian optimization. Optimistic Expected Improvement does not sample data batches, but it evaluates candidate query sets using a lower bound on batch Expected Improvement defined over all distributions matching the same predictive mean and covariance. The relevant sufficient statistic is the second-moment matrix

$k$8

so the acquisition is covariance-aware in the multivariate second-moment sense rather than a variance-weighted batch distribution in stochastic optimization (Rontsis et al., 2017).

Two recurrent misconceptions therefore require separation. First, variance-aware batch sizing is not the same as a variance-weighted batch composition rule. Second, control variates, optimizer-state corrections, or post hoc noise calibration can substantially reduce variance while leaving the batch law uniform. The direct concept is present only when variance information changes represented mass inside the batch, alters the probability with which batch elements or whole batches are selected, or defines the batch target as a joint law whose density is explicitly multiplied by a variance functional.