---
title: Variance-Weighted Batch Distribution
url: https://www.emergentmind.com/topics/variance-weighted-batch-distribution
type: topic
---

# Variance-Weighted Batch Distribution

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 [2007.04532]. 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 [1608.07641][2107.04497][2606.22239].

## 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 [2007.04532].

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 \(N_k^2\) 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 \(K\), minimizing \(\mathrm{Var}[\hat g]\) is equivalent to minimizing
\[
\sum_{k=1}^K N_k^2\,\mathrm{Var}[g_{j,k}],
\]
and, with cluster centers
\[
c_k=\frac1{N_k}\sum_{i:a_i=k} g_i,
\]
this becomes
\[
\min_{\{c_k\},\{a_i\}}
\sum_{k=1}^K
N_k
\sum_{i=1}^N
\|g_i-c_k\|^2\,\mathbf{1}(a_i=k).
\tag{3}
\]
The extra factor \(N_k\) distinguishes the objective from ordinary \(K\)-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 [2007.04532].

## 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
\[
a_i=\arg\min_k N_k\,\|g_i-c_k\|^2,
\tag{4}
\]
and the center update is
\[
c_k=\frac1{N_k}\sum_{i=1}^N g_i\,\mathbf{1}(a_i=k).
\tag{5}
\]
Because \(N_k\) 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 \(N_k\) is updated [2007.04532].

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 \(W\in\mathbb R^{I\times O}\), each per-example gradient has outer-product form
\[
g_i=z_i\delta_i^\top,
\]
with input activations \(z_i\) and backpropagated output gradients \(\delta_i\). Cluster centers are approximated by rank-1 matrices
\[
C_k\approx c_k d_k^\top,
\]
which allows the squared Frobenius distance to be evaluated through low-dimensional inner products:
\[
C_k \odot g_i
=
(c_k d_k^\top)\odot (z_i\delta_i^\top)
=
(c_k\odot z_i)(d_k\odot \delta_i).
\tag{6}
\]
Under a K-FAC-like independence assumption, the factorized center updates are
\[
c_k=\frac1{N_k}\sum_{i=1}^N z_i\,\mathbf{1}(a_i=k),
\qquad
d_k=\frac1{N_k}\sum_{i=1}^N \delta_i\,\mathbf{1}(a_i=k).
\tag{7}
\]
Analogous formulas are given for convolutional layers by exploiting the sum-of-outer-products structure across spatial locations [2007.04532].

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 \(2\times\) a normal backprop step, while memory overhead is like storing \(K\) 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 [2007.04532].

## 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 [2007.04532].

Because raw variance alone is not comparable across tasks and training phases, the paper defines normalized gradient variance as
\[
\text{Normalized Variance}
=
\frac{\mathrm{Var}[g]}{\mathbb E[g^2]}.
\tag{8}
\]
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 \(1\), meaning noise power exceeds signal power; on MNIST and CIFAR-10 it often stays below \(1\). Since the denominator is shared across estimators on a fixed SGD trajectory, normalized variance preserves relative ordering while improving comparability across problems [2007.04532].

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 [2007.04532].

## 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 \(\tau_i\), and one samples \(\tau_i\) nonuniformly with importance correction. For smooth objectives, the recommended batch distribution is
\[
p(\tau_i)=\frac{1}{2d}+\frac{1}{2d}\cdot\frac{L_{\tau_i}}{\overline L_\tau},
\]
where \(L_{\tau_i}\) is the batch Lipschitz constant of \(\nabla g_{\tau_i}\). In least squares this specializes to
\[
p(\tau_i)
=
\frac{b}{2n}
+
\frac12
\cdot
\frac{\|A_{\tau_i}\|^2}{\sum_{i=1}^d \|A_{\tau_i}\|^2},
\]
so batch spectral norms determine the sampling law [1608.07641].

