The Benefits of Balance: From Information Projections to Variance Reduction
Authors: Lang Liu
Ronak Mehta
Soumik Pal
Zaid Harchaoui
Venue: NeurIPS 2024 poster (2024)
Keywords: regularized optimal transport
self-supervised learning
variance reduction
alternating projection
Primary area: learning_theory
TL;DR: We present a data balancing approach to distribution estimation that provides theoretical interpretations of the various self-supervised training schemes.

Abstract: Data balancing across multiple modalities and sources appears in various forms in foundation models in machine learning and AI, e.g., in CLIP and DINO. We show that data balancing across modalities and sources actually offers an unsuspected benefit: variance reduction. We present a non-asymptotic statistical bound that quantifies this variance reduction effect and relates it to the eigenvalue decay of Markov operators. Furthermore, we describe how various forms of data balancing in contrastive multimodal learning and self-supervised clustering can be better understood, and even improved upon, owing to our variance reduction viewpoint.

The Benefits of Balance: From Information Projections to Variance Reduction
Lang Liu"
Ronak Mehta"
Soumik Pal
Zaid Harchaoui
University of Washington
Abstract
Data balancing across multiple modalities and sources appears in various forms foundation models in machine learning and AI, e.g. in CLIP and DINO. We show that data balancing across modalities and sources actually offers an unsuspected benefit: variance reduction. We present quantifies this variance reduction effect and relates the eigenvalue decay Markov operators. Furthermore, we describe how various forms of data balancing in contrastive multimodal learning and self-supervised clustering can be better understood, and even improved upon, owing to our variance reduction viewpoint.
asymptotiiic statistical bound that
1 Introduction
Deep neural networks have shown remarkable success at learning task-specific representations of data when provided supervision from massive amounts of labeled training examples. Recent trends, however, have shifted toward task-agnostic, universal representations that may be easily fine-tuned or have zero-shot capabilities out of the box. Supervised learning, stricto sensu, is too limited framework for these billion-pante paarmmne,,,,tter, data-hungry models, and a question at the heart of modern machine learning learning from unlabeled, partially labeled, or weakly labeled data.
This need has paved the way for the current generation of self-supervised learning (SSL) approaches generic
that circumvent the need for large amounts of "strong" labels. In SSL, model is trained pseudo-tash suited for unlabeled data, such as relating image-caption pairs or augmentations of the same image. Despite modern foundation models such as DINO [Caron et al., 2021] and CLIP [Radford et al., 2021] 1 being trained in this fashion, many aspects of SSL remain mysterious.
In particular, the training process of self-supervised models often transcends the rules of the standard empirical risk minimization (ERM) toolkit. ERM combines two well-understood techniques: mini- batch sampling and gradient-based optimization using backpropagation. On the other hand, SSL adds clever, yet less-understood techniques to the training pipeline. illustrate this, consider a minibatch of independent and identically distributed (i.i.d.) training examples (X1, Yi)..... (Xn, Yr) P. where P is a joint probability measure on sample spaces X' x y (e.g. feature-labe or image-caption pairs) and let P. Σi= 8(x,,)) be the empirical distribution. For a model parameterized by 0 e ggd with loss function hgg a stochastic learning algorithm involves computing the minibatch loss
Ep, [he (X.
Σ he (X,,
and backpropagating through it to produce a minibatch stochastic gradient estimate. The algorithm then proceeds with the stochastic gradient descent (SGD) variant thereof (c.g.. Adam, SGD with momentum, etc). Self-supervised methods often modify this recipe by intervening on the optimization algorithm in minibatch-specific way.
'These authors contributed equally to this work.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

For example, SwaV [Caron et al., 2020] passes the minibatch examples through the model's encoder and clusters output vectors to generate pseudo-labels for prediction task. In teacher-student architectures such as BYOL [Grill et al., 2020] and DINO [Caron et al., 2021], the minibatch is passed through two networks, where the "student" network is updated via backpropagation and the "teacher" network is updated by cloning the student's 's weights in regular intervals. In CLIP [Radford et al., 2021], a model optimizes the sum of two cross-entropy losses, where the predicted class probabilities example i are generated by comparison to all other elements of the minibatch. While introducing such interventions into the procedure has clearly proven useful practically, it remains conceptually unclear what exactly is being optimized by the learning algorithm.
In this work, we aim to gain better theoretical understanding of the objectives and algorithms underlying these empirically effective recipes. In particular, we want to shed a theoretical light on their precise benefits over traditional learning methods. We show that such recipes often enjoy an unsuspected benefit: reducing the variance of the empirical minibatch objective.
Concretely, we formalize the model updates described above as two phrases. Let ..,.... En be minibatch containing data points of arbitrary type (e.g. unlabeled images). In the first phase, this original data source mapped (possibly using model parameterized by 0) another minibatch (X1, YI)..... (X.....,... Yn) of derived pairs in X' x V. For example, in SwaV, each 2, is an image, and we derive (x,, Y;) by setting X; = Z, and letting Y, be the pseudo- based on clustering the vector representations of the images. In CLIP, each Z; is an image-caption pair, and we derive (X, by simply letting X; be the image and Y, be the caption. Note that Y, is not a label in the traditional sense in neither of these examples. In the second phase, we use the model to compute a probability distribution Pn.o over X x V and perform a stochastic gradient update for the objective
Ep,,o (he (X, Y)].
(2)
This reduces to empirical risk minimization on the minibatch objective (1) when Z (X, (each data point is originally observed in X' x )) and Pn.0 = Pa (the empirical distribution of the data is used, regardless of the model). Beyond this setting, one specific example of Pr.e has been applied across various families of self-supervised learning (as detailed in Sec. 2), which we refer to as data balancing simply balancing, the primary subject of this work.
For probability measure Q on X x y, let Qx and Qy be the respective marginals on X' and y and let Qxly and Qyjx denote the respective conditional distributions. Given fixed target marginal distributions Px on X and Py on V, balancing refers to repeatedly applying the operations
R argmin [L[ and R r> arg min (T) Q:Qx=Px
(3)
Q:Qx=Py
in an alternating fashion. After enough iterations, the resulting probability measure approximately marginalizes to Px and Py in each variable. When X' and y are finite with | = m and | = these operations reduce rescaling the rows of an (m m)-matrix by Px/R, or its columns by Ry. This algorithm has a decades-old history and is known in other contexts as Sinkhorn Sinkhorn-Knopp matrix scaling [Sinkhorn, 1967]. iterative proportional or biproportional fitting [Johnston and Pattie, 1993], and taking-ratio estimation [Thompson. 2000]. The marginals Px and Py can represent auxiliary information or inductive bias from users, such as the desire for balanced clusters.
Returning to Pn.o in (2), we show in Sec. 2 that both self-labeling and contrastive approaches in SSL implicitly define Pn.e by the following steps: 1) constructing method-specific "initial" measure pro) on X x V, then 2) applying k iterations of the operations (3) to generate sequence pro), ..... p(k) and finally, 3) setting Pr.0- == Pi In other words, these methods embed learnable balancing operation in their objectives. A natural question to consider is: if the marginals one uses accurately represent the ones of the true probability measure P governing the data, are balanced quantities "better behaved" than their unbalanced counterparts? If so, in what way?
Inspired by this question, we fix the model parameter O (thus dropping the subscript from the quantities above) and analyze the fluctuations of the unbalanced and balanced objectives. The formal problem statement is follows. Let p. P., and p() denote the output of k 2 1 iterations of data balancing (see Sec. 3 for the precise définition). Finally, letting h : X x y - R be a fixed function of interest, we define the population parameter &p and its k-step balanced estimator sp by
x := Er Er (Y) and
pil := Ep (((X,))..
(4)

---

Our goal is to establish theoretical guarantees on the mean squared eтror (MSE) Ep - e) of estimating oo using 9/ with an informative dependence on the sample size n, number of iterations target marginals (Px, Py), and test function h. We are particularly interested in its comparison to the direct estimator based on the empirical measure 4p A Ei= (X,, Y.). as to quantify the effect of the auxiliary information (Px, Py). Our analysis uncovers two surprising facts. Firstly, while originally proposed for different purpose, balancing reduces the variance of the empirical estimate. Secondly, while the balancing iterations are nonlinear operations on the input measure, the variance reduction can be precisely quantified using the spectral decay of two linear Markov operators: the conditional means given X and Y, respectively.
Contributions. In Sec. 2, we detail the mathematical connection between balancing and the modern representation learning techniques mentioned above. In Sec. 3, we prove a new upper bound on the MSE of the balancing estimator sp(t). The bound decomposes into an O( 11 first-order variance term and an OA 10.fn-8// -/22 second-order term. The first-order term is shown to have a strict improvement over the empirical measure baseline with fine-grained dependence on the spectra of two particular Markov operators. The second-order term can be used to compute the asymptotic variance reduction for statistical efficiency comparisons. Our proof technique relies on a recursion decomposition for balancing-base estimators, which may be of independent interest. In Sec. 4, we illustrate how insights from our analysis can be practically applied to CLIP-type objectives and evaluation setups.
2 Data Balancing in Practice
To demonstrate a precise connection to (2), we describe how a collection of training examples Zi...., Zn observed in an original data space Z (c.g. grayscale images) is mapped to probability measure pn. Using the framework introduced in Sec. 1, this amounts to specifying four components: 1) the map from the original data into the derived sample spaces X' and V 2) the initial measure Pro; 3) the function h, and 4) the target marginals (Px. P for this measure to fit. From that point, the iterations of (3) produce P()) ..... P(O) and we set Pn,g P(t). For ease of presentation, we hide the dependence of p(k = Pn.o and h = hg on the model parameter 0. See Fig. 1 for examples of different choices of the sample spaces X' and y.
Example 1: Self-Supervised Clustering. Balancing is used in discriminative clustering and self- supervised clustering; see [Jones al., 2022, Asano et al., 2020, Caron et al., 2020] for variations on this theme. We describe the swapped prediction task of Caron et al. [2020] for concreteness but emphasize that clustering of this form is used
intermediate step (or the task itself) in many SSL pssudo-tasks. At a high level, this approach involves passing elements of a minibatch through two encoders generate vector representations. These representations are then clustered separately, and the features from one encoder predict the cluster label from the other encoders. Denote the encoders fe. ZZ- + R' and fe. Z + R., colloquially known the student and teacher networks. respectively, Here, we let [Z, P be a minibatch of n images, with
X = (Z1
), and y' (1......}.
where m = n and the elements of y' index learnable cluster representation vectors cI...., ........ R". Thus, we consider the overall parameter vector to be 0 := (0., 9,, cı....,c) q). Given temperature hyperparameters c, > O, the initial measure and loss function are given by the expressions
pio (x. y) x (x) (x)Tee) eg/e and 1(x, v) = log E/ (x) "///7"
Llc (o) (o) cu//
Directly optimizing Ex.y Pro (x,, y)h y) without any constraints would lead to collapse, so is balanced before optimization. The target marginals Px and Py are given by the discrete uniform measures on X' and y. This formulation is often derived by solving optimal transport problem with the Sinkhorn-Knon procedure is precisely p The intuition behind the choice of uniform marginal Px is that each data point has an equal amount of mass to be allotted, whereas Py captures that the cluster sizes are equal. The number of iterations l is selected based optimization considerations.
algorithm to assign soft cluster labels, the itcrative solution result from this
Example 2: Contrastive Learning. Contrastive Language-Image Pre-Training [Radford al., 1021], CLIP, is an architecture with an image encoder and text encoder that map joint
3

---

