Variance Matters: Improving Domain Adaptation via Stratified Sampling

AI summary

Overview

Research area: Machine learning — unsupervised domain adaptation (UDA), specifically stochastic variance reduction for domain discrepancy estimators. The work sits at the intersection of clustering-based data sampling and domain-alignment losses (MMD and CORAL).

Technical level: Advanced. The paper is heavily theoretical, presenting lemmas, three theorems with proofs (in appendices), rank-correlation-based error bounds, and a nonconvex quadratic program, followed by DomainBed benchmark experiments.

One-sentence scope: The paper introduces VaRDASS, a stratified-sampling variance reduction technique for unsupervised domain adaptation, derives stratification objectives for the MMD and CORAL losses, proves optimality for the MMD case under stated assumptions, and validates the approach on 4 domain shift datasets.

What This Paper Is About

Unsupervised domain adaptation trains models to be robust when test data comes from a different distribution than training data, typically by minimising a "domain discrepancy" loss between source and target features. The problem is that estimating these discrepancies from small random minibatches is very noisy, which can destabilise training and sometimes produce worse target performance than no alignment at all. This paper's goal is to reduce the variance of those discrepancy estimates by constructing minibatches with stratified sampling rather than uniform random sampling — without introducing bias into the gradients or otherwise compromising the model or learning algorithm.

Key Contributions

1. VaRDASS: The paper presents Variance-Reduced Domain Adaptation via Stratified Sampling (VaRDASS), which it describes as the first specialised stochastic variance reduction technique for UDA. The method only alters how minibatches are constructed, requiring no changes to the model or optimiser.

2. Tailored stratification objectives: The authors derive specific stratification objectives for two widely used UDA losses — correlation alignment (CORAL) and the maximum mean discrepancy (MMD) — reducing both to a dynamically weighted kernel k-means style clustering problem with objective sum over strata h of |S_h| times the within-stratum squared distance to the stratum mean.

3. Theory: Lemma 1 gives the variance of the squared MMD estimate under a normality assumption. Theorem 1 proves that when the covariance of the empirical mean embedding is isotropic (Σ = σ²I) and the Hilbert space has finite dimensionality d, minimising the surrogate objective Var(μ̂_φ,s) has the same minimiser as minimising Var(L̂_MMD), i.e., the surrogate is optimal. Theorems 2 and 3 give worst-case and expected error bounds for the surrogate solution in the anisotropic case, expressed in terms of a Spearman rank correlation ρ and the number of partitionings p = Stirling(n_s, k). Theorem 4 gives exact variance relations for the CORAL surrogate in the Gaussian scalar case.

4. Practical optimisation algorithm: Algorithm 1 solves the stratification problem via Lloyd's-style alternating optimisation using the kernel trick, initialised with kernel k-means++. Because the assignment step is a nonconvex quadratic program with an indefinite Hessian whose global minimum is NP-hard, the authors add Algorithm 2, a greedy incremental row-by-row heuristic that finds a local minimum in O(k n_s) time instead of the O((k n_s)³) complexity of a gradient-based interior point method.

Main Findings

• Estimator variance falls with stratified sampling: Using 1,000 feature embeddings from the Humpbacks dataset with a radial basis function kernel of unit bandwidth, the gap between uniform random sampling and stratified sampling grows with minibatch size k, with around a 50-fold reduction in variance at k = 256. The additional benefit of kernel k-means over linear k-means is only evident at smaller batch sizes.

• The MMD surrogate is well-justified empirically: In simulations for 4 different values of feature dimensionality d (with m = 1 and Σ_μ̂_φ,t = 0), the target and surrogate partitioning objectives show correlations all greater than 0.99; for d = 1 the correlation is exactly 1, matching Theorem 1 and making the error bounds in Theorems 2 and 3 nil.