Deep heteroscedastic regression uses a different construction. Batch Inverse-Variance replaces the uniform minibatch average by a normalized inverse-variance weighted loss
\[
\mathcal L_{\mathrm{BIV}}(D_i,\theta)
=
\left(
\sum_{k=0}^{K}\frac{1}{\sigma_k^2+\epsilon}
\right)^{-1}
\sum_{k=0}^{K}
\frac{
\left(f(\mathbf x_k,\theta)-\tilde y_k\right)^2
}{
\sigma_k^2+\epsilon
}.
\]
Sampling is unchanged; the batch empirical distribution is reweighted internally by per-sample label-noise variances \(\sigma_k^2\). The normalization prevents near-zero-variance labels from causing an uncontrolled increase in the effective learning rate [2107.04497].

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 \(x=(x^{(1)},\dots,x^{(n)})\) is
\[
\pi_A(x)\propto \operatorname{Var}_A(x)\prod_{i=1}^n p_0(x^{(i)}),
\]
where
\[
\operatorname{Var}_A(x)
:=
\frac1n\sum_{i=1}^n \|A(x^{(i)}-\bar x)\|^2.
\]
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 [2606.22239].

| Setting | Weighted object | Representative formulation |
|---|---|---|
| Gradient clustering | Cluster representatives | \(\hat g=\frac1N\sum_{k=1}^K N_k g_{j,k}\) |
| Batched SGD sampling | Batch probabilities | \(p(\tau_i)=\frac{1}{2d}+\frac{1}{2d}\frac{L_{\tau_i}}{\overline L_\tau}\) |
| Deep heteroscedastic regression | Within-batch losses | \(\mathcal L_{\mathrm{BIV}}\propto \sum_k \frac{\ell_k}{\sigma_k^2+\epsilon}\) |
| Diffusion diverse sampling | Joint batch law | \(\pi_A(x)\propto \operatorname{Var}_A(x)\prod_i p_0(x^{(i)})\) |

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
\[
B_k=\left\lceil \frac{c}{F(w_k)-F^\star}\right\rceil
\]
for PL, strongly convex, and convex settings, and
\[
B_k=\left\lceil \frac{c}{\|\nabla F(w_k)\|_2^2}\right\rceil
\]
for smooth nonconvex problems. The key variance relation is
\[
E\!\left[\|\nabla F(w_k)-g_k\|_2^2 \mid w_k\right]
\le
(F(w_k)-F^\star)M_U^2 c^{-1}.
\]
This is variance-sensitive batch sizing, but it is not a variance-weighted distribution over examples or clusters [1910.08222].

Variance-reduced finite-sum methods employ a related idea at the epoch level. History-gradient adaptation sets
\[
N_s=\min\{c_\beta \sigma^2 \beta_s^{-1},\; c_\epsilon \sigma^2\epsilon^{-1},\; n\},
\qquad
\beta_s=\frac1m\sum_{t=1}^m \|v_{t-1}^{\,s-1}\|^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 [1910.09670].

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”:
\[
\left(\frac1\kappa\sum_{k=1}^\kappa g_k\right)^2
\quad\longrightarrow\quad
\frac1\kappa\sum_{k=1}^\kappa g_k^2.
\]
This preserves
\[
E[g^2]=E[g]^2+\operatorname{Var}[g]
\]
across merged micro-batches, so the expected \(v\)-statistic does not change with batch aggregation. The method is variance-preserving in optimizer statistics, but it does not alter the batch sampling law [2402.18824].

## 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
\[
\lambda^*
=
\frac{\mathrm{Cov}_f(\cdot,\cdot)}{\mathrm{Var}_f(\cdot)},
\]
not in a nonuniform batch distribution [2401.00493]. Variance-reduced random batch Langevin dynamics likewise samples shell neighbors uniformly and then subtracts the estimated force covariance from the thermostat noise,
\[
\widetilde{\boldsymbol\xi}_i^n
\sim
\mathcal N\!\left(0,\mathbf D^2\tau-\operatorname{var}(\widetilde{\mathbf F}_i^n)\tau^2\right),
\]
so variance enters as post-sampling correction rather than as a sampling law [2411.01762].

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
\[
\Omega
=
\begin{bmatrix}
\Sigma+\mu\mu^\top & \mu\\
\mu^\top & 1
\end{bmatrix},
\]
so the acquisition is covariance-aware in the multivariate second-moment sense rather than a variance-weighted batch distribution in stochastic optimization [1707.04191].

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.

Source: https://www.emergentmind.com/topics/variance-weighted-batch-distribution