Balanced Clusters
Balanced Modalities
Balanced Metadata
Clusters
Modality
Metadata
Modaliy 2
of..
Minibachh
Triiing Set
Counts
*P"""
MApOOan DE.
Self-Supervised Clustering
Contrastive Learning
Metadata Curation
Figure Data Balancing Examples: Each panel shows possible distribution Q on different choices of ()) The orange histograms are the target marginal Py. Left: Q y) is the affinity of an image x for cluster y. Center: (xx, y) the similarity of image to a text caption y. Right: QIx, y) is the proportion of substring matches between a text caption and keyword y.
embedding space. Trained using image-caption pairs, the loss promotes representations such that images and text that paired in the minibatch close, whereas those that are not paired far. The latter aspect (promoting dissimilarity unpaired images/text) is what prevents collapse in this framework. To knowledge, interpretation of the CLIP objective balancing procedure is novel. Under this interpretation, we demonstrate that the objective is fact a nonlinear function of Pn.0, whereas its gradient will have linear form similar (2). In this case, each Z; (X,, Yi). where X; is an image and Y; is an associated caption. We have that
implicit data
X' = (X....,,
and y (Y .Yi...., Yx).
so that Consider an image encoder fo, X' + R' and text encoder for 3 + R' with parameter vector 0 = 07), The initial, unnormalized measure and the (in this case, vector-valued) function h are chosen based on these encoded representations:
(5) While we usually interpret as a loss function, we will show below that the CLIP loss depends nonlinearly on Pr.g. while the gradient has a linear dependence. If we believe, as in Example 1. that the target marginals (Px, Py) of the images and the text should be roughly uniform, we can apply the balancing iterations (3) with the target marginals being the uniform distributions over X' and ) respectively. Because there no preference for starting
po (I. y) x e (u) (8) and M(x y) = Va(Jo, (y)).
iterations with the or y dimension first, we may consider both orderings. Let Q be one iteration of balancing in the ) dimension and R represent one such iteration the dimension. Then the original CLIP objective L CLIP can be recovered (up to additive constant) as
POO (X,, Yt) P (I, Yi)
P (X,, Y,) Σy P (X,,
L CLIP
Σ
log
log
2
[=1
[log Q (X,, Yi) + log R (X,, Y.)
log n.
(6)
The measure Pr.o 1000 +1 R,) is constructed in this case by averaging the outputs of one iteration of balancing under each modality. Taking the gradient of (6) with respect to (whose dependence is contained in (Q R(")) recovers the expression for h in (5). The objective is often interpreted as an average cross-entropy loss terms, each representing the prediction of one modality's 's original pair from the other. In our formulation, I,CLIP can also be viewed as average negative log-likelihood under the Q and R also of interest to study the effect of using Qu) and R for k: 2 in general, as we show in Sec. 4.

---

Example 3: Metadata Curation. Here, we consider balancing an entire training set, as opposed to particular minibatch. the billion-parte parameeer scale, dataset design can be the primary factor that differentiates performance between foundation models [Fang et al., 2013, Xu et al., 2024, Gadre et al., 2023]. One general approach used in both the original CLIP dataset [Radford et al., 2021] and an open-source replication [Xu et al., 2024] is metadata curation, wherein a text dataset (possibly captions for images) is synthesized using a list of keywords (yI...., . so that
X = (=, Zn}. )' (v1....,, 3r
meaning that m = n. The keywords are matched to texts within X' via substring matching. While the approach of Xu et al. [2024] (dubbed MetaCI.IP) pools all matched keywords on every text to measure the "distribution" of keywords, we consider a version in which each text Z; can only be labeled with a single keyword yj,. This allows for a true joint probability measure on X x Y. The marginal distribution of observed keywords is initially long-tailed (see Fig. 4) (e.g., "the" will match many more texts than "xylophone"), In both Radford et al. [2021] and Xu al. [2024], the data are resampled so that this distribution of keywords over matches is closer to uniformity, i.c. keywords with many matches have their associated texts downsampled during the dataset creation process. While the probability measure may not be computed explicitly (due to scale). this adjustment of the keyword distribution can be viewed as a single iteration of balancing (3) applied to the y marginal. For tasks such as language modeling, we have
p (x, y) = (z v) and Ma, = Eg(2))
where (x( denotes the loss of model evaluated at single text I X' (notice that the keyword is not used). We elucidate this connection by applying direct balancing on a subset of the ImageNet- Captions dataset in Sec. 4, observing the effect on downstream model performance.
Motivated by these scenarios, we address the statistical problem outlined in Sec. by analyzing balancing-baad based estimators. We then return to examples mentioned above in Sec. 4, illustrating how the theoretical analysis can be translated to algorithmic variants.
3 Theoretical Analysis of Variance Reduction
We now present theoretical guarantees the mean squared error (MSE) of the data-balanced estimator (0th)) and highlight relevant points in the proofs. For readers' convenience, a notation table (Tab. 1) is in Appx. A. We first give context the main innovations of the analysis and then outline its high-level steps. These innovations include relating the nonlinear iterations of balancing over probability measures linear operators a vector space and using singular value decomposition of these operators to quantify their effect after a finite number of iterations. Furthermore, by scaling the number of iterations appropriately, we can characterize the estimator using the limit of balancing iterations, which is an object of interest in applications including optimal transport.
Preliminaries. Recall the setting introduced in Sec. 1, in which consider sample spaces (X, )), along with true and unknown joint distribution P on X' x ) with known marginals Py)) For ease of presentation, we assume that | = | = m. although the arguments do rely on equal support sizes. We make the following assumption throughout, which is usually satisfied by the desired marginals Px and Py, such in the uniform cases discussed in Sec. 2: the target marginals (x(r) > 0 and Py (v) > 0 for all x E X' and E V. We define p and for k 2 1 construct
= Pr as the empirical measure
(x( ( v)
1 odd
p(t-1 (a)
(v( (x
A even
p(k-1) (v)
By direct computation, we see that the iterations in (8) are equivalent to applying (3) for A odd and even, respectively. See Fig. 2 (left) for visualization of this procedure. The iterations well-defined for all & under the event that Supp(Pn.x) i.e., all observed row counts and column counts are non-empty.
Supp(Px) and Supp(P Supp(Py).
'Due to this technical consideration, we define p( to be the empirical measure P., when this condition is not satisfied, which show occurs with probbabibll See Appx. D.4 for details.
5

---

P
P
(0)
(2(xx)
pxh
L2(Py)
Qx Px
Qy = Pr
P
L((P))
L LP(P))
Figure 2: Data Balancing, Nonlinear and linear operators associated with each iteration of (8). Left: Visualization of the exact iterations (8) in the space of probability measures. The blue set contains joint distributions with X--mrginaa equal to Px. whereas the orange set contains joint distributions with J'-marginal equal to Py. Right: Visualization of (P)) the operators defining (11), and singular values given in (13).
To provide background, the scheme of alternating the operators (8) is often seen as an iterative algorithm to solve the problem
min (EI((x.PP)
(9)
where II(xx, Py) denotes the set of probability measures on X' x y that marginalize to Px and Py in each variable and KL(-II-) denotes the Kullback-Leibler divergence. The iterations (8) are based on the alternating minimization approach of solving
(x, = arg (Q:Qx=xx> KL Pk--)) odd arg (i KL PI
k even
which inspires the viewpoint of balancing as alternating information projections. As we show in Appx. C, the iterations of (8) can equivalently be defined using the KL, reverse KL, or -dvvergggss This viewpoint relevant as previously, efforts have been made (e.g. Bickel et al. [1991]) to analyze the variance reduction afforded by the solution to (9) directly. However, quantifying the variance reduction (in terms of properties of P) using this approach is challenging, as there is no closed-fomn expression for the solution of (9). A key mathematical outcome of analysis is that the closed-form expressions of the projections (8) can be used to compute the reduction in mean squared error each iteration. Thus, by letting l k((n) - (scaled appropriately against n), can determine the reduction for the solution of (9) for large n. This is the subject of Thm. 1.
From Information Projections to Orthogonal Projections. First, we will show that the variance reduction resulting from each nonlinear iteration of (8) is associated with a linear operator applied to h. Thus, instead of analyzing the alternating information projections may use familiar tools to understand alternating orthogonal projections in a vector space. To define them, first let (PP to be the set of functions h X' y IR satisfying Ep [h²(X, Y)] < co. Even though X y is finite, working within L2P(P) will be analytically convenient. Let (P be the subspace of (PP containing functions that only depend on the first argument E X and define L (Px) analogously. These are the solid-colored subspaces in Fig. 2 (right). Next, let px (P) + L and py LL2(P) + (Py be defined as, for any h E L2(P),
probability measures,
uxh arg min Ep (h(X, f(x)) [uxh](E, = Ep (XXX SEL?(Px)
The operator μ is an orthogonal projection onto (Px) The orthogonal projection operator Jy onto (Pv) defined analogously. We may also define the conditional debiasing operators Cx = I I-ux px and Cy I -y, which each project onto the orthogonal complements of L (Px) and

---

(Py),, visualized as subspaces with dotted border in Fig. 2 (right). To understand the importance of the conditional mean and debiasing operators, we give recursive formula that forms the backbone of our analysis. Define /44- μx for k odd and μı = jy for l even, and define C. similarly. Thus, by using the notation (h) == E.g (x) Y)), we have by linearity of expectation that
[P Ph)
[P(t) Pl(Cah) + (P()) p(()))
[P(t-1) - p)) + P -PC-P)(CCA))
( (.CI...Chh)+37 C.h)+ [P Pa PK-
(10)
first-order term
higher-order terms
To justify the first line, we discuss the case when & is odd. Notice that uxh is only a function of X, so its expectation only depends on Px that is equal to Pre (the X'-maginnal of P())) by (8). The last line follows by unrolling the previous step k 1 times. This recursive expansion is proven formally in Prop. 15 in Appx. D. Given the expansion, the mean squared error can be computed by taking the expectation of squared (10). We show that the second moment of the first-order term in (10) is equal to of/n where
of Var(h) and of == "aar(C Cxh) for k 2 1,
(11)
and all other terms are n//23 Thus, by exactly computing the constant in the dominating term, we may quantify the asymptotic variance reduction. Our first main result concerns the higher-order
terms and shows that it is indeed dominated by the first-order term. Note that the empirical mean
pp - (X;;, Yi) is unbiased, and so its MSE is equal to o n. Define in addition
P. == min (min Px (z), min (v)
which measures the non-uniformity of the target marginals. We have that p. is positive because both Px and Py positive. We now state the first main result.
Theorem 1. For sequence of data balancing estimators (so(h) 1:211 (=2)));:1 defined in (4), there exists an absolute constant C >> O and distribution dependent constant 8 [0. 1) and such the following holds for 0z sup of of: For 2 C'llog2(21) (alogg(2n/pp) (2n/p.) m log (n + 1)1/p² and k 2 1. we have
0
s
Ep [(=)t }p) v
gap O
+
(12)
n
71
3/3
The quantities dgap and are quantified toward the end of this section and are dependent on eigende- cays of the conditional mean operators for each variable under P. Furthermore, o² gap > 0 except for the pathological case of pxh being a constant function. Showing Thm. boils down to showing that the higher-order term in (10) is O(n -1) with high probability. Using the expression (8) and assuming that é is odd, we see that
(x(x) p n, (x)
-1] :
[P P P( -.ICE...CA C+h)=
Σ
[C (x) (z,
The first (blue) term in the product quantifies the disagreement between the X' -mmmrginaal of p and the true marginal, which can be bounded in terms of KT. Px and is shown to be Onn-1/2) with high probability via techniques from information theory. The second (orange) term can be unrolled recursively in a similar fashion to (10) itself, which will consequently be O(n -1/2) as well; this is the most technical part of the analysis (see Appx. D.3). Our analysis also yields bound for the sensitivity of balancing to misspecified marginals; see Appx. D.5.
Given Thm. natural next step is to quantify the gap between of and of. which requires finer- grained properties of Cx and Cy. Notably, we show that as k + co, of approaches a limiting value. Thus, via (12), by using & = o/n2/12) obtains asymptotic variance of the solution to (9). This contrasts with Albertus and Berthet [2019], which the dependence of a quantity similar to (12) exponential in k. meaning that A o(log(n)) is required for convergence under this argument.
7

---

From Orthogonal Projections to Variance Reduction. We now clarify what is precisely meant by the "spectrum" of the conditional mean operators μx and jyy As proven using singular value decomposition (Prop. 3) in Appx. B.1, there exists basis {a;]y=1 of L (Px). a basis ( 1 of L (Py), ), and real values {sz 5- that satisfy
My j Sy and px93 = sya, for E m......m..
(13)
Furthermore, 0n = 1x and 3] = ly leading to the equality (J, (1,0))(((x)) = Epx [/(X]. Finally, $1 and s, is non-negative and non-increasing in For a concrete example, consider m 2. in which case can be written as a matrix in 2 and elements of L2 and L2(Px) are vectors in R.. Then, in the case of uniform marginals, we can verify directly that (13) can be satisfied by setting
31 i 02 B2 [ and P
1 1+ss 1-s 1 - 1+
01
(14)
for 32 (the second largest singular value). Thus, as 1, the distribution becomes "fully dependent" as and are completely determined by one another. As 8 - 0, P approaches the product measure. Geometrically, because 01 31, know that the angle between the subspaces [Px) and (2 is given by the angle between 02 and 32. By computing their inner product in (P))), we have that (arg, (2)) = (P, ay ) = 8 = cos a. Thus, of these subspaces, alluding to the independence of X and Y (see the right panel of Fig. 2).
indicates orthogonality
Returning to m 2 2, we consider the following as sufficient condition for variance reduction: the operators μ.x and Jyy have a positive spectral gap, i.e., 82 < 81. Note that this assumption is satisfied when P(z, y) > O for all (I. y) E X x ) by the Perron Cerron-Frobenius Theorem [Horn and Johnson, 2013, Chapter 8J. Using the intuition from Fig. 2, this rules out pathological cases such as Y being deterministic function of X. Under the spectral gap condition, the singular values {55 1 that
strictly less than 1 will determine geometric rate of decay in variance given in Cor. 2. The left and right singular functions a; : X + R and B, : y + IR will define a useful coordinate system to represent projections of h when analyzing (!!t)
Indeed, let h P(h) be the centered test function. Because uxh L2(xx) and py E (P
we may decompose this function on the two bases to write
uxh
ajay and pyh 1=1
0
(15)
Cor. 2 below relates the (normalized) variance of of the first-order term to the one of the sample mean pp In fact, it shows that the variance reduction oi of decays geometrically to the quantity
1 8,U,)² - 8
Σ j=2
gap
For simplicity, we only present the result for even, i.c.. 22.- Corollary 2. The variance reduction achieved by + iterations of the CyCx operator can be quantified
m 8 (
- s,
m
83 11+2 (uj 8T ug + (1
a 22 ((++))
2 gap
- )
j=2
j=2
Intuitively, the operators Cx and Cy are the main sources of the variance reduction via orthogonality. that the reduction will always be strictly positive as long as yx is not
Since 01 = Lx, constant function. Finally, using s = $2 2 8j for j 2 gives the second term in Thm. 1.
4 Numerical Illustrations
We illustrate how data balancing manifests in the motivating examples mentioned in Sec. 2 with experiments with CLIP-type models. We focus here on zero-shot image classification tasks. Details on these experiments, and additional ones including linear probing and zero-shot retrieval, as well as an empirical investigation of the sensitivity to misspecified marginals, are all contained in Appx. Code to reproduce the data and experiments can be found https://github. http::/githuuccm/roonkdmm/abb wtt.//ooaakmmm/maaaniiinn
8

---

CUP Text Embaddings
BERT Text Embeddinga
GPT-2 Text Embeddinge
128
512
02 1288
512
128
512
CIEAR
0.35
033
04
0.23
04
0.15
0.2
02
02
0.0
0.03
100
0.2
0.00
02
0.01
0:01
0.0
0.5
05
0.3
0.4
66
20
04
0.2
Eeralions
llerations
lteratione
lerratooss
lerations
Iterations
No Belencing
CLIP Balancing (X 1)
MURI-CLIP (A 22)
Figure 3: Zero-Shot Classification Performance across Embeddings, Batch Sizes, and Objectives. The three vertical panels describe different choices of the text encoder fer which increases in quality from left to right; that is, pre-trained GPT-2, BERT, and CLIP embeddings, respectively. Within each vertical panel, examples include batch sizes m 128 and m 512. Rows indicate various evaluation datasets from CIFAR-I0, CIFAR-100, and STL-10. The y-axis of cach plot indicates average per-class recall, whereas the araii indicates training iterations at the given batch size.
Model, Datasets, and Evaluation. Throughout, we consider training variants of CLIP models (see Sec. 2), which require dataset of image-caption pairs. For the training set, use the ImageNet- Captions dataset [Fang et al., 2013], which pairs images from ImageNet [Deng et al., 2009] that were taken from Flickr with their original captions. In the notation of Sec. 2, the model is specified by selecting an image encoder fa, and a text encoder foy- In all cases, we use fixed image/text encoder base vector representation and compose it with trainable feed-forward neural network, i.e., fo = fbccdd o p We fix the hase image encoder as CLIP ViT-B/32 architecture pre-trained on LAION-2B [Schuhmann et al., 2022], and vary the base text encoder across embedding models of varying quality: GPT-2 [Radford et al., 2019], BERT [Devlin et al., 2019], and CLIP- based encodings. When two CLIP encoders are used for the base image/text vector representation, they are taken from separate CLIP models (i.e. the base representations are not dependent). We evaluate models based on zero-shot classification performance using the standard CLIP inference procedure: for any image x, a label E (I....,C, C} is predicted by associating to each e natural language prompt ye, and predicting the scores s(x) = (s1(2),... sc (2)),, with
e (z).for
sc(x) De_ efo, fo),
(16)
for a temperature 7 > 0. Multiple prompting strategies can be used depending the evaluation dataset, for which we average embeddings before applying (16). We use the public CLIP Benchmark repository, using the datasets CIFAR- 10, CIFAR-100, STL-10, with their default caption sets.
Data Balancing Effects. Fig. 3 shows the zero-shot classification performance (in terms of average per-class recall) of variants depending whether the contrastive learning objective from Sec. is used or not. One iteration of balancing already leads to improvement in terms of downstream performance. Multiple balancing iterations lead to further improvements. See Appx. E for more details on this experiment, and for analogous ones with linear probing and zero-shot retrieval.
Fig. 4 then shows how balancing can be used to adjust an entire pre- training set to given marginals based on metadata, as described in Sec. 2 in the metadata curation example. After balancing, the target marginal has less than 2 orders of difference. In terms of downstream performance, data balancing leads to some improvement in the smaller batch regime (m 512) when curating the dataset. See Appx. E for more details on this experiment.

---

Histogram of Metadaiaa Calegariiss in Pre- Pre-Training Dutaset
Zero-Shul Accuracy of Modkls with Diflerent Pre- Training Dala
Oripas Dats
Revelanee Des.
Erodeadd CiNAA-10e
Exaleaad on STL-10
2.30
4.5
107
630
2000
1030
20.05
Tnanin Jontons
Trihning Ramniies
Figure 4: Balancing and Metadata Curation. Depiction of balancing and metadata curation (Example 3 in Sec. 2) on ImageNet-Captions dataset, in which X' represents image- caption pairs and ) represents keywords. Left: Observed marginal P,,y y (orange) and Py (blue), which are sorted by order of increasing probability. Right: Zero-shot evaluation of an embedding model trained using the standard CLIP loss original versus the balanced training set.
Related Work Self-supervised learning has witnessed a surge of recent interest as datasets and computing hardware allow for larger, more capable models (see Balestriero et al. [2023] and references therein). While we highlight in this paper the connections between data balancing and contrastive learning [Radford et al., 2021], we acknowledge that data balancing distillation" approaches more broadly [Grill et al., 2020, Chen and He, 2021, Oquab et al., 2024].
also be related to "self-
Historical motivations for data balancing include census or survey data, in which P. is cross- tabulation of (a limited number of) paired observations and the target marginals were estimated from large amounts of unpaired observations [Deming and Stephan, 1940, Ireland and Kullback, 1968]. This situation is not unlike the present day-yet at a different scale-in which the amount of unstructured modality data (such as images) still dwarfs the amount of high-quality multimodal data [Gadre et al., 2023]. Bickel et al. [1991] proved classical asymptotic results on balancing estimators. Linear operators similar the ones we use in Sec. 3 also appear in their analysis. More recently, Albertus and Berthet [2019] studied such estimators from an asymptotic empirical process viewpoint. Our theoretical results significantly improve those from Albertus and Berthet [2019] primarily in the dependence of the number of iterations it on the sample size n to achieve convergence guarantees (from logarithmic to polynomial).
Matrix scaling is a popular algorithm for solving entropy-regularive optimal transport (EOT). We refer to Peyré and Cuturi [2019] for survey. See also Courty al. [2017]. Shen et al. [2018], Peng et al. [2019] for interesting methods based on EOT in machine learning. Entropy-regularized optimal transport was one of the original inspirations for SSL techniques such as SwaV (see Sec. 2). While EOT itself deterministic optimization problem, a related statistical problem is the large-sample limits of EOT solutions when the marginal measures are estimated from data [Mena and Niles Weed, 2019, Genevay et al., 2019, Klatt et al., 2020]. We emphasize that, while this line of work shares the matrix scaling algorithm with our setting, the statistical problem is entirely distinct; in statistical EOT, the target marginal distributions are computed from observations of independent, unpaired data, and the initial measure can be computed from the cost function. In our setting, the data dependent, forming the random initial measure P.,, whereas Px and Py are fixed auxiliary information.
5 Conclusion
We showed how several disparate techniques used towards the training of foundation models are instances of a data balancing algorithm, which has the unsuspected benefit of reducing the variance of learning objectives involving multiple sources of data. We proved new non-asymptotic bound on the mean- mrean-squared error of balanced estimators they adjust to the given marginals. We also highlight the key roles of conditional expectation operators in quantifying that variance reduction effect. Finally, we translated the marginal balancing interpretation of several training practices for foundation models into algorithmic variants that warrant further investigation. Exploring variants incorporating prior information on the data sources is also an interesting venue for future work.
10

---

Acknowledgements The authors are grateful to G. Ilharco, M. Wortsman, K. Pillutla, L. Schmidt, and J. Wellner for fruitful discussions related this work. This work was supporied by NSF DMS-202316 66, CCF-2019844, DMS-2134012 PIMS 20240827-PRNO1, NIH, and TARPA 2022- 22072200003. Part of this work was done while Liu was with the University of Washington, and while R. Mehta and Z. Harchaoui were visiting the Simons Institute for the Theory of Computing.
Broader Impact While this paper is of theoretical nature, the web-scale pre-training sets used to train foundation models can affect not only the biases of the models themselves but also the behavior of individuals who interact with them. In the case of representation learning, unrefined Internet data may lead to non-uniform performance among protected attributes such as gender, age, etc. For generative models, individuals of all ages may be influenced by harmful images or textual output. Studying the relationship between the balancing procedures considered in this paper and more holistic model evaluations presents a valuable direction for follow-up work.
References
M. Albertus and P. Berthet. Auxiliary information: The raking- ratio empirical process. Electronic Journal of Statistics, 13(1), 2019.
Y. Asano, C. Rupprecht, and A. Vedaldi. Self-labelling via simultaneous clustering and representation learning. In /CLR, 2020.
R. Balestriero, M. Ibrahim, V. Sobal, A. Morcos, S. Shekhar, T. Goldstein, F. Bordes, A. Bardes, G. Mialon, Y. Tian, A. Schwarzschild, A. Wilson, Geiping, Q. Garrido, P. Fernandez, A. Bar, H. Pirsiavash, Y. LeCun, and M. Goldblum. A cookbook of self-supervised learning. arXiv preprint, 2023.
Bickel, Y. Ritov, and J. A. Wellner. Efficient estimation of linear functionals of probability measure P with known marginal distributions. The Annals of Statistics, 1991.
M. Caron, I. Misra, J. Mairal, P. Goyal, P. Bojanowski, and A. Joulin. Unsupervised learning of visual features by contrasting cluster assignments. In NewrIPS, 2020.
M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin. Emerging properties in self-supervised vision transformers. In ICCV, 2021.
X. Chen and K. He. Exploring simple Siamese representation learning. CVPR, 2021.
N. Courty, R. Flamary, A. Habrard, and A. Rakotomamonjy. Joint distribution optimal transportation for domain adaptation. In NeurlPS, 2017.
T. M. Cover. Klements of Information Theory. John Wiley & Sons, 1999.
W. E. Deming and F. Stephan. On a least squares adjustment sampled frequency table when the expected marginal totals are known. Annals of Mathematical Statistics, 11, 1940.
J. Deng, W. Dong, R. Socher, L.JJ. Li, K. Li, and L. Fei-Fei. ImageNet: A large-scale hierarchical image database. In CVPR, 2009.
J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In ACL, 2019.
M. Everingham, L. Van Gool, C. K. I. Williams, J. Winn, and A. Zisserman. The PASCAL Visual Object Classes Challenge 2007 (VOC2007) Results, 2007.
A. Fang, G. Ilharco, M. Wortsman, Y. Wan, V. Shankar, A. Dave, and L. Schmidt. Data determines distributional robustness in contrastive language-image pre-training (CLIP). In ICML, 2013.
S. Y. Gadre, G. Ilharco, A. Fang, J. Hayase, Smyrnis, Nguyen, R. Marten, M. Wortsman, D. Ghosh, J. Zhang, E. Orgad, R. Entezari, G. Daras, S. M. Pratt, V. Ramanujan, Y. Bitton, K. Marathe, S. Mussmann, R. Vencu, M. Cherti, R. Krishna, P. W. Koh, Saukh, A. Ratner, S. Song, H. Hajishirzi, A. Farhadi, R. Beaumont, S. Oh, A. Dimakis, J. Jitsev, Y. Carmon, V. Shankar, and L.. Schmidt. DataComp: search of the next generation of multimodal datasets.
In NewrIPS, 2023.
11

---

A. Genevay, L. Chizat, F. Bach, M. Cuturi, and G. Peyré. Sample complexity of Sinkhorn divergences. In AISTATS, 2019.
I. Gohberg, S. Goldberg, and M. Kaashoek. Classes of Linear Operators Vol. 1. Springer, 1990.
J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. Richemond, E. Buchatskaya, C. Doersch, B. Avila Pires, Z. Guo, M. Gheshlaghi Azar, B. Piot, kavukcuoglu, R. Munos, and M. Valko. Bootstrap your own latent: A new approach to self-supervised learning. In NeurIPS, 2020.
M. Hodosh, P. Young, and J. Hockenmaier. Framing image description ranking task: Data, models and evaluation metrics. Journal of Artificial Intelligence Research, 2013.
R. A. Horn and C. Johnson. Matrix Analysis. Cambridge University Press, 2013.
C. T. Ireland and S. Kullback. Contingency tables with given marginals. Biometrika, 1968.
R. J. Johnston and C. Pattie, The Professional Geographer, 45, 1993.
and the iterative proportional fitting procedure.
C. Jones, V. Roulet, and Z. Harchaoui. Discriminative clustering with representation learning with any ratio of labeled to unlabeled data. Statistics and Computing, 2022.
D. Kingma and J. Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
M. Klatt, C. Tameling, and A. Munk. Empirical regularized optimal transport: Statistical theory and applications. SIAM Journal on Mathematics of Data Science, 2020.
T.-Y. Lin, M. Maire, S. Belongie, L. Bourdev, R. Girshick, J. Hays, P. Perona, D. Ramanan, C. L. Zitnick, and P. Dollár. Microsoft COCO: Common Objects in Context, 2015.
S. Maji, J. Kannala, E. Rahtu, M. Blaschko, and A. Vedaldi. Fine-Grained Visual Classification of Aircraft. Technical report, University of Oxford, 2013.
G. Mena and J. Niles Weed. Statistical bounds for entropic optimal transport: Sample complexity and the central limit theorem. In NeurlPS, 2019.
M. Nutz. Introduction to Entropic Optimal Transport. Lecture notes, Columbia University, 2021.
M. Oquab, T. Darcet, T. Moutakanni, H. V. Vo, M. Szafraniec, V. Khalidov, P. Fernandez, D. HAZIZA, F. Massa, A. El-Nouby, M. Assran, N. Ballas, W. Galuba, R. Howes, P.-Y. Huang, S.-W. Li, I. Misra, M. Rabbat, V. Sharma, G. Synnaeve, H. Xu, H. Jegou, J. Mairal, P. Labatut, A. Joulin, and P. Bojanowski. DINOv2: Learning robust visual features without supervision. Transactions on Machine Learning Research, 2024.
X. Peng, Q. Bai, X. Xia, Z. Huang, K. Saenko, and B. Wang. Moment matching for multi-source domain adaptation. In ICCV, 2019.
G. Peyré and M. Cuturi. Computational Optimal Transport: With Applications to Data Science. Foundations and Trends Machine Learning, 11, 2019.
A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners, 2019.
A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, Goh, S. Agarwal, G. Sastry, A. Askell, Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In ICML, 2021.
C. Schuhmann, R. Beaumont, R. Vencu, C. W. Gordon, R. Wightman, M. Cherti, T. Coombes, A. Katta, C. Mullis, M. Wortsman, P. Schramowski, S. R. Kundurthy, K. Crowson, I.. Schmidt, R. Kaczmarczyk, and J. Jitsev. LAION-5B: An open large-scale dataset for training next generation image-text models. In NewrIPS, 2022.
Shen, Qu, W. Zhang, and Y. Yu. Wasserstein distance guided representation learning for domain adaptation. In AAAI, 2018.
12

---

R. Sinkhorn. Diagonal equivalence to matrices with prescribed row and column sums. American Mathematical Monthly, 74(4), 1967.
M. E. Thompson. Theory of Sample Surveys. Chapman & Hall, 2000.
H. Xu, S. Xie, X. Tan, P.-Y. Huang, R. Howes, V. Sharma, S.-W. Li, G. Ghosh, L. Zettlemoyer, and
C. Feichtenhofer. Demystifying CLIP data. In ICLR, 2024.
13

---

Appendix
Table of Contents
A Notation
15
B Linear Operators and Variance Reduction Singular Value Decomposition
B.2 Proof of Main Results
16
C From Information Projections to Data Balancing
19
C.I Balancing as Information Projections
19
C.2 Proof of Main Results
22
Statistical Analysis of Balancing Estimators
24
D.1 Recursion of Estimation Error
E...............................
25
D.2 Technical Tools & Intermediate Results
.e..............................
26
D.3 Analysis of Higher-Order Term
27
D.4 Proof of Main Results
29
D.5 Misspecified Marginal Distributions
35
E
Experimental Details
Datasets
44
E.2 Model Specification and Hyperparameters
44
E.3
Compute Environment
44
E,4 CLIP and Multi-CLII
45
E.5
Metadata Curation
45
E,6
Additional Experiments
45
F
NeurIPS Paper Checklist
14

---

A
Notation
Symbol
Description
X.Y
Sample spaces for two data sources.
Support sizes m [&] and = D
We sometimes assume m for ease of presentation.
Probability meusure un x y' (the data generating distribution).
Sample size.
(X,, Vi............ Y.)
Independent and identically distributed sample from P. Empirical measure ((X,, Yi
Po
Qx.Qv
Marginals of measure Q on x V. e.g. Rx, Py. Pa, etc. For mensure over 2, the set of values 2 such that ((+)
Supp(Q)
Q(k)
The expected value of under Q, or Rg [6[X, Y)
(P(2))2r
Sequence iterations of (8).
Iteration count of (8).
The (Supp(Pr. (uxp(Px) and Supp(Pa,x) Supp(
Test function h: X x - & of interest.
The estimand Ess., M(x, p(x.))))
The estimator Σn, M(x, (x, y). when well defined.
estimator 1P.. 4/ Is Pi Ptlls..
(1)
Normalized emorr (T 4).
Remainder defined in Prop.
Centered function h Ey [t).
of
Variance term Ep [Ca... C. min(min, Px(x), min, (y))))
LP(P)
Functions l 3' x y + (as X' 3/ is linite).
LP(P))
Subspaces of (PP) containing functions only
&' and e respectively.
Conditional expectation operators (uxh)((r) == YX and (py))(y) Ep ()(
PX-MY
Cx.Cy
Debinsing/centering operators Cx 7 px and Cy = T- py- (xxxCx] for odd and [yy for even. Singular values in Prup. 3.
He.Cz Ca
[sy)]m
{as) //1 (5 pust
Bases for 1 Ix(P)) and (/2()))) Prop.
Table 1: Notation used throughout the paper.
B Linear Operators and Variance Reduction
This section is dedicated to establishing the variance reduction result in Cor. by employing properties of the Markov operators introduced in Sec. 3. In the first part, we establish Prop. 3, the singular value decomposition that defines the quantities appearing in Cor. 2. In the second part, we quantify the difference between of and of for even and odd iterations of k.
B.1 Singular Value Decomposition
Recall the conditional mean operators μx and jyy from Sec. 3,
[uxh[(x) == E [h(X, XX (x) and [uyh](y) == E (M(X (v),
with the corresponding debiasing (a.k.a. centering) operators defined by Cx = - px and Cy = I ny.
Proposition 3. There exists a basis (ayl7m ofL2(Px), basis {3,1m= 2 and real values
{, which satisfy:
Py s,β, and μx s, E {1...., m)
(17)
a1 1x. B1 1y. 1 and sy is non-negative and non-increasing in
15

---

Proof. When ya is restricted to (Py) and py is restricted to (2 these operators are in fact adjoint in L L2(P), as by the tower property we have the relation
(f. (F.1))) E [ [X] = E (X) (uy (a))
Since : ( (Py)) is compact linear operator, by Gohberg et al. [1990, Section IV.1 Theorem 1.1] and Gohberg et al. [1990. Section IV.1 Corollary 1.2]. we have that μuy admits a singular value decomposition satisfying (17). Next, we show that $1 S 1 and that 1x is an eigenvector of μup (Px + L2xPx) with eigenvalue 1, which confirms that $1 = and a1 1x by the definition of singular values (arguing symmetrically achieves $1 = 1y). By the variational representation of singular values [Gohberg et al., 1990, Section IV.1 Equation (2)), we have that
-
sup ||ay fll $1-
Consider any f E L2(Px) such that [ = 1. Define the conditional probability Pxly(zly) = P(x,,)////H/ (y) which is well-defined by assumption. Then, by the Cauchy-Schwarz inequality in LL (Pxjr).
| f = Σ Σ
(()))( (())xx( (r(
Vee
xx)
NEY TEN
Σ (z)
P(r.y)
ZEX
VEY
=1
This proves that S 1. For equality, notice that μXpY 1x px ly 1x, completing the proof.
B.2 Proof of Main Results
Lf and (P respectively. These bases span the range of the operators μa and pyy We will consider the repeated application of the operator CyCx, sequence of two centering operations some function h E L [P),,, and compare
From Prop. 3,
we establish two bases {az}]== and {B,}ا
E [((CyCx)*h]] against E [A²]
for h Ep [b]. We establish the main result by measuring the reduction in variance from a single application, terms of the coordinates of the function of interest each of the subspaces. We will then observe how these coordinates change crration-to-iteratin to give the final result.
Lemma 4. For any h LP(P) such that Ep [h] 0, let
uxh = Σ uja and pyh h =
Bg
j=1
Then, we have that
m
m
E [CyCxh))]=
s,u,)²
j=2
j=2
Proof. By orthogonality, we have that
E [CCC) E [[(I pay Cxh)
EE' [[Cxh)]]] 21E (CCc
TP. [[Cxh)]]] -PP ()
= -E [[Cxh)]] Cx((1))x1))
= E [6a] - (x((uxh))) - - ((xx(x))))
16

---

Because P(A) 0, it holds by the tower property of conditional expectation that Px (uxh) = 0. which implies that
"1 = (1x =0 = Px((uxh)3) - m
Σ U
For the second term, observe that Px(Cxh) 0, so it holds by the tower property that Py (yycCC = 0, so
Py((nx(x1)))= = Σ (uxCx)),)))))))))
Next, compute the term in the square by applying Prop. 3:
(uyCxh, B)) (uy h 3r) ())r) huy
V/j
UACCk; Bj
k=1
L2(Ps)
Σ
1/
L2(Bx)
U/ 8y11y
which completes the proof.
Lem. 4 ensures that we have reduction on each iteration, with formula that depends on the coordinates of the function on each subspace. Because these coordinates change every iteration, we track them in the next lemma. Define ho h and he+1 (CyCx) yhe along with the constants m given by { Tm and [ūt5)]=1
m Σ
m Σ 3
pxh
ag and jy jy h4
We have the following. Lemma 5. For all & 0, it holds that
U++1 83 8y
Ut+1J
Proof. Fix any j E [m]] and use Prop. 3 to write
U++1.
(Ix(I - px- - BY + HyHx) ai
PxXt u.aak. aj m
m
Fx
Be aj
LP(Px)
L2(Px)
sT 8yVt.j+
which proves the first part of the claim. For the second part, note that nyCy = 0. so (urCrC ( 0.
Using Lem. 4 and Lem. 5, we can simply accumulate the reduction incurred every iteration.
17

---

Proposition 6. Define the constants (u,) and (uz) by
Fxh = Σ ujas and by = 1y
Then, we may quantify the variance reduction achieved by + 1 iterations of the Cy-C operator as
+ 07 - 44)
E
h22
u + (v3 10--8 syu,)²
2 j
5=2
(V/ 87 )2
ug +
- 8)
==
as t + 20.
Proof. Apply Lem. 4 + + ti-me that
E (
= E [
[11 + s + t. ur [rJ]
j=2 r=0
m
E [T2]
Σ
vo.j
25, 2s Vo. + (1+
j=2
as by Lem. 5, have that Uz 0 for T >>>>> 0. Next, we unroll the definition of u so that
ur.t j ,jur-1.3 - 83 115
= ( - 8,U7-25) - 1
5 -2 (83 s,00.))
for > 0, yielding
E h
E [[CrCx)++11)
ub.5 + (vo.) s) 2 + (1 + 83) 8y00,))
(s))T-1
Σ
u + (Vo.) - )² 22+ + (1 + +3) -8ymo.)) " (55)T
1=2
7=0
m
n + (Vo.j s) + 21 + 5)) 8y00,) )) - 540)
Σ
-
1=2
v + (vo.. - 3² ( syu0,) s1
1 8
Substitute vo.j My and vo.j = vy to complete the proof.
We also present the corresponding result for odd. The proof follows similarly by repeated application of the operator CyCx. However, the iterations will be compared to of Ep [[Cxh] )2 as we consider Cx h as the "first" iteration to this process. Proposition 7. Define the constants (u,) by
m
CyC h
Σ U3
1=1
18

---

We provide the derivations for three common choices of D: Kullback-Leibler (KL., reverse KI., and x² Using this viewpoint, and simply assuming the positivity of the marginal measures Px and Py, derive an upper bound in Prop. 14 that is constant in k. This is an improvement over the recent work of Albertus and Berthet [2019], in which they show an upper bound that scales exponentially in Then, we may quantify the variance reduction achieved by + 1 iterations of the CxCy operator as
83 s
E [[xx))]] -E [((xx)))**
= Σ
2 Y + (8ju))²
14 1 87
2m
m
+ 8
4z
j=2
as / xc.
In order to have full monotonicity, we also need that of 2 of. This follows by orthogonality,
- EE [(Cxh) + E [(xx)²]] =o + [(xx)] 2 of.
(18)
Thus, we can combine Prop. and (18) to fully quantify the relationship between of and of for odd.
C From Information Projections Data Balancing
This section is dedicated to deriving three representations of the balancing procedure as projections in various statistical divergences, as shown Fig.
We consider two sets of probability measures denoted by IIx = (Q: =x= {Q: : Qy }. The marginal matching steps written projections in terms of statistical divergence (precisely, an f-divergence) in the form
= Px} and Ily =
Px p 1J n.X
Py n.Y
p(k-1) arg min ((Q()P=-"))
& R= arg min min D(Q ((Q QEIly
QENxx
The KL representation will be used in the proof of Prop. 14, which (recalling the sequence (P()))22 from (8)), controls the error between Piyy and Py for & odd and Pit and Px for even.
C.1 Balancing Information Projections
C.1.1 Projection in KI.-Divergence
Proposition 8. Assume that Px & Rx and Py & Ry, and define
Q = arg min (T) p = arg min RT) Qelly
(19)
Qellx
Then, it holds that
(Px(e)Xyxx((lz))
ix >0 ix()) =0
(*(x,v)
(20)
10
and
( (xx
if Rr Ry
(xxx)) =
(21)
0
19

---

Proof. In the case that Q(r = 0, we apply the convention that 0 log 0 = 0. Consider the case Q" the projection R onto Ixx. Write
Σ Σ Q(a,y) log (vxx()))))x()) DEA
RT(Q]RR)
Rr (x2
= Σ Qx(x)
Σ (rix( log ( (11x(11))
TEA
PEY
= =Σ Qx(x)
Σ Qrxx log Qvxx(v))) (xix(
Σ
(2)x( log Qx(w)
Rx(=)
JEA
VEY
DEY
= E (x(x)
Σ
Qrix log Qrix(u))) Rrix(
Qx(x) log @x(x)
Rx(
DEN
NEV
#E
Σ
Qx(x)) KL
KL (I
x
Σ (x(x) Tx
(L((1 + R[x))
where the last line is due to the marginal constraint Q 1Ix. For the above to be well defined, need that Px & Rx so that KL(Px Rxx) < +00. The above is minimized when Qy (y/lr) Ryix (y(z) for all (r, y) E X x y such that Qx(x)= = Px(x) > O. The case of Pa follows analogously when using that Py < Ry.
C.1.2 Projection in Reverse KL-Divergence
Proposition 9. Assume that Py & Rx and Py & Ry, and define
Q = arg min QL( P = arg min KL(R(IQ). QEIly
(22)
Qelxx
Then, it holds that
(x(())y(x((11))
if Rx(x)
(²(x,v)=
(23)
0
if (x) =0
and
Rx)
(y) if Ry (1)
(*(x,
(24)
0
Proof. In the case that R(x, y) apply the convention that 0 Olog 0 0. Note that minimizing KL(RIIQ) over Q is equivalent to minimizing - Ex.y R(I, y) log Q(x, y) (i.c. the cross entropy). Consider the case Q", the projection of R onto IIx. Because R & Q for KL < +oc to hold, we have that R(x) >
((x) 0, that Qy 2) is well-defined. Write
Σ R(x, y) log Q(x,y)
2,
Σ (x) (x((2))
Σ R(x) Σ
(v (v)x(1
nEX
ZEX
NEY
Σ Rx(x)) log Px(x)
Σ (x)
Ry]x(11x)) Qrix
ZEN
The second first term does not depend on Q due to the marginal constraint Q E Tx. The second term the expectation of the cross entropy from Ryjx to Qyix over Rx, which minimized Ryyx Qyix. We have specified Qrix and Qx. completing the proof.
20

---

C.1.3 Projection in -DDierreee
Let 1 denote the function that is identically equal to 1, Consider the following optimization problem, which is the subject of the subsequent lemmas:
min
|| -1l
(25)
tεAx
where
Ax == f: f:&xD + R satisfying
R(z.)) = Px(x) for any x E X
yey
Lemma 10. Assume that Px & Rx, and define The problem (25) is feasible and its solution can be written as
= (1-)++
for any f E (RRR where the linear operator C is specified by
[c y) = g(z. y) Σ
g(x, Ry)xxx
y'EY
Proof. First, we establish feasibility by letting
f(x,y)) (x(x)/Rx(x) if Rx(x)>
otherwise
This function does not depend on the second input y. Because we assumed that Px < Rx, we have that the terms of fxx y) for which Rx (z) = 0 do not affect whether Σyey f(x, ((r,y)((r, y) = Px (z). because Px (z) O in these cases. In the remainder of this proof, we will show that (25) is an affine projection problem, and find its solution by converting to subspace projection problem. Indeed, consider fI....,fy fr E Ax, and oi,,..,,, R such that Σ 03- = 1. Then,
Σ DEY
Σ (2))) 1=1
-((,,))
Σ. aj j=1
fx, g(, = Px(x),
DEy
indicating that Σj=1 (z, v) E Ax and Ax is an affine subset of (RR).. Define
Sx
9 g: x RR satisfying
Eglz, (olx,y)Raa u) = for any a E X
DEY
Then, for any f E Ax, we have that g E Sx if and only if g + Ax. Taking any f E Ax. letting o² be the solution of
min
-il
(26)
eESx
we will have that Ø² + f will be the solution of (25). The remainder of the proof is showing that 6- (1
First, define the operator & by [uxg[(E, v) = -yyen g(z, Rylx(y' (z), and note (by factoring out (x)x)) that E Sx if and only u = 0. In addition, 2²g is linear and idempotent as " = g. so is a projection operator in ( Thus, Sx the orthogonal complement of (ange(($) (n)), ), and the solution of (26) is given by (I (T- -) = C (1 -f), because C T-
The claim is proved.
Lemma 11. Assume that Px < Rx. Define
Q'
== arg min x² (Q) QENx
(27)
and let g* be the solution of problem (25). Then.
Px (x
Rx)
Q*(x,v) (z, 9)
(28)
0
(x) =
21

---

Proof. First, by reparametrizing the problem (27) as finding such that Q(x,y) = E(z, v) R(x, v).
we can compute its solution by solving
min 11 |1-E1lia()) geAx
(29)
Notice that we also have non-negativity constraint, opposed to (25). If 5* solves (25) and happens to be non-negative, then we have that &* solves (29) as well and the first equality of (28) satisfied by definition. We show the second equality of (28) by direct computation, which also establishes the non-negan non-negativity of g* simultaneously.
Apply Lem. 10 with
[xx,y) (x((=)/Xx())
if Rx)
11
otherwise
so that
(x C (1 (x. f(z.z)
Σ Ry)x()) f(x. v)
f(x.))
LEYY
= f(z.v)
for any y V. Thus, the likelihood ratio of Q with respect to is a marginal reweighting. Accordingly,
Px x(vlz)
if > 0 Rx) =
( = ((,,
00
completing the proof.
Proposition 12. Assume that Px < Rx and Py < Ry. Define
Q = arg min x R)) P = arg min (Q)
(30)
Qellx
Qelly
Then, it holds that
Px (xx(x)yyxx((l))
if {x) 0
Q" (z, v)
i (2) 00
Py Rx) (xlv)
Ry > if By =
(²(,,
(31)
0
Proof. The first equality of (31) follows by the claim of Lem. 11. The second equality follows by repeating the argument of Lem. 10 and Lem. with xX and (Y, y) swapped.
C.2 Proof of Main Results
We may now control the errors of the ratio of marginals using the projection interpretation established in the previous sections. Recall the event S as defined in Tab. 1. The following result, the monotonicity of the marginal violation terms in terms of KL, will be useful in the bound.
Proposition 13. [Nutz, 2021, Proposition 6.10] Under the event S, it holds that
XL(P Px) (L(Pv|I > KL(P X 2
We give the following result for X and the analogous claim holds on Y Proposition 14. Assume that Pn,x (x) for all E X. It holds that
(x
|max fn - 1, 1] mmax(1/22 1,
k 1 ifk 1.
max
p
(x)
(32)
n.X
22

---

In addition, we have that
Px (x)
n
(T((Pr,x11xx) =
max TEA
pit (xx
5
KL
i
p
Moreover, when KI,(Pa,x x, Px) S p2/2, we have
(x(x) x (x:
-1 -
2 KI(Pn.xIIPxX
max wEX
P. V2
(33)
Proof. We first show that (x) 2 1/n for k = and ( 2 p. for k > 1. In the case that k the result follows directly from the event S. For k > such that & odd, we have that for x E X,
1 (x) = Σ pit (x. = Σ (r
pi "(x y)
nx
yEy Pryy2 (1)
BEY
2
P 2) (I, p. P (x)
p. Px 2
yEy
The result for even can be proven similarly. We now proceed to prove the inequalities given in the statement, which will rely on the lower bound above.
Proving the first inequality. Then, for any x E XX
(x
(x
Px(x) (x)
max 1, 1,13 =
p(k-1) n,X (x)
max
" nxX (x)
1,1
pia-1) n.
m1x 1,13
if
which is the desired result for the first inequality.
Proving the second and third inequalities. Consider an odd 2 By the definition of total variation distance, it holds that
max T
Px(x)
P n.X
(x) s TV(P
n.
n.X n.x Px).
According to Pinsker's inequality, we have that TV(P(tx") n. Px) V have that
and so we
n.X - (x)x"()) S V n.X
(L(PPXXX|PA) n,
max X
-
V
where the last inequality follows by the monotonicity of Sinkhorn iterations given in Prop. 13. We apply the lower bounds to write
Px (x)
n
(I(P.,xIIPx)
if =
max #
pik p
(I
if
Finally, when thus
KL(Pr..11/2 S p+/2, we have that max
Px(x) n.X <p./2 and
((tt)))
min pik (xx)
min Px(x) max xEX
pik 11
(xx Px(x) 2
P+ 2
NEX
n.X
NEX
n.X
Hence,
(x) pik- (x) - (x)
(x(x)) p 11 (x)
max,ex
x,X
2
mingeX p "(x)
P+
23

---

Now, for i even, set lt = 2t for t 2 O. We have that
1 n.Y ",p s (1.
p
n.Y ( - Pr (v) - TV(Pay
max yE
Invoke Prop. 13 once again to achieve
KL(Pr|IPP n.Y
KL(P,,x(IPx)
which completes the proof.
D Statistical Analysis of Balancing Estimators
This section contains the proof of the main result, namely Thm. 1. We first consolidate notation and then give broad outline of the proof for readability. Let the expectation of function h under probability measure Q on X' x y/ by denoted by
Q(h)
Σ zEX,YEy
(xx,))((E,))
so that
10 p (h),, = P(h),
and
G (h) v P (() vr((pn P(h)).
(34)
Recalling in addition that C. = Cx for & odd and Cx = Cy for even. The event S = (Suup(P>x.> = Supp(Py and Supp( = Supp(Py))
(35)
is used for purely technical reasons in many results.
Proof Outline. We first establish that the recursion formula
(P Pl(h) = [Pr [p(k-1) (l(Ch) + V (Cah)
holds in Prop. 15, where
Ens
Px
(x) -1 h(x, P(k- v) k odd
pla-t
V(-- (h)
(36)
Ex.
1y (w) -1 1(0,, Pxr k even
n.Y
The quantity (Cah" describes an error term that accumulates for each iteration of balancing. which explains why k must be scaled appropriately against n to ensure the eπrоr does not accumulate too fast, Applying the recursion repeatedly the balanced sequence (Px*)) and unrolling the recursion, we see that when l: is odd,
[P(N] (h)
(P)-1) Pl(Cxh) ((--"(Cxh))
[p)-2) P](CrCxh) - V + +V ((x-t
[ Pl Cxhi Σh Vi (C...
(37)
first fidet-order term
higher-urder term
Additionally, let he == Ce Cah,, so that the first-order term can be written as P(o) (h1,,)) (hh,,) higher-order term also be written Ex=1 Vice (he,))) Because our original goal is to upper bound the mean squared error, we use the expansion above to write
E (())()) (h) P(L)) < E P (h1,k) - P(hn,)))
+ 2T Pio0 (I1,,)) P( Ein Vie (he,)) + E Ex Vi "(hea))
n
Regarding the first term, we have that E P. (h1,)) p(h of/n, which is the dominant term in Thm. 1. Thus, the remaining challenge of the proof will be to upper bound the cross term and squared term and show its dependence on n. The dominant term of these two will be the cross term, as we will essentially show that |
p(h,,)) is O(n -1/2) with high probability and that E> 1= V "hex))] is in fact Oln with high probability. As stated in Sec. 3, a key intermediate result controlling the higher-order term Prop. 14, whose proof given in Appx. C. The remaining subsections walk through these steps in detail.
24

---

D.1 Recursion of Estimation Error
We first recall that the sequence (Ph )221 can be computed with the following formula:
Ph
1) "xx v) k odd
pik
P (I.z) == (x 9) and Pa (2,y)
(38)
(y) Ph p)t-2) (z.
k even
Prop. 15 establishes the conditions under which these steps are well-defined (i.e. (x) > 0 and PKty" (v) 0).
Proposition 15. Let (Pr))),,, be sequence computed according to (8). These iterations are well-defined under the event S, and G.0 defined (34) and Vi0 defined in (36), it holds that
(h) = G G(*-" (h) + (nV,
(39)
and
(h) G " (Cah) VVV (Cxh).
(40)
Proof. First, assume that X (x) > 0 and " > 0 for all E X' and E ) so that we may establish the recursion, which we will show by induction toward the end of the proof.
Consider the following steps in the case that k is odd:
1 (h)
Px
Σ M(r, P( Σ
Mx,
v)
(x)
by (38) for k odd
2.
2.9
n.X
Px
Σ 1 h(2, PMk
P.k 1) (I, v) +
Σ
pik 1J (x)
-1
y. (( "x,
z.y
n,X
20 (h) (h).
Arguing analogously for k even and subtracting P(h) on both sides, we have that
[P Pl(h) = [P Ph)h) + Vi vh)-
(41)
We refer this as the "uncentered" recursion, which proves (39).
We can then establish the following "centered" recursion using the following decomposition in the case of & odd.
1 - Pl(h)
=
[p - -P](Cxh) + (P(M) - (])xx)
h = Cxh++ + uxh
[Ppk-- - ((xCh) + V (Cxh"((xh) + [p(h] Pl(ux))
apply (41) to Cx
=Pr 1> P](Cxh) + Vi "(Cxh).
Ps (uxh) = P(uxh)
The last line follows because yx is only function on X, and due to the definition of the marginal rebalancing iterations, P = Px. This gives the desired formula by substituting (34).
We proceed to show that the iterations are well-defined. We will in fact show that "itx'"(a (z) > 0 and 1 (y) > 0 for all X' and y Y. For k (x) Pr,x(x) > 0 and pio) ,, PA, (y) 0 for all X' and y E Y' this holds under the event by assumption. We argue by induction that this holds for all & > 1. Assume that the claim is true for {1,.... k 1}, and that is even. Then,
n. (x) = (x)x) > O.
Px p (x,
may "())"(y)= = Σ pakk "(x,1)=
Σ
(c)
xE
ZEN
n,
Px p n.X
min
(x)
P(
(v)
(( > 0 and p(ty?()) (z) 0 by the inductive hypothesis. Arguing analogously for & odd
achieves the claim.
25

---

D.2 Technical Tools & Intermediate Results
Having established the backbone of the argument, we collect in this subsection some useful tools that are used in the remainder of the proofs.
The following result follows from the method of types in information theory and will be helpful in deriving the dependence of the higher-order term on n.
Theorem 16. [Cover, 1999, Theorem 11.2.1] Let v be discrete probability measure supported on m atoms. Let U....., Un i..dd 0 that
v and Un be the associated empirical measure. Then, we have for any
P(KL(Un u 2 e) S 2 log(n+1)
We then provide a result that counts the number of terms that appear when repeatedly centering via the operators C,...., Cs. This formalizes the pattern
Cx = -x
CyCx = px -0--- py PY
CxCyCx= = - + HYNX + μxpy - μхPУΡX.
and so on. This will be useful when bounding h uniformly.
Lemma 17. For any k 2 I and E {l,...,k).
(k (h-2-1)/2
Mxx - (k-8-1)/2
C, I-
Σ
Σ (uyux)ur
(k-1)/2
(k-1)/2
Σ
(KXHY) +
Σ (DYRX) + (-1)k-(43 pg....
where the sum Ei O when i and is Etd when is not integer by convention.
Proof. We prove the claim by backward induction on E, for the case that l odd. In the case = k the claim holds because Ck I - pgg Next, for any < k., assume that the stated result holds for {e + 1...., k}. Then, & also odd (so that pe = px).
Ce. C,Cx+1..... ...
(k-
(k-B-2)/2
=--
Σ
(μXPY)
Σ
(NYUX) bY
T=U
7=0
(k-2-1)/2
(K-2-1)/2
+ Σ
(
+ Σ
(NYUX) + by
px
k- terms
(K-8-11/2
(K-1-2)12
xx +
Σ
px(μYpx) pY
(k- 1)/2
(k-2-1)/2
(
Σ
#x(myxx) - (AXHY (k- (K- px
The red terms and blue terms cancel out to zero. This leaves
(k (k-1-2)/2
1-2-211/2
C
(μxPY) -
Σ
(NYIX) py
T-U
T=O
(K-R-1)/2
Σ
(mY#x)* + (uYXx) ) (1--2/2
+
(k- 22122
+
Σ
Ux(pypx)IMy + (-1)k- pe... PE
7=0
26

---

wherein we combine the red terms and re-index the blue terms to get
(k-t t ---2-2)² 2)
(k-P-2)
Ce..Ca=I- II-
Σ
(UXPY)
Σ
(nYUX)
(k-0//2
(k-0)/2
+ Σ (NYRX)* + Σ (uXpN) + (-11 pee p4-
Finally, because L even when is odd and is odd, we can set the upper bound of the first two sums to (k: - - without changing the number of terms. This proves the desired result. The result can be proved similarly when is even. As a result, we have proved the claim for any odd and l < k. Similar arguments be used for the case of k even and l k.
D.3 Analysis of Higher-Order Term
Returning to the outline the start of this section, we may now bound the higher-order remainder term in (37), namely
V
(hek)) = Σ
(C....G...C
1=1
f=11
depends on controlling the quantity Vi6-11 the summation, which we recall for convenience:
Ex.y
plk 17 (x) -1 11(,, P. 1xx,,))) y) odd
VO (h)
(42)
Ensyy
(v) -1 h) p k even
Because have established uniform control over the functions Px n. - and Py/P n.Y via Prop. 14 in Appx. C now bound the full remainder in Prop. 20.
We also make use of the following intermediate result, which controls how large the loo the function l can grow after centering.
Lemma 18. |l < 2(k + 1) ||
Proof. Apply Lem. 17 and the triangle inequality, so that we only need to count the number of terms that appear in the sums, adding 2 for the first and last term in the expression. We subtract 1 from the total, as one of either (k f)/2 (k- l 1)/2 will be fraction. This yields 2(k E 1) terms total, the desired result.
We upper bound the sum in Prop. 20. To do so, we introduce some notation. Consider B₁ and Bg defined by
max
P.. p!f-!) (a)
Px(x)
l odd
B Mr and
B2 = max Mf 25E5k
for M =
Py
maX
p(t-")
l even
for k 2 1. We also enumerate the sample spaces X {I1 Fm) and y {yI...., vm. and define the function
1 =====1)) Io odd 1 ,, even
1 (x,
This is an indicator function on the j-th element of either X' or )' depending on whether l is odd or even. Finally, for any function h, use (under the event S) recall the empirical process notation
G Vn (P (h) P(h)).
(43)
Using this notation, we can rewrite the recursion in terms of the quantity ())())) itself. This established in the following lemma.
27

---

Lemma 19. For k odd, it holds that
Px -
G CCx + Σ
G (Crh
Pl (rj)
whereas for even, holds that
m
Py(
G. (Cyh) + (h) G
Σ
G
(Cyh
Put 11
(Vj)
Proof. We give the proof for & odd. By (40) from Prop. 15 and by the definition of (h),,, we need only show that P(Cx/15x) = O. Indeed,
(
[E|CxH|X]()) (1r))
if =
E (x)))NIX(())) - 10
But E [Cxh|X (z²) = by definition of Cx. Taking an expectation over Px gives that P(Cxh1gB) = which implies the desired result. The proof for is even follows symmetrically. O
The higher-order term in (37), can be bounded using Prop. 20.
Proposition 20, For any 2 1, the following holds under the event S:
C.h)l <
Bi G (hi, B2
Σ G (be, 2m
+ m B2 | Vnk(k - 1 + B2 (k + 1)/3]
Proof. First, for any E {1....,A) k). recall the notation he, and by Lem. 19, we have that for odd,
By (39) from Prop. 15
Px pie-n
VnViz" (ha,a)
Σ J=1
(x²) -
GØ-
(44)
n.X
Using the statement above, we have that
m
V (ie S Me
Σ
Gk- 1,))
The bound above holds for & even as well. Then, using the (39) from Prop. 15 again, we have that for 2 2,
[K D Pl(hxx = [Pp 2 Pl(hex + Vie (hexx)j))
which implies that
(he | < (he V
(he,,,,e))) + V (he, + +... +r.
(he,e)e)) + M1 MivnP(o (heal 1,e) +...+ MrVnnP - 15e) (he, + [/]] vn B + BB - 1)] - f (+ 1),
(45)
28

---

by Lem. 18 and M1 S B1 and Mt - B2 for & 2 2. Summing these bounds, we have that
Vu-1 (le,,))
Em1
m
m
S M1
G ( [e)++
Σ
Mt
G
(hex
j=l
8=
j=1
m
Σ Bi
(h je)l + B2
G
2=
m
$ B1
G (h1 (h1 +
B2
(10 (he, [a,) + + 2 vn |B + B2ll 1)) e (+1))
apply (45)
2=
B1 G (b1, + B2 Σ
G (he, +
j=1
8=2
2mB2 Mbloo vn
B1 + B2(t - 12) ++
2=2
because X = m. We sum up the last term:
k-1
Σ Bi (2 1)] (K+ e 1) B1
(k () B2
Σ
f(k e)
1=1
t=1
k(k - 1)
[B1 B2 1)/3)
2
completing the proof.
D.4 Proof of Main Results
We can now show the main result of this section: the bound on the mean squared error of the rebalanced estimator. Recall the event
= (Supp(Pn.,) (hxpp(Pxx))-ee = Supp(Px) and Supp(Px,y) Supp(Py)}
(46)
as introduced in (35). To remind the reader of the high-level steps the proof, we may decompose the сттоr on the event S we used the estimator
$ := ( + 10
so we decompose on the event write
Ep (P, (h) (h)- Ep (Ph 18=
Ep
(P P(h))* 1s]
(47)
Then, the upcoming Prop. 21 bound the first term, which will in turn require showing that occurs with high probability. As for the second term, we will apply Prop. 15 and the derivation (37) write
Ep (P) P(h))² is Ep [TP1s] 2Fp [TTTAA + Ep [T]
(48)
for
T = (P - P| Cahh and T2 =
(C........Ca)))
(49)
t=1
By definition, we have that Ep [1 < Ep T term Ep [Ti s) and squared term Ep [ This is accomplished by Lem. 23 and Lem. 22, respectively.
/ It then remains bound the cross
29

---

Proposition 21. It holds that P(S*) I 2m(1 - P.)) Moreover, for any d E (0, 1). we have
Ep (b)) P(h)) se s < min (2m(1 p+)". + 2 (2/(2 | 2m(11 P+)"
Proof. Define Fx := [Supp(Pr.xx) # (Ppp(Px)} and Fy == ( m,y) + Supp (Pyp)())) Py)}, so that 8 =x Fy. We first control the probability of Fx. Let F3 =- x (I2) - 0} for j E [m]. We
then obtain Fx - Fj, which implies by the union bound that
m
m
P(Fx)
P(F))
Σ (1- Px(a,)) S 5m(1 p.)*
j=1
Similarly, we have that P(Fy) E m(1 - p.)" and thus P(S)) S 2m(1 - p.)", which gives the first claim.
To control the expectation, consider any > 0, and define the event
{ Pi (1) - P(h) S
2 log (2/0)
Eg =
||
n
By Hoeffding's inequality, it holds that P(Es) 2 Furthermore, we get
s[Iss (P.h (h) P(h)) E[ts= cg (P.)) (h) P(h)) + B(1s= Le, (P) (()))))
< 4 1bll Eis 1sy log (2/5)
||Al E 1c,1
< 4 |M min P(E))) 2 log (2/8)
|| P(S*)
S 4 | min - p.)", 83 log (2/5) ||b| 2m(1 p.)
In order to bound the terms appearing in (48), we introduce the events Ef. Eg. and & defined by
1
2 + m
log(n + 1)
E max
||)) - log2
8
F/
[1G (hr. 1,2)l S v /log((2kk/5)2(k + 1) || }, = 1..... k, j = I....,
33 0 Fl t=1
3g |G (h1,)) V 2log(2/6)22 |hill20]
The events are constructed such that P(E)) 2 1 - 8, P(E)) 2 1 8, and P(E)) 2 1 - 6, as we used in the upcoming proofs of Lem. 23, Lem. 22, and Thm. 24.
Lemma 22 (Squared term bound). Let T2 be defined as in (49). For any > 0, assuming that (2/8) + m log(n + 1)l/p3, we have that
2 |Al| m²
Ep [
<
(oog(2/8)) + m log(n + [(k=1)
1-1
A +
8 n²
2mk
(k + 1) + (A: - 1)-1)( + 4)
An +
+ +
8
+
22log
2
p2
p²
Proof. The following computations are done under the event S. First, apply Prop. 20 to write
b
Tz <
B1 G (hi, 152) + B2
Σ
G (he, 15t)l +
/n
j=1
2=2
m m,B ||!l| k(k 1)[BI + B2 + 1)/3).
(50)
30

---

We decompose on the event Ef n Eg. Note that by Thm. 16, we have that P(E 2 1 1 6. It follows from Hoeffding's 's inequality, the union bound, and boundedness of llñ 15ell by Lem. 18 that (PE)) 2 1 As a result, P(E n E 2 1 26.
Bound [7g| under the event S (EE ngg) In this case, we apply (32) from Prop. 14 to get B₁ S and Bg S 1/p2. along with the universal bounds from Lem. 18:
1 G (1 se)1 S 2 2211,1, S 4k 11
(he, 1e)) s 2
|1 o0
4(k - l e+ 1+
2k(k:: 1) ||A||
2=2
E=2
so that by plugging into (50).
k- 1
k+1)
k
T2 S
|1 mk
4n +
+ 2
p²
3p?
and in turn,
k k-1
h: ++1
]²
Ep T (eIS)(e)eE)
S 2 1/ m2kk
4n +
n + 2
2
8.
(51)
p
3p2
Bound [T2] under the event S n Ef n Eg. In this case, we may use that 2 (1log((2/))
+
m log(n + 1)1/p² apply (33) from Prop. 14 get
2 1 p. 2
max B B2) -
KT((M,x|Ixx)
P.vn (2/0) + 2m log(n + 1)
and the bounds based E which give
2mk
G 150)1 S
2 log
2k |1
8
k
2mk
2mk
Σ
(her 15))
S
log
2 - l + 1) 1//l S V 2log
k(k 1) /
8
5
8=2
By plugging into (50),
2m | v log(2mkk/)) logg (2/0) + 2n log(n + +1]
722 S
ki(k: + 1) +
(52)
np,
|1All x
(2 log
(2/8) logg(2/6))
2m log(n + 1)] k(k - 11) 4)
(53)
3mp?
4mk |AA|| [log2 (2/5) + 2m log(n + 111==1(==1)/2
(54)
np?
p. 22 (2mk/5)(k + 1) + (k- 1)(ks +4)
(55)
In turn.
16 |/ m²k² [log2 + m log(n 1))2-1((==1) S
Ep T 1 sx.epnee))
x
n²yp
p- / (olg(22k + 1) + (k- ()kk 4))
(56)
Combining together both (56) and (51) and using that [log. + 2m log(n + 1)] 2 1, we have that
Ep
15
2
2111113 m²
(2gg) m logg( + 1))2-
[
1-1
1: +
2 8 n )
(V2log(2mk//)(( 1) (k: (1):k 4) 2
An +
n 1 2
+
p2
22
the result desired.
31

---

Lemma 23 (Cross term bound). Let I1 and T2 be defined as in (49). For any 6 > 0, assuming that n 22log((20)) + m log(n + 1)]/ we have that
Ep [T]s
2mk 2 | (2g/(22 (ogg(2/00) + 2m log(n + 1))
p²
p. v 2 log (2mk/6) + 1) + (k - ()k-1)(k+ ( + [Anpp (k-1) 1) n++2+ k +
8
3/22
p2
Proof. The following computations are done under the event S. First, apply Prop. 20 to write
TTT2l S
G (h1,))
B1 t (h1 1161 B2 Σ G (he, 1aee
+
vn
mB2 |11l kik - 1)[B1 + B2(k + 1)/3]
(57)
We decompose on the event Efn Note that by Thm. 16 and that 2 log2 +m 1), we have that P(E)) 2 1 - 8. 8. It follows by Hoeffding's inequality and the union bound that
P(E)) 2 5. Similarly, we also have by Hoeffding's inequality that P(Eg) 2 1 - 6. As a result,
P(E n 3 n ES Z 36
Bound [7 under the event ( Ngg ngg). In this case, we apply (32) from Prop. 14 get B1 S and B₂ < 1/p². along with the universal bounds from Lem. 18:
1
[ (h1 $ |l - 4k ||Al
0 (h1,* 1 1x) < 2 ||1
4k ||All
1
0 (he, 1 1,0 2
|1e,
Σ E=2
1(k E + 1) 1/1l/0
24 1) |1
so that by plugging into (57),
k -
1 ++1)
3p )
Ti 2k |//l m
4n +
n 2
and in turn,
12k.2 ||A|| m
4np² + (l - 1) n 2 1
Ep TTIb
8.
(58)
p3
3p?
Bound [Ti2s] under the event S &n n gg ngg. m log(n n + 1)l/p² apply (33) from Prop. 14 get
In this case, we may use that n 2 2|log
-
2 P. KL(Pr,, <
1 /2 log, (2/8) + 2m log(n + 1)
max (B, B2) <
Vn p.
and the bounds based on E n E n Eg which give
(h1.)) < / 22 || C (bi, 100 5
/5)2k ||A]|
2mk
2mk
Σ (hex 150) <
2og
2(k - E 1) |1
S $1
k(k 1) l1
5
5
2
2
32

---

By plugging into (57),
m ||h]] V (lgg(2mk/0) [2gg (2/6) 2m log(n + +1]
72 -
k(k 1)
np,.
m |1hh] (2 log, (2/6) + 2m log(n + 1)] k(k 1)(k +4)
3mp?
mk [1oll loga (2/0) + 2m + 1))² 1))1-1(h=13/2
np?
p- v 2log(2mk/0) 1) (k- + +4]
2mk | 2ogg(2))) (ogg(210) + 2m log(n + 1)]" 11k=1(1=1)/2
[TT x
123/2p²
p. v (8 (k: 1)(k: + +4]
In turn.
2mnk² | v 2 log (2/6) ( + 2m log(n + 1)] 1)k=-((==1/2
Ep 12 (11sstennegne))
x
3/322p
p- log(2mk/8) + + (k - 1)(k + 4)
(59)
Combining together both (59) and (58) and using that [log2(2/8) + 2m log(n + 1)] 2 1, we have that
EP [T]21s]]
2mk |1 / 2 1og(2/0) (2/6) + 2m log(n + 1))=-1(k=1)/2
p
p. 2 log(2mk/ + + (k- 1) 4
+ +)
[mmpp + 1)) n)++
(1)
k+1
32/22
p
the result desired.
We now combine the previous results to prove Thm. 24.
Theorem 24. For sequence of rebalanced distributions (P,, (4)
121. there exists absolute constant
C O such that when n 2 C(logg(21) (2n/p=) p+) ) + m log (n + 1)1/p²
Ep((Pr (K) (h) P(h)) <
CB
+
(60)
n3/22'
where
v log (2n h.m/pp.mm²kk km |1/133
2n logg
p. + m log (n + 1)
2- 2-1(k)
log
2mkn
+
(k: 1)2
p
Proof. We apply the decomposition (47), and subsequently handle the second term using bounds the terms in (48). Set 5 = p²/n". We apply Lem. 22 and Lem. 23 with this choice 8, so that there exists an absolute constants C, Ct, and C2 such that
|111 m² /oo(22n/p.)
Ep [TT21s]] 5 C
[2og2(2npp.) m log(n (k==//2
n3/2p²
(log"
2mnkk k-
P+
p2
|1bll3 m²k4
Ep [T11 S C2
(ogg(2n//)) + m log(n + ()k=) x
n²p2
2mnk
(k - 1)
log
+
p+
p2
33

---

when n 2 (2logg(2//p.. + m log (n + 1)1/p2. This then implies that there is an absolute constant
C3 such that
PR (9)
(h) (m) - p(mi) (n))-P())
Ep
of
< Ep
(P) P(h))² ((h)))114) 1sc
C3 | m²
2n P.
2--((==1)
2mnk
-
1)2
logz
+ log(n 1)
log
+
n2
p+
D
Next, we apply Prop. 21 with the same choice of 8. Because 2[log, (2/6) + m log ( 1)] 2 log(m/0)
and log(1 p+) 2 P+ 2 p., we have that n 2 log(6/m)/ l/ p-), which implies that
m11 p.) R < 5. Combining with the display above, we have that there exists an absolute constant
C' 0 such that
Ep (F (4) (h) P(hi))
of s n
CIIB m²k /log(2//p.))
n²/2p²
1(==1} x [log, m log(n 1)]² 2- 1
(log
2mnk
(k: - 1)2
p. p2
which is the claimed result.
While not shown in the main text, similar techniques to those used above can also control the bias of Pa (k) (h) as in Thm. 25, Interestingly, this bias of order O(n which confirms the intuition that even thought Pn (k) (h) may be biased, the dominant term is the variance.
Theorem 25. For sequence of rebalanced distributions (P()) k21, there exists an absolute constant
C > O such that when n 2 C'logg(2n/ /22pp))>>)/)) + m log (n + 1)
2 s n EP (h) - [h)-P(h)]]
CB
(61)
where B is as defined in Thm. 24.
Proof. First, apply the decomposition (47) so that
Ep Pa (h) - P(t)] P(h)] < EE (h) P(h)) 1s-11
|Ep ((P. P(h)) 1s]l
By using the argument of Prop. 21, we have that
2 1og(2/
Ep (P(hh) P(h) 1s= S 2 2|ll min (2m(1 - p.)",
||/|l,0 2m - p_)
n
Then, by the recursion formula Equation (37), we have that
|E (h) - P(h)) 1s]
(C1 = |Ep [G (h)1s} Ep (1 (G.. Cuh)
vnls Σ Vi (C Cah)
Because G (C1 ... Cah) has zero mean, it follows that
Vn Ep (PPM0) (h) - P(h)) 1s] Ep
+ Ep [1sT2]]
We have by Hoeffding's 's inequality that (P()) > 1 5, and that by Lem. 18 that G (C Cah) s
4k / aoll universally. As result, applying Prop. 21 once again,
Ep [1s G (C ... C,h)]]
Ep Is. I G (C, ... Cahi C...,..h)) + Ep 1s- eg G (C, ... cRAi]| (G...cA))))
4/ //llas min (2m(1 p5) + /2log(2/0)2) |bll 2m - P.)
34

---

Using similar argument to Lem. 22, we have that under SEEE NEg) (which occurs with probability no more than 26).
k-
(n
k+1
[22 S [ mk
4n +
p2 n 2
3p?
and that under S n Ef Eg (which occurs with probability at least 26),
4mnk |el [log, (2/6) + 2m log(n + 1)] (|m
T2 S
nyp
|+ v log(2mk + 1) + (A: - 1)(k + 4)
Applying the decomposition EEp [1 & Ep [Ts(eeene)) T2 + B Ep [ D T2
and setting
5 p achieves the desired result.
n²
D.5 Misspecified Marginal Distributions
We now adapt the main results to cases in which the marginal distributions (Px, Pr) are misspecified. in that the user is provided marginal distributions (Px,cs Py,c) which satisfy the following structure.
Assumption 1. There exist fixed probability mass functions Px and Py for some e E [0. 1).
Px. = (1 P1 + ePx and Py. = Py +
We also have the existence of the positive quantity
p. == minn
Px(x), Px Py (y)) > 0.
Given the existence of p. > we may also define
Pe.c min (min Px,c(x), min Pv)) 2 cp + (1 - c)p
To be precise, the iterations of balancing follow pru Pn and
(x)cxx)
p( (z,z) y)
k odd
p (z, (z.1) v):
P pik- (a)
"
(62)
(v)
pa- (x.
k even
(
We start by deriving a result similar to Prop. 15. Since < 1, the (possibly misspecified) target marginals Px,((w) > and Pv)r(y) > for all x E X' and E ), Define
Σnyy xx
Px p (x) -
M(x, y) P.u- " (x,y)
1 odd
va-n(h) =
mix
(63)
Eny
P( (v) -1 h(x, v) P(L- (x. k even
Py
and
( = Vn (p() (k)
The format of this section will be to derive results analogous to the building blocks the previous section. From that point, the computations from Appx. D.4 will achieve the desired result. For the sake of comparison Thm. we consider error terms containing e only by their dependence on (e, k P....
35

---

D.5.1 Intermediate Results
Proposition 26. Let (P())) be a sequence computed according to (62). Define
c² = =ax { (
These iterations well well-defined under the event S. and for G defined in (43), it holds that
(h) = (h) + vnVz"""h
(64)
and
Vu[Pxe Px](uxh) if odd VilPy. - Py](uyh) if even
(h) = G "(Cxh) + vaVa "(Cch) +
(65)
Furthermore,
(1))
- (Cthi)
+ Vn V (Cah)) +c +cllillare) Vne
cuk nn
(Cahi)
Vn (Cxh)) +0 (vnc).
(66)
Proof. The proof that P (x) > and " (v) > 0 for all x E X and y E y follows the exact same steps as in the proof of Prop. 15. We take this for granted and establish the recursion.
Consider the following steps in the case that k is odd:
Px.e
pi (h)
Σ
h(z y) PIN (I,z)
h(x, 3)
(x) Pik "xx,z)
pix
2,8
3.8
Px,e (1) -
Σ 1 h(r u)
P pik (2, + Σ
1(x,, y) (x,-"((,)))
p(k-1)
2.9
a.
(h) (h) + (h)-
Subtracting P(h) both sides, we have that
| Pl(h) = p P](h) + (*)-"(h).
(67)
n
n
This proves the uncentered recursion formula given (64). We then show the centered version.
1P - Pl(h)
=PP P](Cxh) + (P Pl(xxh)
-P Pl(Cxh) + [Px. - (x)(ux
[ P](Cx)) + (Cxh)(()h)) + [Px,
Pxl(uxh).
Next, we bound the additional error term. By the Cauchy-Schwarz inequality,
[Px.c Px](uxh) <
Px.s Px
||Axh|lb
LP(PX)
Vx2(Px.e -11xx)(lxxll
v xxx 1Px) |||ll
as yx is orthogonal projection in L2(P). Using convexity of f-divergences, we have that S Px)
x2 (
xx (Px
This achieves the desired result.
Using similar ideas, then prove analog of Lem. 19.
36

---

Lemma 27. For k odd, it holds that
Px,
VnVa--"CCa (Cxh) = Σ
(xz) - - "Cxx 1zk),
pik 11
n.X
whereas for k even, it holds that
Py (1z)
VnVa
(Cyh) = Σ
-1 G/--" (Cyh11))
pit
2=1
n,Y
Proof. We give the proof for lt odd. We claim that we need only show that P(Cxh1ju 5k) = 0. This would show that
Px. (z) -1 (2,)((2))))
VnVa
VnVo-- (Cxh)=
Σ
p1 1) n.X
2.y
Px. (I) -1
(x,
n
p(k-1) n,X
1=1 J.y
Px, (Iz) -1) [xxh u)
n.X
Px.e
(xz) -1) vn Σ Cx (x
pik-1 n,X
s.y
Px.e (xy) -1)
(
n.X
where p(Cxh1,k) 0 is employed in the last step. Now the result follows from (65) in Prop. 26 and the definition of G (h). To prove the claim, as Lem. 19, write
E [Cxh1 = (0 [E|cxh|X]
= i
But E [Cxh)X] = by definition of Cx. Taking an expectation over Px gives that P(Cxh1,;) which implies the desired result. The proof for ls even follows symmetrically. A
For the remainder of the argument, we see that (66) can be unrolled so that
G (
(h)
< G (C (C.... Cah) + 7
Vie-" (Ce...Caali Cali) +
O(kvne)
(68)
lirst-order term
misspeciication
higher-order
where use that G ( Next, we need to bound n
Cch) in particular accounting for the marginal violation term. We follow similar steps as in the analysis of the higher-order term in Appx. D.3.
Proposition 28. Assume that Pn,x(1t (x) > for all E X. It holds that
Px.e(x) p(a-11 n. (1r)
(max - 1 max (1/p/1//220 - 1,1}
if = 1 if > 1.
max SEA
<
(69)
In addition, we have that
Px)
log T-c
+ +n 2 KL(PxxxIIxx)
ifk=1
max zEX
-1
<
p(k-1)
n,X
(x)
0
log I=c
+ पे. (T.
37

---

E..
Ryg
Moreover, when (I,(Pn,
and c 5 1-1-00 -xx (-*+)
we have
(x))
00 <2(21. 1
log 10g 1-c
2
KL
(70)
2EA P (x)
Daie
Pa,e
Proof. First, observe that p(o) (2) = - (x) 1 1/n under the event S. For k > such that is
odd, we have that for x E X.
Pr, ( v)
Pixx" (2) = Σ
Pa-1) (r. = Σ
pik 2) (y)
yEY
NEY n,Y
P.,e
Pik (x,,y) p. Pik 2: (x) p. (x, P2. n.X
The result for i even can be proven similarly. We now prove the inequalities listed in the statement using on the lower bounds above. Proving the first inequality. For any I X.
(x) - = max { (x.( 1,1
(x( } pik (x)
max(n - 1,13 if max 1,1] i
pik (x)
n (2)
which is the desired result.
Proving the second and third inequalities. Consider an odd ke 2 1. By the definition of total variation distance, holds that
Px.((x) " S TVV(P n. Px.e). Px.)).
max xEX
According to Pinsker's inequality, we have that TV(P n, Px.x) S v KT. (PTtt))))) PX,))) and so have that
(x) x
- (x) V KL(PEx""||xx.) V (X
max ZEA
where the last inequality follows by the monotonicity of Sinkhorn iterations given in Prop. 13. Notice that the remaining term is (L(P xX = KL(Pm.xI which may not decay to zero as + co. Because < 1, write
Pa,x (x)
Σ (x) log
(1 Px cP
FEX
Pa,x
Pr,x (I) log
Px
TEX
1
= =KL(P + log
1-c
1
E V/2 (T.((...(1Px)) + V/z log 1-c
We can then apply the lower bounds
Px.((x) -1 s n,X "x)
(L(P + log T-c
if
max x
p
Pe.. ( KL + V log 1-c
1
if
Finally, combining the arguments above, we have that
max Px.((x) - p (x)) S KL[Puxx1IPX
1
V
log 1-6
zEX
2
P_,
P.,
Pr,e
4
4
2
38

---

where the last step invoked the assumption that
P..
and c exp (-1a²) )
[I,(P,,||Ixx) v
8
This means that
p(k min
(z) 2 min (x) max
p (2) Pxxc(2))
2
Da.e
aEX n.
ZEX
#Ed
2
Hence,
a, (I) Px)
(x
maxsEX
2
KL ((P.x)..
pin- X,X (x)
millge P,,x"(a (x)
P..
Now, for i even, set l = 2t for t 2 O. We have that
pra n (v) - Pr,a(v) & TV(Per n. ", PY,) S V KL(Pr.ell
pray
n.y 1))
max yEy
D.Y
Invoke Prop. 13 once again achieve
1 (L(P S s / which completes the proof.
KT. (T.(l n.
+ V log
1
Proceeding with similar steps, define the quantities
(x))
maxzex
odd
/ (x)
B = Mi and B₂ = max 25E5k
Mr
for Mr ==
By)
maX
A,Y " (M)
l even
We must now establish an analog of Prop. 20. Proposition 29, For any l 2 1, the following holds under the event S:
B1 G
B2
G (bex
+ B2 | Vnk(A: 1)[ BB 1)/3].
Proof. This proof largely follows the argument of Prop. 20, while accounting for the misspecified marginal error. Using again the notation he, == C₁. C₂h, it follows from Lem. 27 that, for odd ł
m
(z,) 1 pie-
Px,
m
(nvie-((he)) (he,s) = =Σ
CK Cec--"he.a (he, S Me
G
"he, 1,0)
j=1
j=1
The bound above holds for even as well. Then, using (64) from Prop. 26 along with the triangle inequality, we have that for f 2 2,
| - Pl s pa-- Pl(hex + |/ (hex
which implies that
G (he,
(71)
G "her (e))
+ vn Vee ((hexe)5))
6:00 (he,
+ V v ( + +... + vni (
(he,
+ +MVnn (Ihe, 15e) +... Me p " 15
(he
+ 2 21ll vn Bi + Bz(e 1)] (k l+
(72)
39

---

by Lem. 18 and M1 I B1 and Mt I B2 for 2 2. The bound above holds trivially for & = 1. Summing these bounds over and j, we have that
v (hea))
m
< M1
G (h1
Σ
Me 2
G "
P(eex
j=1
8=2
j=1
m
m Σ>
B1
G [11 16) + B2
G " "(hexlse)
2=2
B,
G (
/=1
B2
b ( (hen + +2111blb v vi + Ball - 1)] (k - (k-++1) e+1) apply (72)
2=
B1 (1i, se) | + B2 Σ
1 (he, (ell
+ 2mB2 |1A1l00 Vu Σ
B1 + B2 1)] - -
E=2
because | = m. We sum up the last term:
A-1
(k: - 8) + B2 k-1
Σ B Boll 11- 11] (k (k- 1) B1 Σ
e -
t=1
E=1
k(k 1)
[B + Be(k + +1/3
2
which completes the proof.
D.5.2 Mean Squared Error Bound
Ultimately, we wish to construct an upper bound for
2 Is + Ep (Phn(t) ((h))) 1s
Ep
(73)
as the method returns P. (h) when S is not satisfied. The first term will be controlled by intermediate tools developed above. The second term that includes S no different from the one analyzed in Prop. 21. We handle the second term first. Recall from Prop. 21 that for any 8 (0,
Ep (Px(M)) - P(h))2 1s:
4 || min (2m(1 p..)",. + log((/0) |1bl13 2m - pp)"
(74)
Repeat the argument from the proof of Thm. 24: because 2(10g, + m log(n + 1)] 2 log(m/6)
and log(1 (1 p.) p. 2 p.. have that
2 lo, (2/8) + moog (n + 1)1/p² = n log(5/m)) log(1 - p-).
(75)
This in turn implies that m(1 p.)" < 6, and gives condition on the sample size n. Further in the analysis, we will set 5 = (1). so right-hand side of (74) can then be upper bounded further, resulting in
Ep (h) P(h)) 1s < <411112 (2+ log(2/0) =
pd n
40

---

a higher-order term compared to other components of the bound.
Next, must control the left-hand side of (73). We perform the decomposition based on (68):
Ep ( (h) - P(h)) s
s Ep [s] + 2Rp T1 1ss + Ep [2
(76)
+ +K(vE) Ep [1zil+
[Tal) [
1s + O(kpe)
(77)
for
k
T = [P - PI.CI...... C.h) and T2 =
Σ vre (Ce...Ce
(78)
Recall the events & and E and Eg from Appx. D.4. To perform this computation efficiently, we will split the bounds on each term into two components. In particular, we will show that
Under the event 8 & n 82
< T2 + E2.
Under the event (E ng Tg
S T2 + Es
Under the event S &ng: [7] S Ti.
Under the event S [7i| $ TT.
where any term denoted with "E" will represent all error terms that include and will be written in big-O notation. There are no errors for the bounds on Ti, as this term does not depend on the misspecified marginals. The idea is that for the "T2" terms we may reuse the bounds derived in Appx. D.4 by simply replacing analogous terms from Appx. D.4 depend on p. only through Prop. 14; similarly, the corresponding terms in this section depend on P.,. through Prop. 28. We return to the terms in (76) and (77).
with Pr,e- This is due to the fact that the dependence of the
Decomposing on Eg will result a bound of the form
Olkvee= Ep [[s] S O(kVE) (OT + Ti).
Decomposing on E n Eg will result in a bounds of the form
Ey T 1s
s 220(T)) + T/2 + (6
+ + + (E2 + E2T2))
O(kvC) Ep
S OO(k/) (6 ES) Ta F22)
s
Finally, decomposing on Ef n Eg n Eg will result in a bound of the form
Ep
1
1s < 30T TiT + 0 + E2)).
The leading terms 25 + T and 35T + TiT2 from both bounds should have the exact same form as the terms in Lem. 22 and Lem. 23, with p. replaced by p.,e. thus retaining the same dependence on (n, k). By setting 8 = - n we will achieve a similar result Thm. 24, i.e. that
(1) - P(N) 1s
Ep
2
RG
+Ō
13/22
n
δ (fnna/n) (ES T5 Eg (E + T2 (p..// (p...n)'TTT Es TrE2))
(79)
+ Ô (l:v (D).e/n) 'Ti + Ti + (
/
(T
+
E2) - -k2))
(80)
remains to quantify the Ô terms by computing the order of the constants (T2, F2, Tí, Feg, Tı, T.)
We follow similar steps to Lem. 22 and Lem. 23 achieve this.
41

---

Lemma 30. For 8 (Da,e/n),, assume that n
(110g)22/6) + mlog(( + 1)1/p. and 2
(- P..
- 1eexp
Then, holds that
Too 0
n2 b/
))
E2 =0
Pi,e
T2 Ô
E2
4
(V log T-e log -
Pi.c
Proof. The following computations are done under the event S. First, apply Prop. 29 to write
h
vn
22
-
G (h1,a 15t) + B₂
(0) (he, 15t
j=1
6=
+ mB | k(k - 1)[Bi + B2 + 1)/3}.
(81)
We decompose on the event E/ Eq.
Bound [Tb] under the event S (E ngg) In this case, we apply (69) from Prop. 28 get B1 S and Bg S 1/f2... along with the universal bounds from Lem. 18:
1 g ( 15e)1 S 2 1/1, S 4k //
1
(he, 100) 2
l llo
4(k - e 1) 1)
2k(k: 1) ||All
8=2
2=2
E=2
so that by plugging into (81).
k 1
k+1
S
|h
mk
4n +
n n
2+
p7.
ES
T
Bound [T2l under the event S S 8f Eg. In this case, we may use that n 2 8/p2.. (because (208) + m log(n + 1)] 2 for (0, 1)) and apply (70) from Prop. 28 to get
max B1, B2 0 V log T +
2 P..c
llog (2/0) 2m log(n + 1)
2n
The bounds based on Ed give
2mk
( 1,)) S
2 log
2k ||h
2mk
2mk
C 15r) <
log
2(k -
V 2l
k(k 1) |
8
8
2=2
By plugging into (81), can reuse the steps in the bound from (55) (for all terms without e) to write
4mk |/bll [log 2m log(n + -1
Fg s
np2
|p.. / 2og (2mk/5)(k + 1) + (k- 1)(k + 4)] + Ez.
so that
4mk |11||1 [log2 (208)(2/)) + 2m log(n 11) To
np?..
Pee 2 2og (2mk + 1) + (k - 1) + +1]
x
42

---

We compute Eg by using that
Bi. B2 so
+ +0(
max
G (11, 1() s 0 (k)
Gn (he, 1,t)l - 0 (k?))
2=
which gives
E2 = 0 Pi.e (V log - + log )
We now make the corresponding argument for the term T1.
Lemma 31. For (fn,//)'' holds that
k Vn
T Ö(k), Ti =
Proof. The following computations done under the event S.
Bound (Til under the event SY. Here we simply apply universal bound on the empirical
process term:
G th (h1,))
2 |1 4/: |Alloo
vn
so that 7f =F=4k]o [4/l
Bound [Ti] under the event Sn
Now, we may use the definition the event &g to achieve
2 1lg((2/)) 2k |Abllc = T1.
G 0 (h1,))1 S
Knowing that E3 = 0, we simplify (79) and (79) to read
ò B2
(E2 (E2
T2 Ti)
Ô (KVE ((D../ + TT (D.e//(/)) + E2) k2e).
We combine the bounds from the previous two lemmas compute (79) and (80) state the main result.
Theorem 32. Let Asm. I be true with error E [0, 1). For a sequence of rebalanced distribu-
tions (P(N))21, there exists an absolute constant C > 0 such that when n 2 C'lloga (2m pp..)
m log (n 1)]/ min Pr,c)² we have that
2 Is Ep ( (()(()) - P(h))
+ Er (Pa(M) - P(h))2 1s
o S
46
+
3/2
n
k4
1
42
+ δ
log
+ log
log
+
log
+ Vn
Ph.c
1
p3.0
7
+
2
Ve
Px.c
c4
p2,c
n +
222
2
log
+ log
- )
+ c
n
Vn
4
Pr.
2
1
43

---

E
Experimental Details
We provide the full experimental details of the experimental results from Sec. 4. We report additional evaluations on downstream tasks with linear probing and zero-shot retrieval. Finally, we give illustrations of the sensitivity to misspecified marginals, and of the convergence to the given marginals.
E.1 Datasets
Pre-Trane Tre-Trainng Data. The pre-training data was taken from the public SigeNeet-Captionn dataset [Fang et al., 2013]. We subset the dataset by selecting the 250 classes that were most frequent in the dataset, resulting in 174,594 images and associated Flickr captions. The exact images used and their associated captions are given in the code supplement.
Evaluation Data. We perform zero-shot classification described in Sec. 4), zero-shot retrieval, and linear probing with various image classification and image-caption datasets. We used the default class captions (for classification) and default linear probing parameters from the CLIP Benchmark repo. The datasets (test splits) used were:
CIFAR-10: 10,000 colored natural images labeled with one of 10 classes.
CIFAR-100: 10,000 colored natural images labeled with one of 100 classes.
STL-10: 80,000 colored natural images labeled with one of 10 classes.
MS-COCO: 41,000 colored natural images with associated captions.
Flickr8k: 8,000 colored natural images with associated captions.
Rendered SST2: 1,821 images of typed natural language with sentiment label (2 classes).
VOC2007: 4,952 colored natural images labeled with one of 20 classes.
FGVC Aircraft: 34,000 colored natural images labeled with one of 102 classes.
Evaluation scripts using the various embedding models (described below) are provided.
E.2 Model Specification and Hyperparameters
Architecture and Implementation. The models considered CLIP models [Radford et al., 2021], and are specified by pairs of encoders (fe. ga). representing images and text, respectively. The encoders decompose into fe = fread o ffase (similarly for ge) where fhuse denotes base image encoder and chead denotes a trainable head model. The head models are feed-forward networks with two hidden layers, 256 hidden units, and 28-dimensional output representations. Their input dimensions may be 512 or 768, depending on whether CLIP model or BERT/GPT-2 model is used as the base. For the image base/foundation models, we use the open-source OpenCLIP implementation model with the la laion2b s34b b7 9k model tag. For the text encoder, we use the encoder of the variant of the ViT-B/32 with tag dat a ccomp x1 s13b b90k. For the other text encoders the Huggingface implementations of GPT-2 and BERT were used.
the ViT-B/32
Optimizer. For optimization, models were trained with stochastic gradient descent (SGD) with the learning rate tuned along the grid {1-3 3 3-3 3,, 1,,1-,3- 1_1 and fixed weight decay parameter of 0.01. Momentum-variantt such as Adam [Kingma and Ba, 2015] were used to isolate the effect of varying losses as described in Sec. 4.
E.3 Compute Environment
Experiments were run on a CPU/GPU workstation with virtual cores, 126G memory, and four NVIDIA TITAN Xp GPUs with 12G memory each. The code was written in Python 3 and we use Py PyTorch for automatic differentiation. The OpenCIIP and CLIP Benchmark repositories were used
for zero-shot evaluation.
44

---

E.4 CLIP and Multi-CLIP
We considered the contrastive learning example from Sec. see (6) in particular variant of the CLIP objective in which either zero, or one, or more than one balancing iterations are performed (see (6)), via optimizing
1 2 2=1
L
[log (X,, Yi) + log R (X,, Xi)]
(82)
This contrasts the single-iteration variant LP which in fact reduces to the original CLIP loss. Because these iterations applied in the objective, backpropagation occurs through each iteration.
In Fig. 3, we plot the zero-shot classification performance (in terms of average per-class recall) of the variants trained on _ (the normalized initial measure, No Balancing), L (the original CLIP loss, CLIP balancing), and L (the two-iteration CLIP loss, Multi- CLiti balancing). We also vary the quality of the text encoder for. observing an overall accuracy trend of GPT-2 < BERT < CLIP across variants, which is to be expected given the base representation quality of each model. Interestingly, there is an improvement stemming from performing multiple balancing iterations across choices of the text embedding, the batch size m, and the evaluation dataset.
E.5 Metadata Curation
We considered in the metadata curation example from Sec. 2 how to use balancing to adjust the entire pre-training set, in the spirit of Xu al. [2024]. The target marginal Py is selected by choosing threshold for which frequent keywords have their probability mass truncated, and the probability measure is normalized sum to one. In Fig. 4, we show the observed marginal Pa, and the target marginal Py sorted in increasing order (left). The original marginal on y has approximately 5 orders of magnitude of difference between the most and least probable keyword. After balancing, the target marginal has less than 2 orders of difference. To see how this affects downstream performance, we plot the zero-shot classification accuracy over training iterations in Fig. 4 (right) when using the original dataset (orange) and using the metadata-balance dataset (blue). We observe moderate improvement especially in the small batch regime (m 512) when curating the dataset.
E.6 Additional Experiments
In this section, we provide 1) a synthetic data example that helps elucidate the role of the spectral decomposition introduced in Sec. 3, and 2) additional evaluations on downstream tasks such as zero-shot retrieval and linear probing. For the latter, we maintain the experimental settings as used in the zero-shot classification example from Sec. 4 (Fig. 3). That is. we train variants of CLIP models (see Sec. 2) on the ImageNet-Caption dataset [Fang et al., 2013]. As before, we use a fixed image/text encoder as base vector representation and compose it with trainable ferw-forward neural network, i.e., fe = phead o frase, for = 0, (images) or 8 = Or (text). For the base text embeddings, we maintain three levels of model quality: GPT-2 [Radford al., 2019], BERT [Devlin et al., 2019], and CLIP-based encodings.
Baseline Comparisons. We present synthetic data example to understand the role of the singular values ..,..., Sm and compare our approach to simple baselines that make use of (Px, Py). We also consider misspecification of these target marginals, in that they are chosen by the user but are not the marginal distributions of the tata-generattig distribution P. First, while one can verify by hand that (14) is distribution for which 82 s, construct more general example for We leave the construction to the end of this example. For controllable misspecification, we define E [0, 0.5] to be the misspecification level, so that the corrupted target marginals are set to be
2.
Px == (1 Px cPx and Py == (1 - Py cPy,
(83)
where Px and Py are drawn independently and randomly from the Dirichlet distribution (i.e. uniformly over the probability simplex on m atoms). Finally, other than the empirical measure Pa. we define one additional baseline; the importance weighted independently (IPWI) estimator is defined as
Pv(v) Pn. (x) Pa,Y (v)
Px(x))
IIPWI (x.1)
(r.
(84)
45

---

Marginal Fitting Techniques under Misspecifiation
10
0.25 (IPW)
Mean Squared Error
(IPWI)
(PWI)
0.25 (Bal.)
10
0.125 (Bal.)
<= (Bal.)
Emp. Measure
10
0.0
0.2
0.4
0.8
0.0
Leading Singular Value s
Figure Baseline Comparisons across Dependence and Misspecification Levels. Each line refers to combination of an estimation method (the empirical probability measure Pa, the estimator PIPWI from (84), or the balancing estimator P(a) for k 8) and noise level on the provided marginals (see (83)). The y-axis shows the mean squared error of estimating a linear functional. The r-axis represents the dependence level = 82 (i.e. the leading singular value other than 81 = 1).
This estimator simply reweighs all cells of the empirical measure by the likelihood ratio from each observed marginal to the target marginal. Note that the result may not even be a probability measure, may not sum one. Observe the comparative performance in (see Fig. 5). We notice in particular that the naive pIPWI is outperformed by empirical measure uniformly over s, as by applying both n reweightings simultaneously, the estimator does satisfy either marginal constraint. On the other hand, under the maximum amount of target marginal corruption (c 0.5), the balancing- based estimator suffers an approximately half-order of magnitude in MSE. When s = the MSE of the balancing estimator decreases significantly. We hypothesize that this is because the data sources X and Y nearly a function of one another, and if this function estimable to high precision by small amount of data, then single marginal can identify the entire joint distribution via pushforward calculations. That being said, on so it is difficult to control the singular values without controlling the respective bases.
important to note that the quantities v, and vi in (15) also depend
As for the construction of the probability mass function and test function, let Im and 1mxm denote the identity matrix and matrix of ones in mmxxm For any E (0, and m 2 1, consider the probability mass matrix P given by
1 1mxm + Im
1 1mxm
The eigenvalues of the first matrix in the squared brackets are O,,0,.... 0), it is rank 1 matrix for which 1m is an eigenvector. The second matrix in the square brackets is the centering matrix (the projection matrix that subtracts the mean of vector's 'ss components from the entire vector). Multiplied by s, has eigenvalues (0, s. s.....)) where is associated to the eigenvector 1mm Thus, the matrix in its entirety has eigenvalues (1/m, 8 s/m,-. ... s/m) s/m), where the scaling factor ensures that P sums one. The relation (13) holds for this choice of P and uniform marginals, with 82 = = 8m 8. Thus, by tuning s, may control the level of dependence between X and Y. Finally, because X' and m X m table indexed by i (meaning x;) and (meaning yj). We let h(1,, f/) Zupl where the Zij are independently drawn from standard normal distribution. The resulting mean squared error is estimated with 200 seeds at = 300.
y are finite, we can also specify the test function l via an
Zero-Shot Retrieval. In this evaluation, we assess the ability of the learned representations to match queries from modality their counterparts in another modality. We given test sets X { x xM) J of images and Diest = ( yN } of texts in natural language. We are also given matrix of annotations A {0, Mx where Ag 1 if and only if yj a "relevant" caption
46

---

Zeroshot Retrieval Performance (Recall @ 5)
CLIP Text Embeddings
BERT Text Embeddings
GPT-2 Text Embeddings
MSCOCO (Text) MSCOCO (Image)
0.3
0.100
0.00175
0.075
0.2
0.00150
0.050
0.00125
0.025
0.00100
0.0
0.000
0.008
0.3
0.006
0.2
0.1
0.002
0.0
Flickr-Bh 8k (Image)
0.15
0.0065
0.0060
0.0065
0.2
0.05
0.0050
0.0045
0.0
0.00
Flickr- Bickr-8k 8kk (Text)
0.20
0.020
0.4
0.15
0.015
0.10
0.05
0.005
0.0
0.00
Training Steps
Training Steps
Training Steps
No Balancing
CLIP Balancing (k (k= 1)
Multi-CLIP (k = 2)
Figure 6: Zero-Shot Retrieval Performance across Embeddings and Objectives. The three vertical panels describe different choices of the text encoder fe,, which increases in quality from left to right; that is, pre-trained GPT-2, BERT, and CLIP embeddings, respectively. Rows indicate various datasets, either MS-COCO Flickråk. evaluated under recall at K 5 for image and text retrieval, respectively. The y-axis of each plot indicates the metric (see (85)) for either image or text retrieval, whereas the z-axis indicates training iterations batch size 512.
47

---

CLIP Text Embeddings
BERT Text Embeddings
GPT-2 Text Embeddings
Rendered SST-2
0.62
0.60
0.8
VOC2007
0.6
0.4
0.2
0.4
FGVC Aircraft
0.3
0.2
Training Steps
Training Steps
Training Steps
No Balancing
CLIP Balancing (k = 1)
Multi-CLIP (k = 2)
Figure 7: Linear Probe Performance across Embeddings and Objectives. The three vertical panels describe different choices of the encoder fe, which increases in quality from left right; that is, pre-trained GPT-2, BERT. and CLIP embeddings, respectively. Rows indicate various evaluation datasets from Rendered SST2, VOC2007, and FGVC Aircraft. The y-axis of each plot indicates the average per-class recall, whereas the r-axis indicates training iterations at batch size 512.
for image x; (and vice versa). Given particular query y E Viett we define the top- K neighborhood of
Nv: = arg max
Σ
SC SC[M) iES
iee. the images in the test set that have the closest embeddings under the given model. Then, may define the average recall at K for image retrieval metric as
EiENNtHROn Aij Eve(M Arj
AverageRecall (0)
Σ j=1
(85)
In words, the metric evaluates the retrieval system's 's ability to detect relevant items the dataset, in this case by comparing the closeness of the image-text representations. We can analogously define the average recall at for text retrieval metric by swapping the role of x and y above. The results for both retrieval metrics the MS-COCO [Lin et al., 2015] and Flickr8k [Hodosh al., 2013] benchmarks are given in Fig. 6.
Linear Probe. Here, we evaluate the quality of the model's encoders by fine-tuning single linear layer top of the learned representations for a classification task. In the case of linear probing via image classification, we use only the image encoder fo,- We are given a training set {(11, ei)..... (Ix, ex)} of image-label pairs, where each {l...., C). We fix the model
48

---

2
P.Y
R
R
y
Figure 8: Empirical Marginals of CLIP Contrast Matrix. Depiction of the probability measures @ and R as described in (82) from Sec. 2. The orange bars correspond to the observed marginal after fitting the target uniform distribution neither marginal is set to uniform. Center: Q and R which corresponds to the original CLIP loss. Right: Q and ? which correspond two iterations of the balancing procedure within the loss. The blue bars are slightly non non-uniform.
the given iteration. Left: O and Rg where
parameter fy and solve the regularized multinomial cross entropy (MCE) objective
min WERCxP
[WC = N Σ LogoottmaxxWWW fo, (2:))le, + |
1
A
where À > 0 is regularization parameter, ||-||r denotes the Frobenius norm on Cxrr and LogSoftmax CRC IRC is given by TogSoftmax(2)
log E exp(2z). This results
in classifier
g(x) = arg max Sa, (z))l.-
jE(o]
which can then be evaluated using standard accuracy metrics on held-out test set. The image classification results for the Rendered SST2 [Radford et al., 2021], VOC2007 [Everingham et al., 2007], and FGVC Aircraft (Maji al., 2013] benchmarks are given in Fig.
Empirical Marginals in CLIP Balancing. To further clarify how the iterative balancing procedure baked into the CLIP losses, recall from (82) that objectives decompose into two terms, which depend on Q00 and R800 which differ only based on whether balancing to fit Py or to fit Px is applied first, respectively. Thus, for any model parameterized by and any number of iterations there are four marginal distributions of interest: Q x Q R and Rgy.- Based the order of iterations, we have that Q R = Py, and Rg Q Px. This illustrated in Fig. We see that after only a few iterations, both marginal distributions converge to the uniform distribution.
49

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Theoretical claims, the focus of this paper, are supported with proofs.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We clarify that the setting studied has some dissimilarities with practice in
Sec.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
50

---

Justification:
This is done for all theoretical statements.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: Yes]
Justification: Code is provided to reproduce the main results and an extensive experimental details section is written.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
51

---

Answer:
[Yes]
Justification:
This is written in the public repo provided in Sec. 4.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https: /ttpss//nnis... cc/
pub /ublic//i guides PodeSubmli ssionP oelPolioo olicy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (ht tps: / /nnis.... cc/ publi iici//// //guu/guuu gugguuooo dees/d/// CCde/ooduuaaa Sullesuuiiii as i onPol icy) for more details. cc/pcccllccccccccccccpp
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, we even give list of the specific images of ImageNet used to train the multimodal models.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: Our evaluation metrics are shown with all seeds and their mean plotted in the corresponding figures.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by eπror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (c.g., Normally distributed errors).
52

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Please see Appx.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /tw//nnursp...// cc/ wwt.:///ww.ccccccccccsuuuuu
Answer: [Yes]
Justification: NA
Guidelines:
The answer NA means that the authors have reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
The work is primarily theoretical and retrospective.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
53

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [No]
Justification:
We do not release
models, only small-scale illustrative ones.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We site all software and models used in the study in Appx. E.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name the license (c.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, papers w thcode com da set. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
54

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide list of notebooks and scripts to illustrate our method and connect it to existing software repositories such as OpenCLIP.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
55


Reviews:
[{"review_id": "23pFttbvDD", "reviewer": "Reviewer_G8ZX", "summary": "This work focusses on data balancing strategies in context of self-supervised learning. The main claim of the paper is that data balancing, commonly used to avoid representation collapse, has a variance reduction effect. The authors introduce an upper bound on the MSE of a balancing estimator, relating it to empirical risk minimisation. The main paper covers the key elements of the proofs, which is given in detail (and is extensive) in the appendix. Experiments are conducted to illustrate the impact of data balancing on examples described in the paper.", "questions": "I have the following questions for the authors:\n\n- Line 54: What do you mean by \"X and Y are forms of the data that are related to, but distinct from, the form of Z\"\" given that Z is equal to (X,Y)?\n- p4, example 1: What would the target marginals correspond to? What is \\psi_n^(k) and P_n^(k) here? \n- p4, example 2: What would the target marginals correspond to? What is \\psi_n^(k) and P_n^(k) here?\n- Why do we need \\tilde{\\psi}_n^{(k)} in (12) and how does it relate to \\psi_n^{(k)}?\n- How does (15) relate to the clip example introduced earlier in the paper and why is this a valid and sensible simplification to study?\n- Does the main result have implications in practice in terms of design of algorithm?", "limitations": "The authors provided a brief discussion about future work at the end of the paper. The work presented here is theoretical in nature, attempting to provide new insights in existing approaches. There are not immediate implications in practice, but authors could discuss the scope of their work in more detail.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "This paper attempts to shed light on SSL training and the role of data balancing. The paper formalises the problem and develops extensive theory. The main results is pretty cool and insightful in the sense that the upper bound on the MSE shows that data balancing has a variance reduction effect. The topic is of interest to the community and the work is focussing on a poorly understood paradigm that is becoming dominant.", "weaknesses": "I have three main concerns with this work:\n\n1/ The theory is *very* extensive. The Appendix contains several pages of proofs that are difficult to parse and come on top of the formalism presented in the main paper. It seems like the main body could be simplified and made more to the point to convey the main gist of the contribution and make it more accessible.\n\n2/ It is unclear how the data balancing examples in Section 2 map to the formalism introduced in Section 3. For example, what would (4) look like for example 1 and example 2?\n\n3/ It is unclear what the experiments bring to the table and how they provide evidence to the main result. Making the link more explicit and explaining what are the key take aways from these results would help the reader.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and questions. We address them below.\n\n>**The theory is very extensive. The Appendix contains several pages of proofs that are difficult to parse and come on top of the formalism presented in the main paper.**\n\nWe provide the complete, self-contained, proofs of all of our theoretical results, for reproducibility. The main analytical appendices (B, C, and D) are split by topic and include an outline for readability. As per the NeurIPS Paper Checklist, \"For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?... The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.\" This sketch is given between lines 235 and 259. We are happy to incorporate any suggestions to make them more clear.\n\n>**It is unclear how the data balancing examples in Section 2 map to the formalism introduced in Section 3. For example, what would (4) look like for example 1 and example 2?**\n\nThe connection between the examples and the analysis in Section 3 is first introduced in the introduction, between lines 59 and 81. The main notational difference is that quantities are often indexed by $\\theta$ in Section 2 and 4 as they describe objectives that are evaluated at a specific parameter value. In Section 3, we drop this index as we are considering the statistical properties of balancing estimators (i.e. the loss that is actually optimized in an SSL scenario) at a fixed parameter value. Thus, the analogous value to $\\psi\\_n^{(k)}$ in (4) is always $\\sum_{x, y} h_\\theta(x, y) R_\\theta^{(k)}(x, y)$, which is another way of writing (2).  \n\n>**How does (15) relate to the clip example introduced earlier in the paper and why is this a valid and sensible simplification to study? Does the main result have implications in practice in terms of design of algorithm?**\n\nThe objective in Eq (15) is not a simplification but a generalization of the CLIP objective. It exactly reduces to the CLIP objective when $k = 1$. As seen in experiments, zero-shot accuracy of learned representations can improve when considering more iterations (see Figure 2 of the manuscript and Figure 6 of the attached PDF). The main result is meant to provide theoretical insight into the reasons behind the success of SSL training procedures that are widely popular, but not well-understood (as outlined between lines 23 and 48).  \n\n>**What do you mean by \"$X$ and $Y$ are forms of the data that are related to, but distinct from, the form of $Z$\" given that $Z$ is equal to $(X,Y)$?**\n\n$Z$ is not necessarily equal to $(X, Y)$, except in empirical risk minimization (line 53). We use the examples from Section 2; in the case of self-labeling, $Z$ represents an image, $X=Z$, and $Y$ is a learnable cluster representation (Example 1). For contrastive learning, similarly to empirical risk minimization, we have $Z = (X, Y)$ where $X$ is an image and $Y$ is a caption (Example 2). In Appendix E.4, we have yet another example of metadata curation, in which $Z$ is an image-caption pair, $X = Z$, and $Y$ is an associated keyword.\n\n>**example 1: What would the target marginals correspond to? What is $\\psi_n^{(k)}$ and $P_n^{(k)}$ here?**\n\nThe target marginals are discrete uniform measures on $\\mathcal{X}$ and $\\mathcal{Y}$ (line 148). $P_n^{(k)}$ would be analogous to $R_\\theta^{(k)}$ (line 132), and $\\psi_n^{(k)}$ is (throughout) analogous to (2). We say analogous because the starting measure $R_\\theta^{(0)}$ comes from the output of a model parametrized by $\\theta$, whereas $P_n^{(0)} = P_n$ is used in Section 3 to represent the empirical measure of randomly drawn data.\n\n>**example 2: What would the target marginals correspond to? What is $\\psi_n^{(k)}$ and $P_n^{(k)}$ here?**\n\nThe target marginals are discrete uniform measures on $\\mathcal{X}$ and $\\mathcal{Y}$ (line 167). The other quantities have the same interpretation as in the previous example and others.\n\n>**Why do we need $\\tilde{\\psi}_n^{(k)}$ in (12) and how does it relate to $\\psi_n^{(k)}$?**\n\nBecause balancing iterations are only well-defined under the event $\\mathcal{S}$ (line 192), $\\tilde{\\psi}_n^{(k)}$ is defined simply to handle the technical consideration of $\\mathcal{S}$ not being satisfied.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SkHeONy4Ex", "reviewer": "Reviewer_os3n", "summary": "This paper explores the use of data balancing in various self-supervised learning (SSL) frameworks. The authors argue that this iterative algorithm, which is typically used to avoid representation collapse in SSL models, also provides a benefit of reducing the variance of empirical functionals of the distribution over data sources. The paper establishes non-asymptotic bounds quantifying this variance reduction and relates them to the eigendecays of specific Markov operators.", "questions": "1. Can the findings be extended to other areas?\n2. Can the authors provide more details on how the assumptions, such as the spectral gap condition, hold or when will they not hold in practical scenarios? Are there specific types of data or models where these assumptions are more likely to be satisfied?\n3. For Figure 2, Can the authors provide experiment results with more iterations?\n4. Are there any other variance reduction techniques and how does data balancing compare to those techniques?", "limitations": "In checklist the author mentions that the work is primarily theoretical and has no societal impact. but it will be better to discuss the positive social impact. The author also mentions the limitation that the setting studied has some dissimilarities with practice but haven't addressed the limitations yet.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. the paper provides a new perspective on the benefits of data balancing\n2. provide different examples of data balancing in practice and prove a non-asymptotic bound on the MSE of balanced estimators.\n3. The findings may have implications for improving SSL models", "weaknesses": "1.the experiments are somewhat limited in scope\n2.adding more visualizations or intuitive explanations may be better for understanding the key finding of the paper.\n3.will the assumptions limit the applicability of the findings?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your helpful comments and suggestions. We address them below.\n\nThe upcoming comments concern the interpretation of the spectral gap condition (and the second largest singular value $s_2$). To facilitate this discussion, we introduce a simple example by starting with an arbitrary value of $s = s_2 < 1$ and construct a probability distribution $P$ that satisfies the condition in Eq (8) of the paper. \n\n**Example:** For $m=2$, we have that $\\mathcal{X} = \\\\{x_1, x_2\\\\}$ and $\\mathcal{Y} = \\\\{y_1, y_2\\\\}$, so every element in $h \\in L^2(P)$ can be represented by four numbers $(h(x_1, y_1), h(x_2, y_1), h(x_1, y_2), h(y_1, y_2))$. In the case of uniform marginals, we can verify directly that Eq (8) can be satisfied by setting $\\alpha_1 = \\beta_1 = (1, 1, 1, 1)$, $\\alpha_2 = (1, -1, 1, -1)$, $\\beta_2 = (1, 1, -1, -1)$ and $P(x_1, y_1)  = P(x_2, y_2) = (1+s)/4$ and $P(x_1, y_2) = P(x_2, y_1) = (1-s)/4$ . Thus, as $s \\rightarrow 1$, the distribution becomes “fully dependent” as $Y$ and $X$ are completely determined by one another. As $s \\rightarrow 0$, the $X$ and $Y$ are independent. Intuitively, the marginals would contain the most information when $s = 0$, as the distribution can be computed directly from them. This idea can be generalized to construct similar distributions for $m >2$, which we use for the simulations in the attached PDF.\n\n>**Can the authors provide more details on how the assumptions, such as the spectral gap condition, hold or when will they not hold in practical scenarios? Are there specific types of data or models where these assumptions are more likely to be satisfied? Will the assumptions limit the applicability of the findings?**\n\nAs mentioned in Assumption 2, it will hold if $P(x, y) > 0$ for any $(x, y)$ such that $P_X(x) > 0$ and $P_Y(y) > 0$ by the Perron-Frobenius theorem, but may hold otherwise. As seen above, the $s=1$ cases are often pathological, such as having perfect dependence between $X$ and $Y$. The positivity assumption is standard in even classical references on this subject; see [Bickel, Ritov, Wellner (1991)](https://projecteuclid.org/journals/annals-of-statistics/volume-19/issue-3/Efficient-Estimation-of-Linear-Functionals-of-a-Probability-Measure-P/10.1214/aos/1176348251.full) Eq (P3) for a equivalent condition. Thus, it is safe to assume the condition will hold in all practical scenarios of interest.\n\n>**…adding more visualizations or intuitive explanations may be better for understanding the key finding of the paper.**\n\nThank you for raising this point. Returning to the example above, we also see that because $\\alpha_1 = \\beta_1$, the angle between the subspaces can be measured by the angle between $\\alpha_2$ and $\\beta_2$. By direct computation, we can see that $\\langle \\alpha_2, \\beta_2 \\rangle = s$, which means that the cosine of the angle $a$ between the subspaces is $s = \\cos a$. Thus, the geometric interpretation of $s$ is the cosine of the angle between $L^2(P_X)$ and $L^2(P_Y)$. This can be visualized in Figure 7 of the attached PDF. In fact, this holds even for non-uniform marginals. With $P(x_1) = p_X$ and $P(y_1) = p_Y$, a slightly more tedious computation will show that Eq (8) is satisfied when $\\alpha_1 = \\beta_1 = (1, 1, 1, 1)$, $\\alpha_2 = (\\sqrt{(1-p_X)/p_X}, -\\sqrt{p_X/(1-p_X)}, (\\sqrt{(1-p_X)/p_X}, -\\sqrt{p_X/(1-p_X)})$, $\\beta_2 = (\\sqrt{(1-p_Y)/p_Y}, \\sqrt{(1-p_Y)/p_Y}, -\\sqrt{p_Y/(1-p_Y)}, -\\sqrt{p_Y/(1-p_Y)})$, and $P(x_1, y_1)  = p_Xp_Y + s \\sqrt{p_X(1-p_X)p_Y(1-p_Y)}$. In this case, we still have that $\\langle \\alpha_2, \\beta_2 \\rangle = s$, so the non-uniform marginals do not warp this angle.\n\n>**Are there any other variance reduction techniques and how does data balancing compare to those techniques?**\n\nAnother approach for variance reduction through marginal fitting used in stochastic simulation and one-dimensional optimal transport is to use the [quantile transform](https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.quantile_transform.html). That is, assuming that $X_1, \\ldots, X_n$ are real-valued, we may replace them by $F_n(X_1), \\ldots, F_n(X_n)$, where $F_n$ is the empirical CDF (in other words, $F_n(X_i)$ is the rank of $X_i$ divided by $n$). These random variables are uniformly distributed. By applying $F_X^{-1}$, the resulting (discrete) inform variables are now approximately distributed according to $F_X$. This is an approximation of the inverse CDF sampling technique, as $F_X^{-1}(U)$ for a continuous uniform random variable $U$ will have CDF $F_X$. We call this the “Inverse Probability (IP) Weighting” method and compare it to the empirical distribution and the balanced distribution in a simulation (see Figure 9 of the attached PDF). Across sample sizes, balancing improves on mean squared estimation error. While the values near $s=0$ imply that there is “more information” in the marginals, the inherent variability of the distribution also increases, explaining how various methods do not achieve perfect variance reduction. \n\n>**the experiments are somewhat limited in scope… Can the findings be extended to other areas?**\n\nThank you for raising this point. We also show performance on an image retrieval task in Figure 8 of the attached PDF using the Pascal VOC benchmark. We show mean average precision (MAP) across 30 queries in a 17k image database. Qualitatively similar findings hold, in that the additional iteration of balancing (over the original CLIP objective) yields performance improvements and warrants further investigation.\n\n>**For Figure 2, Can the authors provide experiment results with more iterations?**\n\nWe show this experiment in Figure 6 of the attached PDF. Both from observing the performance of the $k=5$ variant as well as directly observing the marginals in Figure 3 of the manuscript, we see that the marginals stabilize very quickly (after $k=2$ iterations in most cases), and performance remains similar beyond this point.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "te2iCpT8UN", "reviewer": "Reviewer_vTkj", "summary": "This paper introduces a technique called iterative data balancing—altering data distributions to match predefined marginal distributions—that can lead to variance reduction in model predictions. The authors highlight its utility for self-supervised learning, which has been used to train several foundation models. The results demonstrate that iterative rebalancing of data leads to improvements in zero-shot learning performance and a reduction in variance among the empirical marginals with more than one iteration (k>1) of their technique.", "questions": "1. In your work, your target marginals were uniform; how would your method respond to non-uniform marginals?\n2. Your target marginals were accurately specified (as uniform). What if the target marginals of the two distributions were less accurately specified (i.e., the true underlying distributions are not well-known)? How do you think that this would influence the empirical results of your technique (e.g., zero-shot average per-class recall)?\n3. Are there existing methods for variance reduction and data balancing? If so, why did you not include an empirical comparison to existing methods?\n4. You mention that the zero-shot evaluation metrics are difficult to produce intervals for (i.e., you are missing error bars). Why is this the case?", "limitations": "I don't feel the authors addressed the limitations of their work significantly, aside from noting that their methods were different as those \"in practice\". As already mentioned in the questions section, how robust are your results to imbalance and uncertainty in the target marginals? Furthermore, are there situations where your technique fails to improve the empirical results? If so, what are they?", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper has theoretical contributions that include the derivation of non-asymptotic bounds that quantify the variance reduction achieved through their data balancing technique. The authors also present empirical studies that demonstrate the effectiveness of their proposed balancing technique. The authors discuss the utility of data balancing across different tasks, such as image-caption pair matching and self-supervised clustering, identifying the utility of their approach. Their approach has the potential for adoption in various domains, including in the training of foundation models.", "weaknesses": "The authors could expand the range of experiments to include a more diverse set of tasks, which in turn could enhance the generalization of the findings. Furthermore, their iterative data balancing technique relies heavily on predefined (uniform) target marginal distributions (see questions about this in next section). Finally, the iterative nature of the proposed data balancing technique may introduce significant computational demands. The paper could benefit by a more comprehensive overview of how the iterative technique computational overhead is impacted by very large datasets and/or models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments. We address them below.\n\n>**The authors could expand the range of experiments to include a more diverse set of tasks.**\n\nThank you for raising this point. We also show performance on an image retrieval task in Figure 8 of the attached PDF using the Pascal VOC benchmark. We show mean average precision (MAP) across 30 queries in a 17k image database. Qualitatively similar findings hold, in that the additional iteration of balancing (over the original CLIP objective) yields performance improvements and warrants further investigation.\n\n>**Data balancing technique relies heavily on predefined (uniform) target marginal distributions...**\n\nWe first emphasize that neither the theoretical results nor the method described require the marginals to be uniform. The matrix scaling/Sinkhorn algorithm (as balancing is called in other contexts) has been used across domains for arbitrary marginals. The choice of the uniform marginals in the experiments was motivated by the fact that leading SSL methods already in use currently can be shown to be exactly, mathematically equivalent to balancing with uniform marginals.\n\nHowever, as far as inaccurately specified marginal distributions, it is an interesting question what the balanced measure converges to ($n \\rightarrow \\infty$) when $(P_X, P_Y)$ are not the marginals of the true probability measure $P$. This is outside the scope of this work but is a natural follow-up investigation.\n\n>**Finally, the iterative nature of the proposed data balancing technique may introduce significant computational demands.**\n\nThank you for identifying this point. The balancing technique is not a proposal of this paper; it is already in use in large foundation models such as CLIP and DINO, as described in Section 2. In particular, the approach applies to minibatches, and simply reduces to row scaling and column scaling of a matrix. The batch sizes for training the largest CLIP models are on the order of $10^4$, so taking row and column sums of this matrix is usually not the computational bottleneck of these training workloads. \n\n>**Are there existing methods for variance reduction and data balancing?**\n\nThank you for raising this important point. Another approach for variance reduction through marginal fitting used in stochastic simulation and one-dimensional optimal transport is to use the [quantile transform](https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.quantile_transform.html). That is, assuming that $X_1, \\ldots, X_n$ are real-valued, we may replace them by $F_n(X_1), \\ldots, F_n(X_n)$, where $F_n$ is the empirical CDF (in other words, $F_n(X_i)$ is the rank of $X_i$ divided by $n$). These random variables are uniformly distributed. By applying $F_X^{-1}$, the resulting (discrete) inform variables are now approximately distributed according to $F_X$. This is an approximation of the inverse CDF sampling technique, as $F_X^{-1}(U)$ for a continuous uniform random variable $U$ will have CDF $F_X$. We call this the “Inverse Probability (IP) Weighting” method and compare it to the empirical distribution and the balanced distribution in a simulation (see Figure 9 of the attached PDF). Across sample sizes, balancing improves on mean squared estimation error. While the values near $s=0$ imply that there is “more information” in the marginals, the inherent variability of the distribution also increases, explaining how various methods do not achieve perfect variance reduction. \n\n>**You mention that the zero-shot evaluation metrics are difficult to produce intervals for. Why is this the case?**\n\nThank you for identifying this point of improvement. We meant that incorporating error estimates based on bootstrapping the datasets themselves would be computationally infeasible, as each point on the curve requires the evaluation of a foundation model to be used in the CLIP Benchmark suite. However, we have incorporated multiple seeds to account for uncertainty across training runs. Please see both Figure 6 and Figure 8 of the attached PDF to see that qualitatively similar trends hold for the average performance and performance across individual seeds.\n\n>**Are there situations where your technique fails to improve the empirical results? If so, what are they?**\n\nAs mentioned in the response to the Reviewer os3n, the learned representations based on more than 2 iterations of rebalancing (Figure 6 of the attached PDF) show that performance in zero-shot tasks may no longer improve after 2 iterations. Similarly, as seen in the simulation in Figure 9, other methods for balancing may perform as well or better when the random variables in $X$ and $Y$ are nearly independent. Please see the \"Example\" in the response to Reviewer os3n for an intuitive explanation of this simulation. We emphasize that the experiments are meant to be illustrative of the theory established in Sections 2 and 3, and we do not make claims of state-of-the-art empirical methods.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their hard work reviewing our paper and providing concrete comments! We collect the broad points made below and address other reviewer concerns in the individual responses. To summarize, our paper provides three theoretical innovations:
1. The first quantitative and non-asymptotic analysis of the statistical effect of data balancing (a.k.a. matrix scaling or biproportional fitting), which spans self-supervised learning (SSL), optimal transport, and statistics.
2. A novel interpretation of contrastive objectives relating them to data balancing with subsequent algorithmic implications.
3. A mathematical recursion formula for balanced probability measures that can be used as an independent technical tool.

The main contributions of the paper are of a theoretical nature, for which we provide complete proofs in the Appendix sections, and we also provide numerical experiments as illustrations. While the paper should be assessed as such, we have addressed all feedback from the reviewers.

**The main concerns across reviews were the following.**

**Additional experimentation:** The reviewers requested additional experiments, including performance on tasks other than zero-shot image classification, comparison to other balancing approaches, and extensions of the balanced-CLIP variant for more than two iterations. These experiments are described in the responses and the associated results can be found in the attached PDF. 

**Intuition behind theoretical assumptions:** The reviewers also requested clarity on theoretical assumptions, such as the spectral gap condition or whether uniform marginals are needed. We address these with both intuitive explanations and companion simulation experiments.

Please see the individual responses for more details. Thank you!

Meta-review: This work studies data balancing in SSL and show that avoiding representation collapse also leads to a variance reduction. The paper establishes non-asymptotic bounds quantifying this variance reduction. The authors addressed the concerns raised by the reviewers and produced the additional experiments they requested. I found those additional results helpful; authors should incorporate the additional material in the final paper or Appendices. All reviewers voted for acceptance.
Decision: Accept (poster)