• The CORAL surrogate also holds up: Simulations across a range of dimensionalities for 2 distributions (normal and lognormal, chosen to show an asymmetric non-Gaussian case) with k = 8 and cluster size ratio 1, …, k show near-monotonic rank correlation between the target Var(R̂_s) and surrogate Var(R̂'_s).

• The heuristic assignment is practical: On a problem built from Humpbacks embeddings with n_s = 100 and k = 16, Algorithm 2 was compared against a commercial interior point solver and an unweighted closest-centroid assignment, with performance also examined as random trials varied from 1 to 10^5 and parallel assignments varied from 1 to 20.

• Target accuracy improves: On the four datasets, VaRDASS achieved the highest accuracy in six out of the eight CORAL/MMD cases. Examples include Camelyon17 MMD (95.2 ± 0.6 vs 94.5 ± 0.7 for uniform random), Humpbacks MMD (92.0 ± 0.5 vs 87.7 ± 1.3), Spawrious MMD (71.6 ± 2.2 vs 64.4 ± 1.1), Spawrious CORAL (67.7 ± 0.8 vs 62.5 ± 0.6), and Office-Home CORAL (50.7 ± 0.2 vs 47.3 ± 0.3).

• Clustering inputs is not effective: The approach of clustering the input data directly (Zhao and Zhang, 2014) did not improve test-domain accuracy, which the authors attribute to the required smoothness conditions not holding for deeper models. Clustering features with plain k-means gave a decent improvement, and the diversity-based samplers k-means++ and DPP also improved accuracy, though inconsistently. No significant difference appeared between k-means++ and DPP, though k-means++ is computationally faster.

• VaRDASS beats other UDA baselines: In Table 2, VaRDASS performed best on all four datasets when compared against ERM, DANN, CDAN, CDAN + SDAT, CDAN + ELS, ARM, and MCC, as well as plain CORAL and MMD.

• ERM is unaffected by the sampler: The performance of non-adaptive empirical risk minimisation does not significantly change when varying the sampler, which the authors take as evidence that the CORAL and MMD gains come from the sampler's effect on the adaptation component of the loss.

• A higher computational cost: Wall-clock training times show VaRDASS is around an order of magnitude slower than simple random sampling and 2-3 times slower than vanilla k-means clustering. For example, on Office-Home: uniform random 625 ± 30 seconds, k-means (features) 1375 ± 68, VaRDASS 2163 ± 38, k-means++ 1977 ± 65, DPP 2529 ± 120.

• Post-training discrepancy is unchanged: Appendix D measures domain discrepancy between held-out datapoints not seen during training. These results do not demonstrate a difference in post-training discrepancy between standard and variance-reduced samplers, and all discrepancy values are already extremely small (the text states "two to three orders of" before being truncated in the provided content).

Methodology in Plain English

Start from the standard UDA setup: you have labelled source examples and unlabelled target examples, and you train a model that decomposes into a featuriser and a prediction head, minimising a task loss on the source plus a domain adaptation loss (weighted by λ) between source and target features.

Instead of drawing each minibatch uniformly at random, the method partitions the dataset into k strata ahead of time and then draws exactly one example uniformly from each stratum per batch. Because the sample is spread deliberately across the data rather than clumping, the resulting estimate of the domain discrepancy is less noisy. The authors prove when this works: for the MMD, under an isotropy assumption, minimising the variance of the empirical mean embedding is equivalent to minimising the variance of the MMD estimate itself.

To decide how to form the strata, the authors reduce both the MMD and CORAL objectives to a variant of kernel k-means clustering. In this variant, each cluster's contribution is weighted by its size, so large clusters are penalised more — encouraging, though not strictly enforcing, balanced clusters. For CORAL they cannot cluster on the covariance directly because it is not a kernel mean embedding, so they substitute a surrogate covariance that assumes the population mean is known, and use the mapping φ_c(z) = (z − μ_s)(z − μ_s)ᵀ with kernel κ_c(z, z′) = ((z − μ_s)ᵀ(z′ − μ_s))².

Clusters are optimised by alternating between computing distances to centroids (via the kernel trick) and reassigning points. The reassignment step is a nonconvex quadratic program; rather than solving it exactly with an interior point method at cubic cost, the authors use a greedy algorithm that fills in the assignment matrix row by row, each time putting a point in the cluster minimising D_ij(n_j + 1) where n_j is the interim cluster size. Randomising the row order and running many trials in parallel compensates for the greedy algorithm's tendency to land in local minima.

Since the optimal strata depend on the network weights, the authors re-cluster periodically every fixed T iterations rather than every step, relying on the observation that embeddings change smoothly during training.

Experiments use the DomainBed framework across four datasets: Camelyon17-WILDS (tumour detection, 1 training-evaluation split, 5 domains, 14,000 examples), Humpbacks (whale vocalisation detection, 4 data splits, 4 domains, 8,000 examples), Spawrious (4 dog breeds across backgrounds, 6 data splits, 6 domains, 18,664 examples), and Office-Home (65 object categories in different image styles, 12 data splits, 4 domains, 15,500 examples). ResNet-18 pre-trained on ImageNet is the backbone for Spawrious, Camelyon17, and Office-Home; Humpbacks uses the audio frontend and architecture from Napoli and White (2023). The MMD uses an RBF mixture kernel with bandwidths {0.001, 0.01, 0.1, 1, 10}. Algorithm 2 is run with 100 random trials and 10 parallel assignments, and T = 100. Models are trained with Adam for 3,000 iterations; hyperparameters including learning rate, weight decay, minibatch size k, and λ are tuned by a random search of size 10 on an in-distribution validation set, independently per sampler, with k drawn from 2^Uniform(3,7) and λ from 10^Uniform(−1,1). The full experiment set is repeated 5 times.

Why This Matters

Impact on research: The paper reframes a well-known practical failure mode of UDA — noisy discrepancy estimates in high dimensions and small minibatches — as a stratification problem, and supplies theory linking a tractable surrogate objective to the true variance. It connects UDA to the stratified and diverse sampling literature, and shows the technique adds value on top of Adam and implicit shrinkage through λ tuning. The empirical result that sampler choice does not move ERM accuracy is a useful control for interpreting future sampling work.

Real-world applications:

• Medical imaging across institutions: Camelyon17-WILDS is tumour detection in microscopic tissue images across samples from different hospitals, exactly the cross-site generalisation problem that limits clinical deployment.
• Marine bioacoustics: The Humpbacks dataset concerns detecting whale vocalisations across data from different acoustic monitoring programs, where recording hardware and conditions shift between sources.
• Robust image classification: Spawrious covers classifying 4 dog breeds across background environments such as desert, jungle, and snow, representing the spurious-correlation problem.
• Cross-style object recognition: Office-Home covers 65 everyday object categories across Art, Clipart, Product, and Real World image styles.

Industry relevance: Any deployment where labelled data comes from one source but the model meets another — cameras, sensors, hospitals, recording sites — faces domain shift. Because VaRDASS changes only data sampling, it can be dropped into existing training pipelines built around CORAL or MMD without modifying model architecture or optimiser. The cost is real, though: roughly an order of magnitude more wall-clock time than random sampling, and 2-3 times more than vanilla k-means clustering.

Future Directions

• Cheaper re-clustering schedules: The paper uses a fixed period of T = 100 iterations for simplicity, but notes that Liu et al. (2020b) proposed a low-cost criterion for deciding when to re-cluster. Adopting or improving such a criterion could reduce the wall-clock overhead that currently makes VaRDASS around an order of magnitude slower than random sampling.

• Closing the gap on the anisotropic case: Theorem 1's optimality guarantee depends on isotropic covariance. Theorems 2 and 3 bound the error otherwise via Spearman correlation, but whether a tighter or differently parameterised objective could raise the observed correlations further for larger feature dimensionalities remains open.

• Explaining the discrepancy paradox: Appendix D finds no difference in post-training domain discrepancy between standard and variance-reduced samplers, and reports that discrepancy values are already very small at that stage. The paper does not explain why reduced estimation variance during training produces better target accuracy without a matching drop in the measured discrepancy — a natural follow-up question.

• Extending beyond MMD and CORAL: The framework is derived specifically for these two losses. Whether the same dynamic-weighted kernel k-means stratification can be derived for other discrepancy measures — and whether the CORAL surrogate for the unknown-mean case can be tightened — is left to future work.

Target Audience

This paper is best suited to machine learning researchers and graduate students working on domain adaptation, transfer learning, or stochastic optimisation, particularly those already familiar with kernels, reproducing kernel Hilbert spaces, MMD and CORAL, and the DomainBed evaluation protocol. Practitioners building cross-domain production systems may also find the empirical tables useful, but will need the theoretical background to follow Sections 2.2 through 2.4. Readers looking for an easy off-the-shelf sampling trick should note the substantial wall-clock cost reported in Table 3.

Authors’ abstract

Domain shift remains a key challenge in deploying machine learning models to the real world. Unsupervised domain adaptation (UDA) aims to address this by minimising domain discrepancy during training, but the discrepancy estimates suffer from high variance in stochastic settings, which can stifle the theoretical benefits of the method. This paper proposes Variance-Reduced Domain Adaptation via Stratified Sampling (VaRDASS), the first specialised stochastic variance reduction technique for UDA. We consider two specific discrepancy measures -- correlation alignment and the maximum mean discrepancy (MMD) -- and derive ad hoc stratification objectives for these terms. We then present expected and worst-case error bounds, and prove that our proposed objective for the MMD is theoretically optimal (i.e., minimises the variance) under certain assumptions. Finally, a practical k-means style optimisation algorithm is introduced and analysed. Experiments on four domain shift datasets demonstrate improved discrepancy estimation accuracy and target domain performance.