The Benefits of Balance: From Information Projections to Variance Reduction

Authors: lang  liu
ronak  mehta
soumik  pal
zaid  harchaoui
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 05:17:54

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.

Main paper: 
 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 even have zero-shot capabilities out of the box. Supervised learning, stricto sensu, is too limited a framework for these billion-parameter, data-hungry models, and a question at the heart of modern machine learning is 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 that circumvent the need for large amounts of "strong" labels. In SSL, a model is trained on a generic pseudo-task 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] 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: minibatch sampling and gradient-based optimization using backpropagation. On the other hand, SSL adds clever, yet less-understood techniques to the training pipeline. To illustrate this, consider a minibatch (X 1 , Y 1 ), . . . , (X n , Y n ) of training examples observed in X × Y (e.g. feature-label or image-caption pairs) and let P n be the empirical distribution of the minibatch. For a model parameterized by θ ∈ R d with loss function h θ , a standard stochastic learning algorithm involves computing the minibatch lossE Pn [h θ (X, Y )] = 1 n n i=1 h θ (X i , Y i ) (1)and backpropagating through it to produce a minibatch stochastic gradient estimate. The algorithm then proceeds with the stochastic gradient training or a variant thereof (e.g., Adam). Self-supervised methods often modify this recipe by intervening on the optimization algorithm in a minibatch-specific way. 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 a 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 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 on example i is 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 a better theoretical understanding of the objectives and algorithms underlying these empirically effective recipes. In particular, we want to shed a theoretical light on their marginal benefits over traditional learning methods. In this work, we show that such recipes often enjoy an unsuspected benefit: reducing the variance of the empirical minibatch objective.To be more concrete, we formalize the broad steps of the training algorithms above as follows.Let Z 1 , . . . , Z n be a minibatch containing data points of arbitrary type (e.g. unlabeled images).Possibly using a model parameterized by θ, this original data source is mapped to another minibatch (X 1 , Y 1 ), . . . , (X n , Y n ) of derived pairs in X × Y. For example, in SwaV, each Z i is an image with vector representation X i and pseudo-label Y i based on clustering. In CLIP, each Z i is an image-caption pair and X i and Y i are the vector representations of its image and text components, respectively. We may also use the parameterized model to compute a probability distribution P n,θ over X × Y, and view each training step above as exactly or approximately optimizing the objectiveEQUATIONThis reduces to empirical risk minimization on the minibatch objective (1) when Z = (X, Y ) (each data point is originally observed in X × Y) and P n,θ = P n (the empirical distribution of the data is used, regardless of the model). Beyond this setting, one specific example of P n,θ has been applied across various families of self-supervised learning (as detailed in Sec. 2), which we refer to as data balancing or simply balancing, the primary subject of this work.For a probability measure Q on X × Y, let Q X and Q Y be the respective marginals on X and Y and let Q X|Y and Q Y |X denote the respective conditional distributions. Given an initial measure P (0) n over X × Y which may depend on the model2 , and fixed marginal distributions P X on X and P Y on Y, balancing refers to generating the sequence P (0) n , . . . , P (k) n by applying the operationsEQUATIONin an alternating fashion. After enough iterations, the resulting probability measure approximately marginalizes to P X and P Y in each variable. When X and Y are finite with |X | = m and |Y| = l, these operations reduce to rescaling the rows of an (m × l)-matrix by P X /Q X or its columns by P Y /Q Y . This algorithm has a decades-old history and is known in other contexts as Sinkhorn-Knopp matrix scaling [Sinkhorn, 1967] , iterative proportional or biproportional fitting [Johnston and Pattie, 1993] , and raking-ratio estimation [Thompson, 2000] . The marginals P X and P Y can represent auxiliary information or inductive bias from users, such as the desire for balanced clusters.As we describe in detail in Sec. 2, both self-labeling and contrastive approaches in SSL implicitly use P n,θ = P (k) n in (2) for different choices of P (0) n . In other words, they embed a learnable balancing operation in their objectives, where the objective (as in (2)) or the stochastic gradient estimate depends linearly on P n,θ . 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 θ (thus dropping the subscript from the quantities above) and analyze the fluctuations of the unbalanced and balanced objectives. The formal problem statement is as follows. Let X and Y be two finite sample spaces on which there is an unknown probability measure P with known marginals (P X , P Y ). We observe independent data (X 1 , Y 1 ), . . . , (X n , Y n ) ∼ P , defining the empirical measure P n = 1 n n i=1 δ (Xi,Yi) . Let P (0) n = P n and P (k) n denote the output of k ≥ 1 iterations of data balancing (see Sec. 3 for the precise definition). Finally, letting h : X × Y → R be a function of interest, we define the population parameter φ and its k-step balanced estimator φ (k) n by φ := E P [h(X, Y )] and φ (k) n :EQUATION)Our goal is to establish theoretical guarantees on the mean squared error (MSE) E P [(φ (k) n -φ) 2 ] of estimating φ using φ (k) n , with an informative dependence on the sample size n, number of iterations k, target marginals (P X , P Y ), and test function h. We are particularly interested in its comparison to the direct estimator based on the empirical measure φ (0) n = 1 n n i=1 h(X i , Y i ), as to quantify the effect of the auxiliary information (P X , P Y ). Our analysis uncovers two surprising facts. Firstly, while originally proposed for a 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 clarify the mathematical connection between the classical data balancing methods and the modern representation learning techniques mentioned above. In Sec. 3, we prove a new upper bound on the MSE of the balancing estimator φ (k) n . The bound decomposes into an O(n -1 ) first-order variance term and an O(n -3/2 ) second-order term. The first-order term is shown to have a strict improvement over the empirical measure baseline with a fine-grained dependence on the spectra of two particular Markov operators. Our proof technique relies on a recursion decomposition for balancing-based 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. 
data balancing in practice: 
 We now expand on the examples of data balancing applications mentioned in Sec. 1. To demonstrate a precise connection to (2), we describe how a collection of training examples Z 1 , . . . , Z n observed in an original data space Z (e.g. grayscale images) is mapped to a probability measure P n,θ . 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 Y, 2) the initial measure P (0) n , 3) the function h, and 4) the target marginals (P X , P Y ) for this measure to fit. From that point, the iterations of (3) produce P (1) n , . . . , P (k) n , and we set P n,θ := P (k) n . For ease of presentation, we hide the dependence of P (k) n = P n,θ and h ≡ h θ on the model parameter θ when constructing φ (k) n as described in (4). Note that the sample spaces X and Y do not necessarily correspond to feature-label pairs in each example below; each example given below corresponds to a different choice of X and Y (see Fig. 1 ).Example 1: Self-Supervised Learning. Balancing is used in discriminative clustering and selfsupervised clustering; see [Jones et 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 as an intermediate step (or as the task itself) in many SSL pseudo-tasks. At a high level, this approach involves passing elements of a minibatch through two encoders to 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 f θs : Z → R r and f θt : Z → R r , colloquially known as the student and teacher networks, respectively. Here, we let {Z i } n i=1 be a minibatch of n images, with X = {Z 1 , . . . , Z n } and Y = {1, . . . , l} , where m = n and the elements of Y index learnable cluster representation vectors c 1 , . . . , c l ∈ R r . Thus, we consider the overall parameter vector to be θ := (θ s , θ t , c 1 , . . . , c l ). Given temperature hyperparameters ϵ, τ > 0, the initial measure and loss function are given by the expressionsP (0) n (x, y) ∝ e f θs (x) ⊤ cy/ϵ and h(x, y) = log e f θ t (x) ⊤ cy /τ l y ′ =1 e f θ t (x) ⊤ c y ′ /τ .Directly optimizing x,y P (0) n (x, y)h(x, y) without any constraints would lead to collapse, so it is balanced before optimized. The target marginals P X and P Y are given by the discrete uniform measures on X and Y. This formulation is often derived by solving an optimal transport problem with the Sinkhorn-Knopp algorithm to assign soft cluster labels, the iterative solution result from this procedure is precisely P (k) n . The intuition behind the choice of uniform marginal P X is that each data point has an equal amount of mass to allot to each cluster, whereas P Y captures that the cluster sizes are equal. The number of iterations k is selected based on optimization considerations. Example 2: Contrastive Learning. Contrastive Language-Image Pre-Training [Radford et al., 2021] , or CLIP, is an architecture with an image encoder and a text encoder that map to a joint embedding space. Trained using image-caption pairs, the loss promotes representations such that images and text that are paired in the minibatch are close, whereas those that are not paired are far. The latter aspect (promoting dissimilarity of unpaired images/text) is what prevents collapse in this framework. Our interpretation of the CLIP objective as an implicit data balancing procedure is novel, to our knowledge. Under this interpretation, we demonstrate that the objective is in fact a nonlinear function of P n,θ , whereas its gradient will have a linear form similar to (2). In this case, each Z i = (X i , Y i ), where X i is an image and Y i is an associated caption. We have that X = {X 1 , . . . , X n } and Y = {Y 1 , . . . , Y n } , so that m = n. Consider an image encoder f θ I : X → R r and text encoder f θ T : Y → R r with parameter vector θ = (θ I , θ T ). The initial, unnormalized measure and the (in this case, vector-valued) function h are chosen based on these encoded representations:EQUATIONWhile we usually interpret h as a loss function, we will show below that the CLIP loss depends nonlinearly on P n,θ , while the gradient has a linear dependence. If we believe, as in Example 1, that the target marginals (P X , P Y ) 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 Y, respectively. Because there is no preference for starting the iterations with the X or Y dimension first, we may consider both orderings. Let Q (1) n be one iteration of balancing in the Y dimension and R (1) n represent one such iteration in the X dimension. Then the original CLIP objective L CLIP n can be recovered (up to an additive constant) asEQUATIONThe measureP n,θ = 1 2 Q (1) n + 1 2 R (1)n 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 (1) n , R (1) n )) recovers the expression for h in (5). The objective is often interpreted as an average of cross-entropy loss terms, each representing the prediction of one modality's original pair from the other. In our formulation, L CLIP n can also be viewed as an average negative log-likelihood under the Q (1) n and R (1) n . It is also of interest to study the effect of using Q (k) n and R (k) n for k ≥ 0 in general, as we show in Sec. 4.Example 3: Metadata Curation. Here, we consider taking M n to be an entire training set, as opposed to a particular mini-batch. At the billion-parameter 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 M n (possibly captions for images) is synthesized using a list of keywords {y 1 , . . . , y l } so thatX = {Z 1 , . . . , Z n } , Y = {y 1 , . . . , y l } ,meaning that m = n. The keywords are matched to texts within X via substring matching. While the approach of Xu et al. [2024] (dubbed MetaCLIP) pools all matched keywords on every text to measure the "distribution" of keywords, we consider a version in which each text Z i can only be labeled with a single keyword y j . This allows for a true joint probability measure on 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 et al. [2024] , the data are resampled so that this distribution of keywords over matches is closer to uniformity, i.e. 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 haveEQUATIONwhere ℓ(x) denotes the loss of a model evaluated at a single text x ∈ 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. 1 by analyzing balancing-based estimators. We then return to examples mentioned above in Sec. 4, illustrating how the theoretical analysis can be translated to algorithmic variants. 
theoretical analysis of variance reduction: 
 We now present theoretical guarantees on the mean squared error (MSE) of the data-balanced estimator φ (k) n and highlight relevant points in the proofs. For readers' convenience, a notation table (Tab. 1) is in Appx. A. We first give context on 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 to linear operators on a vector space and using a 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 we consider sample spaces (X , Y), along with true and unknown joint distribution P on X × Y with known marginals (P X , P Y ). For ease of presentation, we assume that |X | = |Y| = m, although the arguments do not rely on equal support sizes. We make the following assumption throughout, which is usually satisfied by the desired marginals P X and P Y , such as in the uniform cases discussed in Sec. 2: the target marginals P X (x) > 0 and P Y (y) > 0 for all x ∈ X and y ∈ Y. We define P (0) n = P n as the empirical measure and for k ≥ 1 constructEQUATION)By direct computation, we see that the iterations in (8) are equivalent to applying (3) for k odd and even, respectively. See Fig. 2 (left) for a visualization of this procedure. The iterations are well-defined for all k under the event that Supp(P n,X ) = Supp(P X ) and Supp(P n,Y ) = Supp(P Y ), i.e., all observed row counts and column counts are non-empty. 3To provide background, the scheme of alternating the operators ( 8) is often seen as an iterative algorithm to solve the problem minEQUATIONwhere Π(P X , P Y ) denotes the set of probability measures on X × Y that marginalize to P X and P Y in each variable and KL(•∥•) denotes the Kullback-Leibler divergence. The iterations (8) can be derived from the alternating minimization subproblemsP (k) n (x, y) := arg min {Q:Q X =P X } KL(Q∥P (k-1) n ) k odd arg min {Q:Q Y =P Y } KL(Q∥P (k-1) n ) 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 χ 2 -divergences. This viewpoint is relevant as previously, efforts have been made (e.g. in Bickel et al. [1998] ) 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-form expression for the solution of (9). A key mathematical outcome of our analysis is that the closed-form expressions of the projections (8) can be used to compute the reduction in mean squared error at each iteration. Thus, by letting k → ∞ (scaled appropriately against n), we can determine the reduction for the solution of (9) for large n. This is the subject of the upcoming Cor. 2.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 over probability measures, we may use familiar tools to understand alternating orthogonal projections in a vector space. To define them, we first let L 2 (P ) to be the set of functions h :X × Y → R satisfying E P h 2 (X, Y ) < ∞.Even though X × Y is finite, working within L 2 (P ) will be analytically convenient. Let L 2 (P X ) be the subspace of L 2 (P ) containing functions that only depend on the first argument x ∈ X and define L 2 (P Y ) analogously. These are the solid-colored subspaces in Fig. 2 (right). Next, let µ X : L 2 (P ) → L 2 (P X ) and µ Y : L 2 (P ) → L 2 (P Y ) be defined as, for any h ∈ L 2 (P ),µ X h = arg min f ∈L 2 (P X ) E P (h(X, Y ) -f (X)) 2 =⇒ [µ X h](x, y) := E P [h(X, Y )|X] (x)The operator µ X is an orthogonal projection onto L 2 (P X ). The orthogonal projection operator µ Y onto L 2 (P Y ) is defined analogously. We may also define the conditional debiasing operators C X = I -µ X and C Y = I -µ Y , which each project onto the orthogonal complements of L 2 (P X ) and L 2 (P Y ), visualized as subspaces with dotted border in Fig. 2 (right). To understand the importance of the conditional mean and debiasing operators, we give a recursive formula that forms the backbone of our analysis. Define µ k = µ X for k odd and µ k = µ Y for k even, and define C k similarly. Thus, by using the notationQ(h) := E Q [h(X, Y )],we have by linearity of expectation thatEQUATION)To justify the first line, we discuss the case when k is odd. Notice that µ X h is only a function of X, so its expectation only depends on P X that is equal to P (k) n,X 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 σ 2 k /n where σ 2 0 := Var(h) andEQUATIONand all other terms are O(n -3/2 ). 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φ (0) n = 1 n n i=1 h(X i , Y i )is unbiased, and so its MSE is equal to σ 2 0 /n. Define in addition p ⋆ := min{min x P X (x), min y P Y (y)} which measures the non-uniformity of the target marginals. We have that p ⋆ is positive because both P X and P Y are positive. We now state the first main result. Theorem 1. For a sequence of data-balancing estimators (φ (k) n ) k≥1 as defined in (4), there exists an absolute constant C > 0 and distribution dependent constant s ∈ [0, 1) and such the following holds for σ 2 gap = σ 2 0 -σ 2 k : For n ≥ C[log 2 (2n/p ⋆ ) + m log (n + 1)]/p 2 ⋆ and k ≥ 1, we haveEQUATION)The quantities σ 2 gap and s are quantified toward the end of this section and are dependent on eigendecays of the conditional mean operators for each variable under P . Furthermore, σ 2 gap > 0 except for the pathological case of µ X h being a constant function. Showing Thm. 1 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 ℓ ≥ 1 is odd, we see that[P (ℓ) n -P (ℓ-1) n ](C ℓ . . . C k h) = x,y P X (x) P (ℓ-1) n,X (x) -1 • [C ℓ . . . C k h](x, y)P (ℓ-1) n (x, y).The first (blue) term in the product quantifies the disagreement between the X -marginal of P (ℓ-1) n and the true marginal, which can be bounded in terms of KL(P (0) n,X ∥P X ) and is shown to be O(n -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 a bound for the sensitivity of balancing to misspecified marginals; see Appx. D.5. Given Thm. 1, a natural next step is to quantify the gap between σ 2 0 and σ 2 k , which requires finergrained properties of C X and C Y . Notably, we show that as k → ∞, σ 2 k approaches a limiting value. Thus, via (12), by using k = o(n 1/12 ) obtains asymptotic variance of the solution to (9). This contrasts with Albertus and Berthet [2019] , in which the dependence of a quantity similar to (12) is exponential in k, meaning that k = o(log(n)) is required for convergence under this argument.From Orthogonal Projections to Variance Reduction. We now clarify what is precisely meant by the "spectrum" of the conditional mean operators µ X and µ Y . As proven using a singular value decomposition (Prop. 3) in Appx. B.1, there exists a basis {α j } m j=1 of L 2 (P X ), a basis {β j } m j=1 of L 2 (P Y ), and real values {s j } m j=1 , that satisfy µ Y α j = s j β j and µ X β j = s j α j for j ∈ {1, . . . , m} .(13) Furthermore, α 1 = 1 X and β 1 = 1 Y leading to the equality ⟨f,α 1 ⟩ L 2 (P X ) = E P X [f (X)].Finally, s 1 = 1 and s j is non-negative and non-increasing in j. For a concrete example, consider m = 2, in which case P can be written as a matrix in R 2×2 and elements of L 2 (P X ) and L 2 (P X ) are vectors in R 2 . Then, in the case of uniform marginals, we can verify directly that (13) can be satisfied by settingα 1 = β 1 = 1 1 , α 2 = β 2 = 1 -1 , and P = 1 4 1 + s 1 -s 1 -s 1 + s (14)for s = s 2 (the second largest singular value). Thus, as s → 1, the distribution becomes "fully dependent" as Y and X are completely determined by one another. As s → 0, P approaches the product measure. Geometrically, because α 1 = β 1 , we know that the angle a between the subspaces L 2 (P X ) and L 2 (P Y ) is given by the angle between α 2 and β 2 . By computing their inner product in L 2 (P ), we have that ⟨α 2 , β 2 ⟩ L 2 (P ) = P, α 2 β ⊤ 2 = s = cos a. Thus, s = 0 indicates orthogonality of these subspaces, alluding to the independence of X and Y .Returning to m ≥ 2, we consider the following as a sufficient condition for variance reduction: the operators µ X and µ Y have a positive spectral gap, i.e., s 2 < s 1 . Note that this assumption is satisfied when P (x, y) > 0 for all (x, y) ∈ X × Y by the Perron-Frobenius Theorem [Horn and Johnson, 2013, Chapter 8] . Using the intuition from Fig. 2 , this rules out pathological cases such Y being a deterministic function of X, for instances Under the spectral gap condition, the singular values {s j } m j=2 that are strictly less than 1 will determine a geometric rate of decay in variance given in Cor. 2. The left and right singular functions α j : X → R and β j : Y → R will define a useful coordinate system to represent projections of h when analyzing φ (k) n . Indeed, let h -P (h) be the centered test function. Because µ X h ∈ L 2 (P X ) and µ Y h ∈ L 2 (P Y ) , we may decompose this function on the two bases to writeEQUATIONCor. 2 below relates the (normalized) variance σ 2 k of the first-order term to the one of the sample mean φ (0) n . In fact, it shows that the variance reduction σ 2 0 -σ 2 k decays geometrically to the quantityσ 2 gap := m j=2 u 2 j + (v j -s j u j ) 2 1 -s 2 j .For simplicity, we only present the result for k even, i.e., σ 2 2t . Corollary 2. The variance reduction achieved by t + 1 iterations of the C Y C X operator can be quantified asσ 2 0 -σ 2 2(t+1) = σ 2 gap - m j=2 s 2 j (v j -s j u j ) 2 1 -s 2 j s 4t j = m j=2 u 2 j + (1 -s 4t+2 j ) (v j -s j u j ) 2 1 -s 2 j .Intuitively, the operators C X and C Y are the main sources of the variance reduction via orthogonality.Since α 1 = 1 X , we can see that the reduction will always be strictly positive as long as µ X h is not a constant function. Finally, using s := s 2 ≥ s j for j ≥ 2 gives the second term in Thm. 1. 
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. E. Code to reproduce the data and experiments can be found at https://github.com/ronakdm/balancing. Model, Datasets, and Evaluation. Throughout, we consider training variants of CLIP models (see Sec. 2), which require a dataset of image-caption pairs. For the training set, we 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 f θ I and a text encoder f θ T . In all cases, we use a fixed image/text encoder as a base vector representation and compose it with a trainable feed-forward neural network, i.e., f θ = f head θ • f base . We fix the base 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 c ∈ {1, . . . , C} is predicted by associating to each c a natural language prompt y c , and predicting the scores s(x) = (s 1 (x), . . . , s C (x)), withs c (x) = e ⟨fθ I (x),f θ T (yc)⟩/τ C c ′ =1 e ⟨fθ I (x),f θ T (y c ′ )⟩/τ (16)for a temperature τ > 0. Multiple prompting strategies can be used depending on 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 on whether the contrastive learning objective from Sec. 2 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. 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 can also be related to "selfdistillation" approaches more broadly [Grill et al., 2020 , Chen and He, 2021 , Oquab et al., 2024] .Historical motivations for data balancing include census or survey data, in which P n is a crosstabulation 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 single-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 to 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 on those from Albertus and Berthet [2019] primarily in the dependence of the number of iterations k on the sample size n to achieve convergence guarantees (from logarithmic to polynomial).Matrix scaling is a popular algorithm to solve entropy-regularized optimal transport (EOT). We refer to Peyré and Cuturi [2019] for a survey. See also Courty et 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 is itself a 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 are dependent, forming the random initial measure P n , whereas P X and P Y are fixed auxiliary information. 
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 a new non-asymptotic bound on the mean-squared error of balanced estimators as 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. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [No] Justification: We do not release general-purpose models, only small-scale illustrative ones.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 a best faith effort.12. Licenses for existing assets Question: 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?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 use 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 a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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. 
 
 Appendix: Symbol Description X , Y Sample spaces for two data sources.Support sizes m = |X | and l = |Y|.We sometimes assume m = l for ease of presentation.Probability measure on X × Y (the data-generating distribution).n Sample size.(X 1 , Y 1 ), . . . , (X n , Y n ) Independent and identically distributed sample from P .P n Empirical measure of {(X i , Y i )} n i=1 . Q X , Q Y Marginals of measure Q on X × Y, e.g. R X , P Y , P n,X , etc. Supp(Q) For measure Q over Z, the set of values z ∈ Z such that Q(z) > 0. Q(h) The expected value of h under Q, or E Q [h(X, Y )]. (P (k) n ) k≥1Sequence of iterations of (8).k Iteration count of (8).The event {Supp(P n,X ) = Supp(P X ) and Supp(Pn,Y ) = Supp(P Y )}. h Test function h : X × Y → R of interest.φThe estimand x,y h(x, y)P (x, y).nThe estimator x,y h(x, y)P (k) n (x, y), when well-defined.φ (k) nThe estimator φ (k) n := φ(k)n 1 S + φ (0) n 1 S c . G (k) n (h) Normalized error √ n( φ(k) n -φ). V (k) n (h) Remainder defined in Prop. 15. h Centered function h -E P [h]. σ 2 k Variance term E P (C 1 , . . . C k h) 2 . p ⋆ min{min x P X (x), min y P Y (y)}. L 2 (P ) Functions h : X × Y → R (as X × Y is finite). L 2 (P X ), L 2 (P Y )Subspaces of L 2 (P ) containing functions only of x ∈ X and y ∈ Y, respectively.µ X , µ Y Conditional expectation operators [µ X h](x) := E P [h(X, Y )|X] (x)and[µ Y h](y) := E P [h(X, Y )|Y ] (y). C X , C Y Debiasing/centering operators C X = I -µ X and C Y = I -µ Y . µ k , C k (µ X , C X ) for k odd and (µ Y , C Y ) for k even. {s j } m j=1Singular values in Prop. 3.{α j } m j=1 , {β j } m j=1Bases for L 2 (P X ) and L 2 (P Y ) in Prop. 3.Table 1 : Notation used throughout the paper.This section is dedicated to establishing the variance reduction result in Cor. 2 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 σ 2 0 and σ 2 k for even and odd iterations of k.Recall the conditional mean operators µ X and µ Y from Sec. 3, µ Y α j = s j β j and µ X β j = s j α j for j ∈ {1, . . . , m} ,[µ X h](x) := E [h(X, Y )|X] (x) and [µ Y h](y) := E [h(X, Y )|Y ] (y),EQUATIONand s j is non-negative and non-increasing in j.Proof. When µ X is restricted to L 2 (P Y ) and µ Y is restricted to L 2 (P X ), these operators are in fact adjoint in L 2 (P ), as by the tower property we have the relation ⟨f, 17). Next, we show that s 1 ≤ 1 and that 1 X is an eigenvector of µ X µ Y : L 2 (P X ) → L 2 (P X ) with eigenvalue 1, which confirms that s 1 = 1 and α 1 = 1 X by the definition of singular values (arguing symmetrically achieves β 1 = 1 Y ). By the variational representation of singular values [Gohberg et al., 1990 , Section IV.1 Equation ( 2)], we have that supµ X g⟩ L 2 (P X ) = E [f (X)E [g(Y )|X]] = E [E [f (X)|Y ] g(Y )] = ⟨µ Y f, g⟩ L 2 (P Y ) . Since µ Y : L 2 (P X ) → L 2 (P Y ) is a compactf :∥f ∥ L 2 (P X ) =1 ∥µ Y f ∥ L 2 (P Y ) = s 1 .Consider any f ∈ L 2 (P X ) such that ∥f ∥ L 2 (P X ) = 1. Define the conditional probability P X|Y (x|y) = P (x, y)/P Y (y) which is well-defined by assumption. Then, by the Cauchy-Schwarz inequality in L 2 (P X|Y ),∥µ Y f ∥ 2 L 2 (P Y ) = y∈Y x∈X f (x)P X|Y (x|y) 2 P Y (y) ≤ y∈Y x∈X f 2 (x)P X|Y (x|y)P Y (y) = x∈X f 2 (x) y∈Y P (x, y) = ∥f ∥ 2 L 2 (P X ) = 1. This proves that s 1 ≤ 1. For equality, notice that µ X µ Y 1 X = µ X 1 Y = 1 X , completing the proof.From Prop. 3, we establish two bases {α j } m j=1 and {β j } m j=1 of L 2 (P X ) and L 2 (P Y ), respectively. These bases span the range of the operators µ X and µ Y . We will consider the repeated application of the operator C Y C X , a sequence of two centering operations on some function h ∈ L 2 (P ), and compareE ((C Y C X ) t h) 2 against E h2 for h = h -E P [h].We establish the main result by measuring the reduction in variance from a single application, in terms of the coordinates of the function of interest on each of the two subspaces. We will then observe how these coordinates change iteration-to-iteration to give the final result. Lemma 4. For any h ∈ L 2 (P ) such thatE P [h] = 0, let µ X h = m j=1 u j α j and µ Y h = m j=1 v j β j .Then, we have thatE (C Y C X h) 2 = E h 2 - m j=2 u 2 j - m j=2 (v j -s j u j ) 2 .Proof. By orthogonality, we have thatE (C Y C X h) 2 = E ((I -µ Y )C X h) 2 = E (C X h) 2 -2E [(C X h)(µ Y C X h)] + E (µ Y C X h) 2 = E (C X h) 2 -2P Y ((µ Y C X h) 2 ) + P Y ((µ Y C X h) 2 ) = E (C X h) 2 -P Y ((µ Y C X h) 2 ) = E h 2 -P X ((µ X h) 2 ) -P Y ((µ Y C X h) 2 ).Because P (h) = 0, it holds by the tower property of conditional expectation that P X (µ X h) = 0, which implies thatu 1 = ⟨µ X h, α 1 ⟩ L 2 (P X ) = 0 =⇒ P X ((µ X h) 2 ) = m j=2 u 2 j .For the second term, observe that P X (C X h) = 0, so it holds by the tower property thatP Y (µ Y C X h) = 0, so P Y ((µ Y C X h) 2 ) = m j=2 ⟨µ Y C X h, β j ⟩ L 2 (P Y ) 2 .Next, we compute the term in the square by applying Prop. 3:⟨µ Y C X h, β j ⟩ L 2 (P Y ) = ⟨µ Y h, β j ⟩ L 2 (P Y ) -⟨µ Y µ X h, β j ⟩ L 2 (P Y ) = v j -µ Y m k=1 u k α k , β j L 2 (P Y ) = v j - m k=1 u k s k β k , β j L 2 (P Y ) = v j -s j u j ,which completes the proof.Lem. 4 ensures that we have a reduction on each iteration, with a 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 h 0 = h and h t+1 = (C Y C X )h t , along with the constants {u t,j } m j=1 and {v t,j } m j=1 given byµ X h t = m j=1 u t,j α j and µ Y h t = m j=1 v t,j β j .We have the following.Lemma 5. For all t ≥ 0, it holds that u t+1,j = s 2 j u t,j -s j v t,j , v t+1,j = 0.Proof. Fix any j ∈ [m], and use Prop. 3 to writeu t+1,j = ⟨µ X C Y C X h t , α j ⟩ L 2 (P X ) = ⟨µ X (I -µ X -µ Y + µ Y µ X )h t , α j ⟩ L 2 (P X ) = ⟨µ X µ Y µ X h t , α j ⟩ L 2 (P X ) -⟨µ X µ Y h t , α j ⟩ L 2 (P X ) = µ X µ Y m k=1 u t,k α k , α j L 2 (P X ) -µ X m k=1 v t,k β k , α j L 2 (P X ) = s 2 j u t,j -s j v t,j ,which proves the first part of the claim. For the second part, note thatµ Y C Y = 0, so ⟨µ Y C Y C X h t , α j ⟩ L 2 (P Y ) = 0.Using Lem. 4 and Lem. 5, we can simply accumulate the reduction incurred on every iteration.Proposition 6. Define the constants (u j ) m j=1 and (v j ) m j=1 byµ X h = m j=1 u j α j and µ Y h = m j=1 v j β j .Then, we may quantify the variance reduction achieved by t + 1 iterations of the C Y C X operator asE h2 -E ((C Y C X ) t+1 h) 2 = m j=2 u 2 j + (v j -s j u j ) 2 1 + s 2 j (1 -s 4t j ) 1 -s 2 j → m j=2 u 2 j + (v j -s j u j ) 2 1 -s 2 j as t → ∞.Proof. Apply Lem. 4 (t + 1)-times so thatE ((C Y C X ) t+1 h) 2 = E h2 - m j=2 t τ =0 (1 + s 2 j )u 2 τ,j + v 2 τ,j -2s j u τ,j v τ,j = E h2 - m j=2 v 2 0,j -2s j u 0,j v 0,j + t τ =0 (1 + s 2 j )u 2 τ,jas by Lem. 5, we have that v τ,j = 0 for τ > 0. Next, we unroll the definition of u τ,j so thatu τ,j = s 2 j u τ -1,j -s j v τ -1,j = s 2 j (s 2 j u τ -2,j -s j v τ -2,j ) -s j v τ -1,j = s 2τ -2 j (s 2 j u 0,j -s j v 0,j ) for τ > 0, yielding E h2 -E ((C Y C X ) t+1 h) 2 = m j=2 u 2 0,j + (v 0,j -s j u 0,j ) 2 + (1 + s 2 j )(s 2 j u 0,j -s j v 0,j ) 2 t τ =1 (s 4 j ) τ -1 = m j=2 u 2 0,j + (v 0,j -s j u 0,j ) 2 + (1 + s 2 j )(s 2 j u 0,j -s j v 0,j ) 2 t-1 τ =0 (s 4 j ) τ = m j=2 u 2 0,j + (v 0,j -s j u 0,j ) 2 + s 2 j (1 + s 2 j )(v 0,j -s j u 0,j ) 2 (1 -s 4t j ) 1 -s 4 j = m j=2 u 2 0,j + (v 0,j -s j u 0,j ) 2 + s 2 j (v 0,j -s j u 0,j ) 2 (1 -s 4t j ) 1 -s 2 j .Substitute u 0,j = u j and v 0,j = v j to complete the proof.We also present the corresponding result for k odd. The proof follows similarly by repeated application of the operator C Y C X . However, the iterations will be compared to σ 2 1 = E P (C X h) 2 , as we consider C X h as the "first" iteration to this process. Proposition 7. Define the constants (u j ) m j=1 byµ Y C X h = m j=1 u j β j .Then, we may quantify the variance reduction achieved by t + 1 iterations of the C X C Y operator asE (C X h) 2 -E ((C X C Y ) t+1 C X h) 2 = m j=2 u 2 j + (s j u j ) 2 1 + s 2 j (1 -s 4t j ) 1 -s 2 j → m j=2 1 + s 2 j 1 -s 2 j u 2 j as t → ∞.In order to have full monotonicity, we also need that σ 2 0 ≥ σ 2 1 . This follows by orthogonality, asEQUATIONThus, we can combine Prop. 7 and ( 18) to fully quantify the relationship between σ 2 0 and σ 2 k for k odd.This section is dedicated to deriving three representations of the balancing procedure as projections in various statistical divergences, as shown in Fig. 2 .We consider two sets of probability measures denoted byΠ X = {Q : Q X = P X } and Π Y = {Q : Q Y = P Y }.The marginal matching steps are written as projections in terms of a statistical divergence D (precisely, an f -divergence) in the formP X P (k-1) n,X ⊗ P (k-1) n = arg min Q∈Π X D(Q∥P (k-1) n ), P Y P (k-1) n,Y ⊗ R = arg min Q∈Π Y D(Q∥P (k-1) n ).We provide the derivations for three common choices of D: Kullback-Leibler (KL), reverse KL, and χ 2 . Using this viewpoint, and simply assuming the positivity of the marginal measures P X and P Y , we 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 k.The KL representation will be used in the proof of Prop. 14, which (recalling the sequence (P (k) n ) k≥1 from (8)), controls the error between P (k) n,Y and P Y for k odd and P (k) n,X and P X for k even.C.1.1 Projection in KL-Divergence Proposition 8. Assume that P X ≪ R X and P Y ≪ R Y , and defineEQUATION)Then, it holds thatEQUATIONandEQUATION)Proof. In the case that Q(x, y) = 0, we apply the convention that 0 log 0 = 0. Consider the case Q ⋆ , the projection of R onto Π X . WriteKL(Q∥R) = x∈X y∈Y Q(x, y) log Q Y |X (y|x)Q X (x) R Y |X (y|x)R X (x) = x∈X Q X (x)   y∈Y Q Y |X (y|x) log Q Y |X (y|x)Q X (x) R Y |X (y|x)R X (x)   = x∈X Q X (x)   y∈Y Q Y |X (y|x) log Q Y |X (y|x) R Y |X (y|x) + y∈Y Q Y |X (y|x) log Q X (x) R X (x)   = x∈X Q X (x)   y∈Y Q Y |X (y|x) log Q Y |X (y|x) R Y |X (y|x)   + x∈X Q X (x) log Q X (x) R X (x) = x∈X Q X (x) KL(Q Y |X (•|x)∥R Y |X (•|x)) + KL(Q X ∥R X ) = x∈X P X (x) KL(Q Y |X (•|x)∥R Y |X (•|x)) + KL(P X ∥R X ),where the last line is due to the marginal constraint Q ∈ Π X . For the above to be well defined, we need that P X ≪ R X so that KL(P X ∥R X ) < +∞. The above is minimized whenQ Y |X (y|x) = R Y |X (y|x) for all (x, y) ∈ X × Y such that Q X (x) = P X (x) > 0.The case of P ⋆ follows analogously when using that P Y ≪ R Y . Then, it holds thatEQUATIONandEQUATION)Proof. In the case that R(x, y) = 0, we apply the convention that 0 log 0 = 0. Note that minimizing KL(R∥Q) over Q is equivalent to minimizingx,y R(x, y) log Q(x, y) (i.e. the cross entropy).Consider the case Q ⋆ , the projection of R onto Π X . Because R ≪ Q for KL(R∥Q) < +∞ to hold, we have that R(x) > 0 =⇒ Q(x) > 0, so that Q Y |X (y|x) is well-defined. Write - x,y R(x, y) log Q(x, y) = - x∈X R X (x) log Q X (x) - x∈X R(x) y∈Y R Y |X (y|x) log Q Y |X (y|x) = - x∈X R X (x) log P X (x) + x∈X R X (x)   - y∈Y R Y |X (y|x) log Q Y |X (y|x)   .The second first term does not depend on Q due to the marginal constraint Q ∈ Π X . The second term is the expectation of the cross entropy fromR Y |X to Q Y |X over R X , which is minimized if R Y |X = Q Y |X . We have specified Q Y |X and Q X , completing the proof. C.1.3 Projection in χ 2 -DivergenceLet 1 denote the function that is identically equal to 1. Consider the following optimization problem, which is the subject of the subsequent lemmas:EQUATIONwhereA X :=    f : X × Y → R satisfying y∈Y f (x, y)R(x, y) = P X (x) for any x ∈ X    .Lemma 10. Assume that P X ≪ R X , and define The problem (25) is feasible, and its solution can be written asξ ⋆ = C R X (1 -f ) + f for any f ∈ L 2 (R), where the linear operator C R X is specified by [C R X g](x, y) = g(x, y) - y ′ ∈Y g(x, y ′ )R Y |X (y ′ |x).Proof. First, we establish feasibility by lettingf (x, y) := P X (x)/R X (x) if R X (x) > 0 1 otherwise .This function does not depend on the second input y. Because we assumed that P X ≪ R X , we have that the terms of f (x, y) for which R X (x) = 0 do not affect whether y∈Y f (x, y)R(x, y) = P X (x), because P X (x) = 0 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 it to a subspace projection problem. Indeed, consider f 1 , . . . , f r ∈ A X , and α 1 , . . . , α r ∈ R such thatr j=1 α j = 1. Then, y∈Y   r j=1 α j f j (x, y)   • R(x, y) = r j=1 α j   y∈Y f j (x, y)R(x, y)   = P X (x),indicating that r j=1 α j f j (x, y) ∈ A X and A X is an affine subset of L 2 (R). DefineS X :=    g : X × Y → R satisfying y∈Y g(x, y)R(x, y) = 0 for any x ∈ X    .Then, for any f ∈ A X , we have that g ∈ S X if and only if g + f ∈ A X . Taking any f ∈ A X , letting ϕ ⋆ be the solution ofEQUATIONwe will have that ϕ ⋆ + f will be the solution of ( 25). The remainder of the proof is showing that ϕ⋆ = C R X (1 -f ). First, define the operator µ R X by [µ X g](x, y) = y ′ ∈Y g(x, y ′ )R Y |X (y ′ |x), and note (by factoring out R X (x)) that g ∈ S X if and only if µ R X g = 0. In addition, µ R X g is linear and idempotent as µ R X µ R X g = µ R X g, so it is a projection operator in L 2 (R). Thus, S X is the orthogonal complement of range(µ R X ), and the solution of ( 26) is given by (EQUATION)and let ξ ⋆ be the solution of problem (25). Then,EQUATION)Proof. First, by reparametrizing the problem (27) as finding ξ such that Q(x, y) = ξ(x, y)R(x, y), we can compute its solution by solvingEQUATIONNotice that we also have a non-negativity constraint, as opposed to (25). If ξ ⋆ solves (25) and happens to be non-negative, then we have that ξ ⋆ solves (29) as well and the first equality of ( 28) is satisfied by definition. We show the second equality of (28) by direct computation, which also establishes the non-negativity of ξ ⋆ simultaneously.Apply Lem. 10 withf (x, y) := P X (x)/R X (x) if R X (x) > 0 1 otherwise . so that ξ ⋆ (x, y) = C R X (1 -f ) (x, y) + f (x, y) = z∈Y f (x, z)R Y |X (z|x) -f (x, y) + f (x, y) = f (x, y ′ )for any y ′ ∈ Y. Thus, the likelihood ratio of Q ⋆ with respect to R is a marginal reweighting. Accordingly,Q ⋆ (x, y) = ξ ⋆ (x, y)R(x, y) = P X (x)R Y |X (y|x) if R X (x) > 0 0 if R X (x) = 0 ,completing the proof.Proposition 12. Assume that P X ≪ R X and P Y ≪ R Y . DefineEQUATIONThen, it holds thatEQUATION)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. 11 with (X, x) and (Y, y) swapped.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 thatKL(P (0) n,X ∥P X ) ≥ KL(P Y ∥P (1) n,Y ) ≥ KL(P (2) n,X ∥P X ) ≥ . . .We give the following result for X and the analogous claim holds on Y. Proposition 14. Assume that P n,X (x) > 0 for all x ∈ X . It holds thatmax x∈X P X (x) P (k-1) n,X (x) -1 ≤ max{n -1, 1} if k = 1 max{1/p 2 ⋆ -1, 1} if k > 1.(32)In addition, we have thatmax x∈X P X (x) P (k-1) n,X (x) -1 ≤    n 1 2 KL(P n,X ∥P X ) if k = 1 1 p 2 ⋆ 1 2 KL(P n,X ∥P X ) if k > 1 .Moreover, when KL(P n,X ∥P X ) ≤ p 2 ⋆ /2, we haveEQUATIONProof. We first show thatP (k-1) n (x) ≥ 1/n for k = 1 and P (k-1) n (x) ≥ p 2 ⋆ for k > 1.In the case that k = 1, the result follows directly from the event S. For k > 1 such that k is odd, we have that for x ∈ X ,P (k-1) n,X (x) = y∈Y P (k-1) n (x, y) = y∈Y P Y (y) P (k-2)n,Y (y)P (k-2) n (x, y) ≥ p ⋆ y∈Y P (k-2) n (x, y) = p ⋆ P (k-2) n,X (x) = p ⋆ P X (x) ≥ p 2 ⋆ .The result for k 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 ∈ X ,EQUATION-1 = max P X (x)P (k-1) n,X (x) -1, 1 - P X (x) P (k-1) n,X (x) ≤ max{n -1, 1} if k = 1 max{1/p 2 ⋆ -1, 1} if k > 1, which is the desired result for the first inequality.Proving the second and third inequalities. Consider an odd k ≥ 1. By the definition of total variation distance, it holds that max x∈X P X (x) -P (k-1) n,X (x) ≤ TV(P (k-1) n,X , P X ).According to Pinsker's inequality, we have that TV(P (k-1) n,X , P X ) ≤ 1 2 KL(P (k-1) n,X ∥P X ), and so we have thatmax x∈X P X (x) -P (k-1) n,X (x) ≤ 1 2 KL(P (k-1) n,X ∥P X ) ≤ 1 2 KL(P (0) n,X ∥P X ),where the last inequality follows by the monotonicity of Sinkhorn iterations given in Prop. 13. We apply the lower bounds to writemax x∈X P X (x) P (k-1) n,X (x) -1 ≤    n 1 2 KL(P n,X ∥P X ) if k = 1 1 p 2 ⋆ 1 2 KL(P n,X ∥P X ) if k > 1 .Finally, when 1 2 KL(P n,X ∥P X ) ≤ p ⋆ /2, we have that max x∈X P X (x) -P (k-1) n,X (x) ≤ p ⋆ /2 and thusmin x∈X P (k-1) n,X (x) ≥ min x∈X P X (x) -max x∈X P (k-1) n,X (x) -P X (x) ≥ p ⋆ 2 .Hence,max x∈X P X (x) P (k-1) n,X (x) -1 ≤ max x∈X P (k-1) n,X (x) -P X (x) min x∈X P (k-1) n,X (x) ≤ 2 p ⋆ 1 2 KL(P n,X ∥P X ).Now, for k even, set k = 2t for t ≥ 0. We have thatmax y∈Y P (2t-1) n,Y (y) -P Y (y) ≤ TV(P (2t-1) n,Y , P Y ) ≤ 1 2 KL(P Y ∥P (2t-1) n,Y ).Invoke Prop. 13 once again to achieve1 2 KL(P Y ∥P (2t-1) n,Y ) ≤ 1 2 KL(P n,X ∥P X ),which completes the proof.This section contains the proof of the main result, namely Thm. 1. We first introduce some additional notation and then give a broad outline of the proof for readability. Let the expectation of a function h under a probability measure Q on X × Y by denoted byEQUATIONEQUATIONRecalling in addition that C k = C X for k odd and C k = C Y for k even. Finally, the eventEQUATIONis used as a condition in many results.Proof Outline. We first establish that the recursion formula[P (k) n -P ](h) = [P (k-1) n -P ](C k h) + V (k-1) n (C k h)holds in Prop. 15. Applying this result repeatedly to the balanced sequence (P (k) n ) k≥1 and unrolling the recursion, we see that when k is odd,EQUATIONAdditionally, let h ℓ,k := C ℓ . . . C k h, so that the first-order term can be written as P (0) n (h 1,k ) -P (h 1,k ) higher-order term can also be written ask ℓ=1 V (ℓ-1) n (h ℓ,k ).Because our original goal is to upper bound the mean squared error, we use the expansion above to writeE |P (k) n (h) -P (h)| 2 ≤ E |P (0) n (h 1,k ) -P (h 1,k )| 2 + 2E |P (0) n (h 1,k ) -P (h 1,k )| k ℓ=1 V (ℓ-1) n (h ℓ,k ) + E k ℓ=1 V (ℓ-1) n (h ℓ,k ) 2Regarding the first term, we have thatE |P (0) n (h 1,k ) -P (h 1,k )| 2 = σ 2 k /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 (0) n (h 1,k ) -P (h 1,k )| is O(n -1/2 ) with high probability and that | k ℓ=1 V (ℓ-1) n (h ℓ,k )| is in fact O(n -1) with high probability. As stated in Sec. 3, a key intermediate result in controlling the higher-order term is Prop. 14, whose proof is given in Appx. C. The remaining subsections walk through these steps in detail.We first recall that the sequence (P (k) n ) k≥1 can be computed with the following formula:P (0)n (x, y) := P n (x, y) andEQUATIONProp. 15 establishes the conditions under which these steps are well-defined (i.e. P (k-1) n,X (x) > 0 and P (k-1) n,Y (y) > 0). LetEQUATION)Proposition 15. Let (P (k) n ) k≥1 , be a sequence computed according to (8). These iterations are well-defined under the event S, and for G (k) n defined in (44), it holds thatEQUATION)andEQUATIONProof. First, assume that P (k-1) n,X (x) > 0 and P (k-1) n,Y (y) > 0 for all x ∈ X and y ∈ Y 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:P (k) n (h) =x,y h(x, y)P (k) n (x, y) =x,y h(x, y) P X P (k-1) n,X (x)P (k-1) n (x, y) by ( 37) for k odd =x,y 1 • h(x, y)P (k-1) n (x, y) +x,yP X P (k-1) n,X (x) -1 • h(x, y)P (k-1) n (x, y) = P (k-1) n (h) + V (k-1) n (h),where we substitutedEQUATION)Arguing analogously for k even and subtracting P (h) on both sides, we have thatEQUATIONWe refer to this as the "uncentered" recursion., which proves (39).We can then establish the following "centered" recursion using the following decomposition in the case of k odd.[P (k) n -P ](h) = [P (k) n -P ](C X h) + [P (k) n -P ](µ X h) h = C X h + µ X h = [P (k-1) n -P ](C X h) + V (k-1) n (C X h) + [P (k) n -P ](µ X h) apply (42) to C X h = [P (k-1) n -P ](C X h) + V (k-1) n (C X h). P (k) n (µ X h) = P (µ X h)The last line follows because µ X h is only a function on X , and due to the definition of the marginal rebalancing iterations, P (k) n,X = P X . This gives the desired formula by substituting (34).We proceed to show that the iterations are well-defined. We will in fact show that P (k-1) n,X (x) > 0 and P (k-1) n,Y (y) > 0 for all x ∈ X and y ∈ Y. For k = 1, P (0) n,X (x) = P n,X (x) > 0 and P (0) n,Y (y) = P n,Y (y) > 0 for all x ∈ X and y ∈ Y this holds under the event S by assumption. We argue by induction that this holds for all k > 1. Assume that the claim is true for {1, . . . , k -1}, and that k is even. Then,P (k-1)n,X (x) = P X (x) > 0,P (k-1) n,Y (y) = x∈X P (k-1) n (x, y) = x∈X P X P (k-2) n,X (x)P (k-2) n (x, y) ≥ min x∈X P X P (k-2) n,X (x) • P (k-2) n,Y (y) > 0 as P (k-2)n,X (x) > 0 and P (k-2) n,Y (y) > 0 by the inductive hypothesis. Arguing analogously for k odd achieves the claim.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 ν be a discrete probability measure supported on m atoms. Let U 1 , . . . , U n i.i.d.∼ ν and ν n be the associated empirical measure. Then, we have for any ϵ > 0 thatP (KL(ν n ∥ν) ≥ ϵ) ≤ 2 -n(ϵ-m log(n+1) n ) .We then provide a result that counts the number of terms that appear when repeatedly centering via the operators C 1 , . . . , C k . This formalizes the patternC X = I -µ X C Y C X = I -µ X -µ Y + µ Y µ X C X C Y C X = I -µ X -µ Y + µ Y µ X + µ X µ Y -µ X µ Y µ X ,and so on. This will be useful when bounding h ℓ,k uniformly.Lemma 17. For any k ≥ 1 and ℓ ∈ {1, . . . , k},C ℓ . . . C k = I - (k-ℓ-1)/2 τ =0 (µ X µ Y ) τ µ X - (k-ℓ-1)/2 τ =0 (µ Y µ X ) τ µ Y + (k-ℓ)/2 τ =1 (µ X µ Y ) τ + (k-ℓ)/2 τ =1 (µ Y µ X ) τ + (-1) k-ℓ+1 µ ℓ . . . µ k ,where the sum j τ =i is 0 when i > j and is ⌊j⌋ τ =i when j is not an integer by convention.Proof. We prove the claim by backward induction on ℓ, for the case that k is odd. In the case ℓ = k, the claim holds because C k = I -µ k . Next, for any ℓ < k, assume that the stated result holds for {ℓ + 1, . . . , k}. Then, if ℓ is also odd (so that µ ℓ = µ X ),C ℓ . . . C k = C ℓ C ℓ+1 . . . C k = I - (k-ℓ-2)/2 τ =0 (µ X µ Y ) τ µ X - (k-ℓ-2)/2 τ =0 (µ Y µ X ) τ µ Y + (k-ℓ-1)/2 τ =1 (µ X µ Y ) τ + (k-ℓ-1)/2 τ =1 (µ Y µ X ) τ + µ Y . . . k-ℓ terms µ X -µ X + (k-ℓ-2)/2 τ =0 (µ X µ Y ) τ µ X + (k-ℓ-2)/2 τ =0 µ X (µ Y µ X ) τ µ Y - (k-ℓ-1)/2 τ =1 (µ X µ Y ) τ - (k-ℓ-1)/2 τ =1 µ X (µ Y µ X ) τ -(µ X µ Y ) (k-ℓ)/2 µ XThe red terms and blue terms cancel out to zero. This leavesC ℓ . . . C k = I - (k-ℓ-2)/2 τ =0 (µ X µ Y ) τ µ X - (k-ℓ-2)/2 τ =0 (µ Y µ X ) τ µ Y + (k-ℓ-1)/2 τ =1 (µ Y µ X ) τ + (µ Y µ X ) (k-ℓ)/2 + (k-ℓ-2)/2 τ =0 µ X (µ Y µ X ) τ µ Y + (-1) k-ℓ+1 µ ℓ . . . µ kwherein we combine the red terms and re-index the blue terms to getC ℓ . . . C k = I - (k-ℓ-2)/2 τ =0 (µ X µ Y ) τ µ X - (k-ℓ-2)/2 τ =0 (µ Y µ X ) τ µ Y + (k-ℓ)/2 τ =1 (µ Y µ X ) τ + (k-ℓ)/2 τ =1 (µ X µ Y ) τ + (-1) k-ℓ+1 µ ℓ . . . µ k .Finally, because k -ℓ is even when k is odd and ℓ is odd, we can set the upper bound of the first two sums to (k -ℓ -1)/2 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 k and ℓ ≤ k. Similar arguments can be used for the case of k even and ℓ ≤ k.Returning to the outline at the start of this section, we may now bound the higher-order remainder term in (36), namelyk ℓ=1 V (ℓ-1) n (h ℓ,k ) = k ℓ=1 V (ℓ-1) n (C ℓ . . . C k h),depends on controlling the quantity V (k-1) n in the summation, which we recall for convenience:EQUATION)Because we have established uniform control over the functions P X /P (k-1) n,X -1 and P Y /P (k-1) n,Y -1, via Prop. 14 in Appx. C we can now bound the full remainder in Prop. 20.We also make use of the following intermediate result, which controls how large the ℓ ∞ -norm of the function h can grow after centering.Lemma 18. ∥h ℓ,k ∥ ∞ ≤ 2(k -ℓ + 1) ∥h∥ ∞ .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 -ℓ)/2 or (k -ℓ + 1)/2 will be a fraction. This yields 2(k -ℓ + 1) terms total, the desired result.We upper bound the sum in Prop. 20. To do so, we introduce some notation. Consider B 1 and B 2 defined byB 1 := M 1 and B 2 := max 2≤ℓ≤k M ℓ for M ℓ :=        max x∈X P X (x) P (ℓ-1) n,X (x) -1 ℓ odd max y∈Y P Y (y) P (ℓ-1) n,Y (y) -1 ℓ even for k ≥ 1.We also enumerate the sample spaces as X = {x 1 , . . . , x m } and Y = {y 1 , . . . , y m }, and define the function1 jk (x, y) := 1 {x = x j } k odd 1 {y = y j } k even .This is an indicator function on the j-th element of either X or Y depending on whether k is odd or even. Finally, for any function h, use (under the event S) recall the empirical process notationEQUATION)Using this notation, we can rewrite the recursion in terms of the quantity G (k) n (h) itself. This is established in the following lemma. Lemma 19. For k odd, it holds thatG (k) n (h) = G (k-1) n (C X h) + m j=1 P X (x j ) P (k-1) n,X (x j ) -1 G (k-1) n (C X h 1 jk ),whereas for k even, it holds thatG (k) n (h) = G (k-1) n (C Y h) + m j=1 P Y (y j ) P (k-1) n,Y (y j ) -1 G (k-1) n (C Y h 1 jk ),Proof. We give the proof for k odd. By (40) from Prop. 15 and by the definition of G (k) n (h), we need only show that P (C X h1 jk ) = 0. Indeed,E [C X h1 jk |X] (x) = E [C X h|X] (x j ) if x = x j 0 if x ̸ = x j . But E [C X h|X] (x j ) = 0 by definition of C X .Taking an expectation over P X gives that P (C X h1 jk ) = 0, which implies the desired result. The proof for k even follows symmetrically.The higher-order term in (36), can be bounded using Prop. 20. Proposition 20. For any k ≥ 1, the following holds under the event S:√ n k ℓ=1 |V (ℓ-1) n (C ℓ . . . C k h)| ≤ m j=1 B 1 |G (0) n (h 1,k 1 jℓ )| + B 2 k ℓ=2 |G (0) n (h ℓ,k 1 jℓ )| + mB 2 ∥h∥ ∞ √ nk(k -1)[B 1 + B 2 (k + 1)/3].Proof. First, for any ℓ ∈ {1, . . . , k}, recall the notation h ℓ,k := C ℓ . . . C k h. By (39) from Prop. 15 and by Lem. 19, we have that for ℓ odd,EQUATION)Using the statement above, we have that√ n |V (ℓ-1) n (h ℓ,k )| ≤ M ℓ m j=1 |G (ℓ-1) n (h ℓ,k 1 jℓ )| .The bound above holds for ℓ even as well. Then, using the (39) from Prop. 15 again, we have that for ℓ ≥ 2,[P (ℓ-1) n -P ](h ℓ,k 1 jℓ ) = [P (ℓ-2) n -P ](h ℓ,k 1 jℓ ) + V (ℓ-2) n (h ℓ,k 1 jℓ )which implies thatEQUATIONby Lem. 18 and M 1 ≤ B 1 and M ℓ ≤ B 2 for ℓ ≥ 2. Summing these bounds, we have that√ n k ℓ=1 |V (ℓ-1) n (h ℓ,k )| ≤ M 1 m j=1 |G (0) n (h 1,k 1 jℓ )| + k ℓ=2 M ℓ m j=1 |G (ℓ-1) n (h ℓ,k 1 jℓ )| ≤ B 1 m j=1 |G (0) n (h 1,k 1 jℓ )| + B 2 k ℓ=2 m j=1 |G (ℓ-1) n (h ℓ,k 1 jℓ )| ≤ B 1 m j=1 |G (0) n (h 1,k 1 jℓ )| + B 2 k ℓ=2 m j=1 |G (0) n (h ℓ,k 1 jℓ )| + 2 ∥h∥ ∞ √ n [B 1 + B 2 (ℓ -1)] (k -ℓ + 1) apply (46) = m j=1 B 1 |G (0) n (h 1,k 1 jℓ )| + B 2 k ℓ=2 |G (0) n (h ℓ,k 1 jℓ )| + 2mB 2 ∥h∥ ∞ √ n k ℓ=2 [B 1 + B 2 (ℓ -1)] (k -ℓ + 1), because |X | = m.We sum the last term:k ℓ=2 [B 1 + B 2 (ℓ -1)] (k -ℓ + 1) = B 1 k-1 ℓ=1 (k -ℓ) + B 2 k-1 ℓ=1 ℓ(k -ℓ) = k(k -1) 2 [B 1 + B 2 (k + 1)/3] .completing the proof.We can now show the main result of this section: the bound on the mean squared error of the rebalanced estimator. Recall the event S := {Supp(P n,X ) = Supp(P X ) and Supp(P n,Y ) = Supp(P Y )} (47) as introduced in (35). To remind the reader of the high-level steps of the proof, we may decompose the error on the event S we used the estimatorφ(k) n := φ (k) n 1 S + φ (0) n 1 S cso we decompose on the event S to writeE P Pn (k) (h) -P (h) 2 = E P (P n (h) -P (h)) 2 1 S c + E P (P (k) n (h) -P (h)) 2 1 S . (48)Then, we use the upcoming Prop. 21 to bound the first term, which will in turn require showing that S occurs with high probability. As for the second term, we will apply Prop. 15 and the derivation (36) to writeEQUATION)forEQUATION)By definition, we have thatE P T 2 1 1 S ≤ E P T 2 1 = σ 2 k /n.It then remains to bound the cross term E P [T 1 T 2 1 S ] and squared term E P T 2 2 1 S . This is accomplished by Lem. 23 and Lem. 22, respectively. Proposition 21. It holds that P (S c ) ≤ 2m(1 -p ⋆ ) n . Moreover, for any δ ∈ (0, 1), we haveE P (P n (h) -P (h)) 2 1 S c ≤ 4 ∥h∥ 2 ∞ min {2m(1 -p ⋆ ) n , δ} + 2 log(2/δ) n ∥h∥ 2 ∞ 2m(1 -p ⋆ ) n .Proof. Define F X := {Supp(P n,X ) ̸ = Supp(P X )} andF Y := {Supp(P n,Y ) ̸ = Supp(P Y )}, so that S c = F X ∪ F Y . We first control the probability of F X . Let F j := {P n,X (x j ) = 0} for j ∈ [m].We then obtain F X = ∪ m j=1 F j , which implies by the union bound thatP (F X ) ≤ m j=1 P (F j ) = m j=1 (1 -P X (x j )) n ≤ m(1 -p ⋆ ) n .Similarly, we have that P (F Y ) ≤ m(1 -p ⋆ ) n and thus P (S c ) ≤ 2m(1 -p ⋆ ) n , which gives the first claim.To control the expectation, consider any δ > 0, and define the eventE δ := P (0) n (h) -P (h) ≤ 2 log (2/δ) n ∥h∥ ∞ .By Hoeffding's inequality, it holds that P (E δ ) ≥ 1 -δ. Furthermore, we getE[1 S c (P (0) n (h) -P (h)) 2 ] = E[1 S c 1 E c δ (P (0) n (h) -P (h)) 2 ] + E[1 S c 1 E δ (P (0) n (h) -P (h)) 2 ] ≤ 4 ∥h∥ 2 ∞ E[1 S c 1 E c δ ] + 2 log (2/δ) n ∥h∥ 2 ∞ E[1 S c 1 E δ ] ≤ 4 ∥h∥ 2 ∞ min{P (S c ), P (E c δ )} + 2 log (2/δ) n ∥h∥ 2 ∞ P (S c ) ≤ 4 ∥h∥ 2 ∞ min{2m(1 -p ⋆ ) n , δ} + 2 log (2/δ) n ∥h∥ 2 ∞ 2m(1 -p ⋆ ) n .In order to bound the terms appearing in (49), we introduce the events E δ 1 , E δ 2 , and E δ 3 , defined byE δ 1 := max {KL(P n,X ∥P X ), KL(P n,Y ∥P Y )} ≤ 1 n log 2 2 δ + m log(n + 1) n F δ ℓ := |G (0) n (h ℓ,k 1 jℓ )| ≤ 2 log(2mk/δ)2(k -ℓ + 1) ∥h∥ ∞ , ℓ = 1, . . . , k, j = 1, . . . , m E δ 2 := k ℓ=1 F δ ℓ E δ 3 := |G (0) n (h 1,k )| ≤ 2 log(2/δ)2k ∥h∥ ∞ .The events are constructed such that P(E δ 1 ) ≥ 1 -δ, P(E δ 2 ) ≥ 1 -δ, and P(E δ 3 ) ≥ 1 -δ, as we used in the upcoming proofs of Lem. 23, Lem. 22, and Thm. 24. Lemma 22 (Squared term bound). Let T 2 be defined as in (50). For any δ > 0, assuming that n ≥ 2[log 2 (2/δ) + m log(n + 1)]/p 2 ⋆ , we have thatE P T 2 2 1 S ≤ 2 ∥h∥ 2 ∞ m 2 k 2 p 2 ⋆ [log 2 (2/δ) + m log(n + 1)] 2-1{k=1} ×   4n + k -1 p 2 ⋆ n + 2 + k + 1 p 2 ⋆ 2 δ + 8 n 2 2 log 2mk δ (k + 1) + (k -1)(k + 4) p 2 ⋆ 2   .Proof. The following computations are done under the event S. First, apply Prop. 20 to writeEQUATION)We decompose on the event E δ 1 ∩ E δ 2 . Note that by Thm. 16, we have that P(E δ 1 ) ≥ 1 -δ. It follows from Hoeffding's inequality, the union bound, and boundedness of ∥h ℓ,k 1 jℓ ∥ by Lem. 18 thatP(E δ 2 ) ≥ 1 -δ As a result, P(E δ 1 ∩ E δ 2 ) ≥ 1 -2δ. Bound |T 2 | under the event S\(E δ 1 ∩ E δ 2 ).In this case, we apply (32) from Prop. 14 to get B 1 ≤ n and B 2 ≤ 1/p 2 ⋆ , along with the universal bounds from Lem. 18:1 √ n |G (0) n (h 1,k 1 jℓ )| ≤ 2 ∥h 1,k ∥ ∞ ≤ 4k ∥h∥ ∞ 1 √ n k ℓ=2 |G (0) n (h ℓ,k 1 jℓ )| ≤ 2 k ℓ=2 ∥h ℓ,k ∥ ∞ ≤ k ℓ=2 4(k -ℓ + 1) ∥h∥ ∞ = 2k(k -1) ∥h∥ ∞so that by plugging into (51),|T 2 | ≤ ∥h∥ ∞ mk 4n + k -1 p 2 ⋆ n + 2 + k + 1 3p 2 ⋆ ,and in turn,EQUATION)Bound |T 2 | under the event S ∩ E δ 1 ∩ E δ 2 .In this case, we may use that n ≥ 2[log 2 (2/δ) + m log(n + 1)]/p 2 ⋆ apply (33) from Prop. 14 to getmax {B 1 , B 2 } ≤ 2 p ⋆ 1 2 KL(P n,X ∥P X ) ≤ 1 p ⋆ √ n 2 log 2 (2/δ) + 2m log(n + 1)and the bounds based on E δ 2 which give|G (0) n (h 1,k 1 jℓ )| ≤ 2 log 2mk δ 2k ∥h∥ ∞ k ℓ=2 |G (0) n (h ℓ,k 1 jℓ )| ≤ k ℓ=2 2 log 2mk δ 2(k -ℓ + 1) ∥h∥ ∞ ≤ 2 log 2mk δ k(k -1) ∥h∥ ∞ ,By plugging into (51),EQUATIONIn turn,EQUATION)Combining together both ( 57) and ( 52) and using that [log 2 (2/δ) + 2m log(n + 1)] ≥ 1, we have thatE P T 2 2 1 S ≤ 2 ∥h∥ 2 ∞ m 2 k 2 p 2 ⋆ [log 2 (2/δ) + m log(n + 1)] 2-1{k=1} × 4n + k -1 p 2 ⋆ n + 2 + k + 1 p 2 ⋆ 2 δ + 8 n 2 2 log(2mk/δ)(k + 1) + (k -1)(k + 4) p 2 ⋆ 2 ,the result as desired.Lemma 23 (Cross term bound). Let T 1 and T 2 be defined as in (50). For any δ > 0, assuming that n ≥ 2[log 2 (2/δ) + m log(n + 1)]/p 2 ⋆ , we have thatE P [T 1 T 2 1 S ] ≤ 2mk 2 ∥h∥ 2 ∞ 2 log(2/δ) [log 2 (2/δ) + 2m log(n + 1)] 1-1{k=1}/2 p 2 ⋆ × p ⋆ 2 log (2mk/δ)(k + 1) + (k -1)(k + 4) n 3/2 + 6 4np 2 ⋆ + (k -1) n + 2 + k + 1 p 2 ⋆ δ ,Proof. The following computations are done under the event S. First, apply Prop. 20 to writeEQUATIONWe decompose on the event E δ 1 ∩E δ 2 ∩E δ 3 . Note that by Thm. 16 and that n ≥ log 2 (2/δ)+m log(n+1), we have that P(E δ 1 ) ≥ 1 -δ. It follows by Hoeffding's inequality and the union bound that P(E δ 2 ) ≥ 1 -δ. Similarly, we also have by Hoeffding's inequality that P(E δ 3 ) ≥ 1 -δ. As a result,P(E δ 1 ∩ E δ 2 ∩ E δ 3 ) ≥ 1 -3δ. Bound |T 1 T 2 | under the event S\(E δ 1 ∩ E δ 2 ∩ E δ 3 ).In this case, we apply (32) from Prop. 14 to get B 1 ≤ n and B 2 ≤ 1/p 2 ⋆ , along with the universal bounds from Lem. 18:1 √ n |G (0) n (h 1,k )| ≤ 2 ∥h 1,k ∥ ∞ ≤ 4k ∥h∥ ∞ 1 √ n |G (0) n (h 1,k 1 jℓ )| ≤ 2 ∥h 1,k ∥ ∞ ≤ 4k ∥h∥ ∞ 1 √ n k ℓ=2 |G (0) n (h ℓ,k 1 jℓ )| ≤ 2 k ℓ=2 ∥h ℓ,k ∥ ∞ ≤ k ℓ=2 4(k -ℓ + 1) ∥h∥ ∞ = 2k(k -1) ∥h∥ ∞ ,so that by plugging into (58),|T 1 T 2 | ≤ 4k 2 ∥h∥ 2 ∞ m 4n + k -1 p 2 ⋆ n + 2 + k + 1 3p 2 ⋆ ,and in turn,EQUATION)Bound |T 1 T 2 | under the event S ∩ E δ 1 ∩ E δ 2 ∩ E δ 3 .In this case, we may use that n ≥ 2[log 2 (2/δ) + m log(n + 1)]/p2 ⋆ apply (33) from Prop. 14 to getmax {B 1 , B 2 } ≤ 2 p ⋆ 1 2 KL(P n,X ∥P X ) ≤ 1 √ n 1 p ⋆ 2 log 2 (2/δ) + 2m log(n + 1)and the bounds based onE δ 2 ∩ E δ 2 ∩ E δ 3 which give |G (0) n (h 1,k )| ≤ 2 log(2/δ)2k ∥h∥ ∞ |G (0) n (h 1,k 1 jℓ )| ≤ 2 log(2mk/δ)2k ∥h∥ ∞ k ℓ=2 |G (0) n (h ℓ,k 1 jℓ )| ≤ k ℓ=2 2 log 2mk δ 2(k -ℓ + 1) ∥h∥ ∞ ≤ 2 log 2mk δ k(k -1) ∥h∥ ∞ ,By plugging into (58),|T 2 | ≤ m ∥h∥ ∞ 2 log(2mk/δ) [2 log 2 (2/δ) + 2m log(n + 1)] np ⋆ k(k + 1) + m ∥h∥ ∞ [2 log 2 (2/δ) + 2m log(n + 1)] 3np 2 ⋆ k(k -1)(k + 4) ≤ mk ∥h∥ ∞ [log 2 (2/δ) + 2m log(n + 1)] 1-1{k=1}/2 np 2 ⋆ × p ⋆ 2 log(2mk/δ)(k + 1) + (k -1)(k + 4) |T 1 T 2 | ≤ 2mk 2 ∥h∥ 2 ∞ 2 log(2/δ) [log 2 (2/δ) + 2m log(n + 1)] 1-1{k=1}/2 n 3/2 p 2 ⋆ × p ⋆ 2 log(2mk/δ)(k + 1) + (k -1)(k + 4) ,In turn,EQUATIONCombining together both ( 60) and ( 59) and using that [log 2 (2/δ) + 2m log(n + 1)] ≥ 1, we have thatE P [T 1 T 2 1 S ] ≤ 2mk 2 ∥h∥ 2 ∞ 2 log(2/δ) [log 2 (2/δ) + 2m log(n + 1)] 1-1{k=1}/2 p 2 ⋆ × p ⋆ 2 log(2mk/δ)(k + 1) + (k -1)(k + 4) n 3/2 + 6 4np 2 ⋆ + (k -1) n + 2 + k + 1 p 2 ⋆ δ ,the result as desired.We now combine the previous results to prove Thm. 24.Theorem 24. For a sequence of rebalanced distributions ( Pn (k) ) k≥1 , there exists an absolute constantEQUATIONwhereB = log (2n/p ⋆ )m 2 k 4 ∥h∥Proof. We apply the decomposition (48), and subsequently handle the second term using bounds on the terms in (49). Set δ = p 4 ⋆ /n 4 . We apply Lem. 22 and Lem. 23 with this choice of δ, so that there exists an absolute constants C, C 1 , and C 2 such thatE P [T 1 T 2 1 S ] ≤ C 1 ∥h∥ 2 ∞ m 2 k 3 log(2n/p ⋆ ) n 3/2 p 2 ⋆ [log 2 (2n/p ⋆ ) + m log(n + 1)] 1-1{k=1}/2 × log 2mnk p ⋆ + k -1 p 2 ⋆ E P T 2 2 1 S ≤ C 2 ∥h∥ 2 ∞ m 2 k 4 n 2 p 2 ⋆ [log 2 (2n/p ⋆ ) + m log(n + 1)] 2-1{k=1} × log 2mnk p ⋆ + (k -1) 2 p 2 ⋆ , when n ≥ C[log 2 (2n/p ⋆ ) + m log (n + 1)]/p 2 ⋆ .This then implies that there is an absolute constant C 3 such thatE P Pn (k) (h) -P (h) 2 ≤ E P (P (0) n (h) -P (h)) 2 1 S c + σ 2 k n + C 3 ∥h∥ 2 ∞ m 2 k 4 log(2n/p ⋆ ) n 3/2 p 2 ⋆ log 2 2n p ⋆ + m log(n + 1) 2-1{k=1} log 2mnk p ⋆ + (k -1) 2 p 2 ⋆ .Next, we apply Prop. 21 with the same choice of δ.Because 2[log 2 (2/δ)+m log(n+1)] ≥ log(m/δ) and -log(1 -p ⋆ ) ≥ p ⋆ ≥ p 2 ⋆ , we have that n ≥ log(δ/m)/ log(1 -p ⋆ ), which implies that m(1 -p ⋆ ) n ≤ δ.Combining with the display above, we have that there exists an absolute constant C > 0 such thatE P Pn (k) (h) -P (h) 2 ≤ σ 2 k n + C ∥h∥ 2 ∞ m 2 k 4 log(2n/p ⋆ ) n 3/2 p 2 ⋆ × [log 2 (2/δ) + m log(n + 1)] 2-1{k=1} log 2mnk p ⋆ + (k -1) 2 p 2 ⋆ ,which is the claimed result.While not shown in the main text, similar techniques to those used above can also control the bias of Pn (k) (h) as in Thm. 25. Interestingly, this bias is of order O(n -2 ) which confirms the intuition that even thought Pn (k) (h) may be biased, the dominant term is the variance.Theorem 25. For a sequence of rebalanced distributions (P (k) ) k≥1 , there exists an absolute constantEQUATION)where B is as defined in Thm. 24.Proof. First, apply the decomposition (48) so thatE P Pn (k) (h) -P (h) ≤ |E P [(P n (h) -P (h)) 1 S c ]| + |E P [(P (k) n (h) -P (h)) 1 S ]| .By using the argument of Prop. 21, we have that|E P [P n (h) -P (h)] 1 S c | ≤ 2 ∥h∥ ∞ min {2m(1 -p ⋆ ) n , δ} + 2 log(2/δ) n ∥h∥ ∞ 2m(1 -p ⋆ ) n .Then, by the recursion formula Equation (36), we have that√ n |E P [(P (k) n (h) -P (h)) 1 S ]| = |E P [G (k) n (h)1 S ]| = E P (1 -1 S c )G (0) n (C 1 . . . C k h) + √ n1 S k ℓ=1 V (ℓ-1) n (C ℓ . . . C k h) . Because G (0) n (C 1 . . . C k h) has zero mean, it follows that √ n |E P [(P (k) n (h) -P (h)) 1 S ]| ≤ |E P [1 S c G (0) n (C 1 . . . C k h)]| + √ n |E P [1 S T 2 ]|We have by Hoeffding's inequality that P(E δ 3 ) ≥ 1 -δ, and that by Lem. 18 thatG (0) n (C 1 . . . C k h) ≤ 4k√ n ∥h∥ ∞ universally. As a result, applying Prop. 21 once again,|E P [1 S c G (0) n (C 1 . . . C k h)]| ≤ E P 1 S c 1 E δ 3 G (0) n (C 1 . . . C k h) + E P 1 S c 1 E δ 3 G (0) n (C 1 . . . C k h) ≤ 4k √ n ∥h∥ ∞ min {2m(1 -p ⋆ ) n , δ} + 2 log(2/δ)2k ∥h∥ ∞ 2m(1 -p ⋆ ) n .Using a similar argument to Lem. 22, we have that under S\(E δ 1 ∩ E δ 2 ) (which occurs with probability no more than 2δ),|T 2 | ≤ ∥h∥ ∞ mk 4n + k -1 p 2 ⋆ n + 2 + k + 1 3p 2 ⋆ ,and that under S ∩ E δ 1 ∩ E δ 2 (which occurs with probability at least 1 -2δ),|T 2 | ≤ 4mk ∥h∥ ∞ [log 2 (2/δ) + 2m log(n + 1)] 1-1{k=1}/2 np 2 ⋆ p ⋆ 2 log(2mk/δ)(k + 1) + (k -1)(k + 4) .P [1 S T 2 ]| ≤ E P 1 S\(E δ 1 ∩E δ 2 ) T 2 + E P 1 S∩E δ 1 ∩E δ 2 T 2 and setting δ = p 2⋆ n 2 achieves the desired result.We now adapt the main results to cases in which the marginal distributions (P X , P Y ) are misspecified, in that the user is provided marginal distributions ( PX,ϵ , PY,ϵ ) which satisfy the following structure.Assumption 1. There exist fixed probability mass functions PX and PY for some ϵ ∈ [0, 1), PX,ϵ = (1 -ϵ)P X + ϵ PX and PY,ϵ = (1 -ϵ)P Y + ϵ PY .We also have the existence of the positive quantity p⋆ := min{min x PX (x), min y PY (y)} > 0.Given the existence of p⋆ > 0, we may also definep⋆,ϵ = min{min x PX,ϵ (x), min y PY,ϵ (y)} ≥ ϵp ⋆ + (1 -ϵ)p ⋆ > 0.To be precise, the iterations of balancing follow P (0) n = P n andEQUATION• P (k-1) n (x, y) k odd PY,ϵ (y)EQUATION•EQUATION)We start by deriving a result similar to Prop. 15. Since ϵ < 1, the (possibly misspecified) target marginals PX,ϵ (x) > 0 and PY,ϵ (y) > 0 for all x ∈ X and y ∈ Y. DefineV (k-1) n (h) :=        x,y PX,ϵ P (k-1) n,X (x) -1 h(x, y) P (k-1) n (x, y) k odd x,y PY,ϵ P (k-1) n,Y (y) -1 h(x, y) P (k-1) n (x, y) k even (64) and Ĝ(k) n (h) := √ n P (k) n (h) -P (h) .The format of this section will be to derive analogous results to the building blocks of the previous section. From that point, the computations from Appx. D.4 will achieve the desired result. For the sake of comparison to Thm. 1 we consider error terms containing ϵ only by their dependence on (ϵ, k, n, p⋆,ϵ ).Proposition 26. Let ( P (k) n ) k≥1 be a sequence computed according to (63). Definec 2 = max χ 2 ( PX ∥P X ), χ 2 ( PY ∥P Y ) .These iterations are well-defined under the event S, and for G (k) n defined in (44), it holds that Ĝ(k)EQUATIONandEQUATIONEQUATIONProof. The proof that P (k-1) n,X (x) > 0 and P (k-1) n,Y (y) > 0 for all x ∈ X and y ∈ 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: Subtracting P (h) on both sides, we have thatP (k) n (h) =EQUATIONThis proves the uncentered recursion formula given in (65). We then show the centered version.[P (k) n -P ](h) = [ P (k) n -P ](C X h) + [ P (k) n -P ](µ X h) = [ P (k) n -P ](C X h) + [ PX,ϵ -P X ](µ X h) = [ P (k-1) n -P ](C X h) + V (k-1) n (C X h) + [ PX,ϵ -P X ](µ X h).Next, we bound the additional error term. By the Cauchy-Schwarz inequality,[ PX,ϵ -P X ](µ X h) ≤ PX,ϵ P X -1 L 2 (P X ) • ∥µ X h∥ L 2 (P X ) = χ 2 ( PX,ϵ ∥P X ) • ∥µ X h∥ L 2 (P X ) ≤ χ 2 ( PX,ϵ ∥P X ) • ∥h∥ L 2 (P ) ,as µ X is an orthogonal projection in L 2 (P ). Using convexity of f -divergences, we have thatχ 2 ( PX,ϵ ∥P X ) ≤ ϵχ 2 ( PX ∥P X ) + (1 -ϵ)χ 2 (P X ∥P X ) = ϵχ 2 ( PX ∥P X ).This achieves the desired result.Using similar ideas, we then prove an analog of Lem. 19.Lemma 27. For k odd, it holds that√ n V (k-1) n (C X h) = m j=1 PX,ϵ P (k-1) n,X (x j ) -1 Ĝ(k-1) n (C X h 1 jk ),whereas for k even, it holds that√ n V (k-1) n (C Y h) = m j=1 PY,ϵ P (k-1) n,Y (x j ) -1 Ĝ(k-1) n (C Y h 1 jk ).Proof. We give the proof for k odd. We claim that we need only show that P (C X h1 jk ) = 0. This would show that√ n V (k-1) n (C X h) = x,y PX,ϵ P (k-1) n,X (x) -1 [C X h](x, y) P (k-1) n (x, y) = √ n m j=1 x,y PX,ϵ P (k-1) n,X (x) -1 [C X h 1 jk ](x, y) P (k-1) n (x, y) = √ n m j=1 x,y PX,ϵ P (k-1) n,X (x j ) -1 [C X h 1 jk ](x, y) P (k-1) n (x, y) = m j=1 PX,ϵ P (k-1) n,X (x j ) -1 √ n x,y [C X h 1 jk ](x, y) P (k-1) n (x, y) = m j=1 PX,ϵ P (k-1) n,X (x j ) -1 Ĝ(k-1) n (C X h 1 jk ),where P (C X h1 jk ) = 0 is employed in the last step. Now the result follows from (66) in Prop. 26 and the definition of Ĝ(k) n (h). To prove the claim, as in Lem. 19, writeE [C X h1 jk |X] (x) = E [C X h|X] (x j ) if x = x j 0 if x ̸ = x j .But E [C X h|X] (x j ) = 0 by definition of C X . Taking an expectation over P X gives that P (C X h1 jk ) = 0, which implies the desired result. The proof for k even follows symmetrically.For the remainder of the argument, we see that (67) can be unrolled so thatEQUATIONwhere we use that G (0) n = Ĝ(0) n .Next, we need to bound V (ℓ-1) n (C ℓ . . . C k h) , 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 P n,X (x) > 0 for all x ∈ X . It holds that max x∈X PX,ϵ (x)P (k-1) n,X (x) -1 ≤ max{n -1, 1} if k = 1 max{1/p 2 ⋆,ϵ -1, 1} if k > 1.(70)In addition, we have thatEQUATION, we haveEQUATIONProof. First, observe that P (0) n,X (x) = P (0) n,X (x) ≥ 1/n under the event S. For k > 1 such that k is odd, we have that for x ∈ X ,P (k-1) n,X (x) = y∈Y P (k-1) n (x, y) = y∈Y PY,ϵ (y) P (k-2)n,Y (y)P (k-2) n (x, y) ≥ p⋆,ϵ y∈Y P (k-2) n (x, y) = p⋆,ϵ P (k-2) n,X (x) = p⋆,ϵ PX,ϵ (x) ≥ p2 ⋆,ϵ .The result for k 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 x ∈ X , PX,ϵ (x)EQUATION-1 = max PX,ϵ (x)EQUATION-1, 1 -PX,ϵ (x)P (k-1) n,X (x) ≤ max{n -1, 1} if k = 1 max{1/p 2 ⋆,ϵ -1, 1} if k > 1, which is the desired result.Proving the second and third inequalities. Consider an odd k ≥ 1. By the definition of total variation distance, it holds that max x∈X PX,ϵ (x) -P (k-1) n,X (x) ≤ TV( P (k-1) n,X , PX,ϵ ).According to Pinsker's inequality, we have that TV( P (k-1) n,X , PX,ϵ ) ≤ 1 2 KL( P (k-1) n,X ∥ PX,ϵ ), and so we have thatmax x∈X PX,ϵ (x) -P (k-1) n,X (x) ≤ 1 2 KL( P (k-1) n,X ∥ PX,ϵ ) ≤ 1 2 KL(P (0) n,X ∥ PX,ϵ ),where the last inequality follows by the monotonicity of Sinkhorn iterations given in Prop. 13. Notice that the remaining term is KL(P (0) n,X ∥ PX,ϵ ) = KL(P n,X ∥ PX,ϵ ), which may not decay to zero as n → ∞. Because ϵ < 1, writeKL(P n,X ∥ PX,ϵ ) = x∈X P n,X (x) log P n,X (x) (1 -ϵ)P X (x) + ϵ PX (x) ≤ x∈X P n,X (x) log P n,X (x) (1 -ϵ)P X (x) = KL(P n,X ∥P X ) + log 1 1 -ϵ =⇒ 1 2 KL(P n,X ∥ PX,ϵ ) ≤ 1 2 KL(P n,X ∥P X ) + 1 2 log 1 1 -ϵ .We can then apply the lower boundsmax x∈X PX,ϵ (x) P (k-1) n,X (x) -1 ≤    n 1 2 KL(P n,X ∥P X ) + 1 2 log 1 1-ϵ if k = 1 1 p2 ⋆,ϵ 1 2 KL(P n,X ∥P X ) + 1 2 log 1 1-ϵ if k > 1 .Finally, combining the arguments above, we have thatmax x∈X PX,ϵ (x) -P (k-1) n,X (x) ≤ 1 2 KL(P n,X ∥P X ) + 1 2 log 1 1 -ϵ ≤ p⋆,ϵ 4 + p⋆,ϵ 4 = p⋆,ϵ 2 ,where the last step invoked the assumption thatEQUATION. This means thatmin x∈X P (k-1) n,X (x) ≥ min x∈X PX,ϵ (x) -max x∈X P (k-1) n,X (x) -PX,ϵ (x) ≥ p⋆,ϵ 2 .Hence,max x∈X PX,ϵ (x) P (k-1) n,X (x) -1 ≤ max x∈X P (k-1) n,X (x) -PX,ϵ (x) min x∈X P (k-1) n,X (x) ≤ 2 p⋆,ϵ 1 2 KL(P n,X ∥ PX,ϵ ).Now, for k even, set k = 2t for t ≥ 0. We have thatmax y∈Y P (2t-1) n,Y (y) -PY,ϵ (y) ≤ TV( P (2t-1) n,Y , PY,ϵ ) ≤ 1 2 KL( PY,ϵ ∥ P (2t-1) n,Y ).Invoke Prop. 13 once again to achieve1 2 KL( PY,ϵ ∥ P (2t-1) n,Y ) ≤ 1 2 KL(P n,X ∥ PX,ϵ ) ≤ 1 2 KL(P n,X ∥P X ) + 1 2 log 1 1 -ϵ ,which completes the proof.Proceeding with similar steps, define the quantitiesEQUATION-1 ℓ odd max y∈Y PY,ϵ (y)EQUATION-1 ℓ even .We must now establish an analog of Prop. 20. Proposition 29. For any k ≥ 1, the following holds under the event S:√ n k ℓ=1 V (ℓ-1) n (C ℓ . . . C k h) ≤ m j=1 B1 |G (0) n (h 1,k 1 jℓ )| + B2 k ℓ=2 |G (0) n (h ℓ,k 1 jℓ )| + m B2 ∥h∥ ∞ √ nk(k -1)[ B1 + B2 (k + 1)/3].Proof. This proof largely follows the argument of Prop. 20, while accounting for the misspecified marginal error. Using again the notation h ℓ,k := C ℓ . . . C k h, it follows from Lem. 27 that, for odd ℓ,√ n V (ℓ-1) n (h ℓ,k ) = m j=1 PX,ϵ P (ℓ-1) n,X (x j ) -1 Ĝ(ℓ-1) n (h ℓ,k 1 jℓ ) ≤ Mℓ m j=1 Ĝ(ℓ-1) n (h ℓ,k 1 jℓ ) .The bound above holds for ℓ even as well. Then, using (65) from Prop. 26 along with the triangle inequality, we have that for ℓ ≥ 2,EQUATIONby Lem. 18 and M1 ≤ B1 and Mℓ ≤ B2 for ℓ ≥ 2. The bound above holds trivially for ℓ = 1. Summing these bounds over ℓ and j, we have that√ n k ℓ=1 V (ℓ-1) n (h ℓ,k ) ≤ M1 m j=1 |G (0) n (h 1,k 1 jℓ )| + k ℓ=2 Mℓ m j=1 Ĝ(ℓ-1) n (h ℓ,k 1 jℓ ) ≤ B1 m j=1 |G (0) n (h 1,k 1 jℓ )| + B2 k ℓ=2 m j=1 Ĝ(ℓ-1) n (h ℓ,k 1 jℓ ) ≤ B1 m j=1 |G (0) n (h 1,k 1 jℓ )| + B2 k ℓ=2 m j=1 Ĝ(0) n (h ℓ,k 1 jℓ ) + 2 ∥h∥ ∞ √ n B1 + B2 (ℓ -1) (k -ℓ + 1) apply (73) = m j=1 B1 |G (0) n (h 1,k 1 jℓ )| + B2 k ℓ=2 |G (0) n (h ℓ,k 1 jℓ )| + 2m B2 ∥h∥ ∞ √ n k ℓ=2 B1 + B2 (ℓ -1) (k -ℓ + 1), because |X | = m.We sum the last term:k ℓ=2 B1 + B2 (ℓ -1) (k -ℓ + 1) = B1 k-1 ℓ=1 (k -ℓ) + B2 k-1 ℓ=1 ℓ(k -ℓ) = k(k -1) 2 B1 + B2 (k + 1)/3 ,which completes the proof.Ultimately, we wish to construct an upper bound forEQUATIONFinally, decomposing on E δ 1 ∩ E δ 2 ∩ E δ 3 will result in a bound of the formE P T 1 T2 1 S ≤ 3δT c 1 T c 2 + T 1 T 2 + Õ (δT c 1 E c 2 + T 1 E 2 ) .The leading terms 2δ(T c 2 ) 2 + T 2 2 and 3δT c 1 T c 2 + T 1 T 2 from both bounds should have the exact same form as the terms in Lem. 22 and Lem. 23, with p ⋆ replaced by p⋆,ϵ , thus retaining the same dependence on (n, k). By setting δ = p4 ⋆,ϵ /n 4 , we will achieve a similar result to Thm. 24, i.e., thatEQUATIONEQUATIONIt remains to quantify the Õ terms by computing the order of the 6 constants (T 2 , E 2 , T c 2 , E c 2 , T 1 , T c 1 ). We follow similar steps to Lem. 22 and Lem. 23 to achieve this. Lemma 30.For δ = (p ⋆,ϵ /n) 4 , assume that n ≥ 8[log 2 (2/δ) + m log(n + 1)]/p 2 ⋆,ϵ and ϵ ≤ 1 -exp - p2 ⋆,ϵ 8 . Then, it holds that T c 2 = Õ k 2 p2 ⋆,ϵ n + k p2 ⋆,ϵ , E c 2 = 0 T 2 = Õ k 3 np 2 ⋆,ϵ , E 2 = Õ k 3 p2 ⋆,ϵ 1 n log 1 1-ϵ + log 1 1-ϵ .Proof. The following computations are done under the event S. First, apply Prop. 29 to writeEQUATIONWe decompose on the event E δ 1 ∩ E δ 2 .Bound |T 2 | under the event S\(E δ 1 ∩ E δ 2 ). In this case, we apply (70) from Prop. 28 to get B1 ≤ n and B2 ≤ 1/p 2 ⋆,ϵ , along with the universal bounds from Lem. 18:1 √ n |G (0) n (h 1,k 1 jℓ )| ≤ 2 ∥h 1,k ∥ ∞ ≤ 4k ∥h∥ ∞ 1 √ n k ℓ=2 |G (0) n (h ℓ,k 1 jℓ )| ≤ 2 k ℓ=2 ∥h ℓ,k ∥ ∞ ≤ k ℓ=2 4(k -ℓ + 1) ∥h∥ ∞ = 2k(k -1) ∥h∥ ∞ so that by plugging into (82), T2 ≤ ∥h∥ ∞ mk 4n + k -1 p2 ⋆,ϵ n + 2 + k + 1 3p 2 ⋆,ϵ T c 2 + 0 E c 2 . Bound |T 2 | under the event S ∩ E δ 1 ∩ E δ 2 .In this case, we may use that n ≥ 8/p 2 ⋆,ϵ (because [log 2 (2/δ) + m log(n + 1)] ≥ 1 for δ ∈ (0, 1)) and apply (71) from Prop. 28 to getEQUATIONThe bounds based on E δ 2 give|G (0) n (h 1,k 1 jℓ )| ≤ 2 log 2mk δ 2k ∥h∥ ∞ k ℓ=2 |G (0) n (h ℓ,k 1 jℓ )| ≤ k ℓ=2 2 log 2mk δ 2(k -ℓ + 1) ∥h∥ ∞ ≤ 2 log 2mk δ k(k -1) ∥h∥ ∞ .By plugging into (82), we can reuse the steps in the bound from (56) (for all terms without ϵ) to writeT2 ≤ 4mk ∥h∥ ∞ [log 2 (2/δ) + 2m log(n + 1)] 1-1{k=1}/2 np 2 ⋆,ϵ × p⋆,ϵ 2 log (2mk/δ)(k + 1) + (k -1)(k + 4) + E 2 , so that T 2 = 4mk ∥h∥ ∞ [log 2 (2/δ) + 2m log(n + 1)] 1-1{k=1}/2 np 2 ⋆,ϵ × p⋆,ϵ 2 log (2mk/δ)(k + 1) + (k -1)(k + 4) .We compute E 2 by using thatmax B1 , B2 ≤ O 1 p⋆,ϵ log 1 1-ϵ + Õ 1 p⋆,ϵ √ n |G (0) n (h 1,k 1 jℓ )| ≤ Õ (k) k ℓ=2 |G (0) n (h ℓ,k 1 jℓ )| ≤ Õ k 2 ,which givesE 2 = Õ k 3 p2 ⋆,ϵ 1 n log 1 1-ϵ + log 1 1-ϵ.We now make the corresponding argument for the term T 1 . Lemma 31. For δ = (p ⋆,ϵ /n) 4 , it holds thatT c 1 = Õ(k), T 1 = Õ k √ n .Proof. The following computations are done under the event S.Bound |T 1 | under the event S\E δ 3 . Here we simply apply a universal bound on the empirical process term:1 √ n |G (0) n (h 1,k )| ≤ 2 ∥h 1,k ∥ ∞ ≤ 4k ∥h∥ ∞ , so that T c 1 = 4k ∥h∥ ∞Bound |T 1 | under the event S ∩ E δ 3 . Now, we may use the definition of the event E δ 3 to achieve Optimizer. For optimization, models were trained with stochastic gradient descent (SGD) with the learning rate tuned along the grid 1 -3 , 3 -3 , 1 -2 , 3 -2 , 1 -1 and a fixed weight decay parameter of 0.01. Momentum-variants such as Adam [Kingma and Ba, 2015] were not used to isolate the effect varying losses as described in Sec. 4.1 √ n |G (0) n (h 1,k )| ≤ 2 log(2/δ) n 2k ∥h∥ ∞ = T 1 .Experiments were run on a CPU/GPU workstation 12 virtual cores, 126G of memory, and four NVIDIA TITAN Xp GPUs with 12G memory each. The code was written in Python 3 and we use PyTorch for automatic differentiation. The OpenCLIP and CLIP Benchmark repos were used for zero-shot evaluation.We considered in the contrastive learning example from Sec. 2 -see ( 6) in particular -a variant of the CLIP objective in which either zero, or one, or more than one balancing iterations are performed (see ( 6)), via optimizingEQUATIONThis contrasts the single-iteration variant L (1) n which in fact reduces to the original CLIP loss. Because these iterations are 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 L (0) n (the normalized initial measure, No Balancing), L (1) n (the original CLIP loss, CLIP balancing), and L (2) n (the two-iteration CLIP loss, Multi-CLIP balancing). We also vary the quality of the text encoder f θ T , 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.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 et al. [2024] . The target marginal P Y is selected by choosing a threshold for which frequent keywords have their probability mass truncated, and the probability measure is normalized to sum to one. In Fig. 4 , we show the observed marginal P n,Y and the target marginal P Y 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-balanced dataset (blue). We observe moderate improvement especially in the small batch regime (m = 512) when curating the dataset.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-Captions dataset [Fang et al., 2013] . As before, we use a fixed image/text encoder as a base vector representation and compose it with a trainable feed-forward n for k = 8) and a noise level on the provided marginals (see ( 84)). The y-axis shows the mean squared error of estimating a linear functional. The x-axis represents the dependence level s = s 2 (i.e. the leading singular value other than s 1 = 1). neural network, i.e., f θ = f head θ • f base , for θ = θ I (images) or θ = θ T (text). For the base text embeddings, we maintain three levels of model quality: GPT-2 [Radford et al., 2019] , BERT [Devlin et al., 2019] , and CLIP-based encodings.Baseline Comparisons. We present a synthetic data example to understand the role of the singular values s 2 , . . . , s m and compare our approach to simple baselines that make use of (P X , P Y ). We also consider misspecification of these target marginals, in that they are chosen by the user but are not the marginal distributions of the data-generating distribution P . First, while one can verify by hand that ( 14) is a distribution for which s 2 = s, we construct a more general example for m ≥ 2. We leave the construction to the end of this example. For controllable misspecification, we define ϵ ∈ [0, 0.5] to be the misspecification level, so that the corrupted target marginals are set to beEQUATIONwhere PX and PY are drawn independently and randomly from the Dirichlet(1 m ) distribution (i.e. uniformly over the probability simplex on m atoms). Finally, other than the empirical measure P n , we define one additional baseline; the importance weighted independently (IPWI) estimator is defined asP IPWI n (x, y) = PX (x)P n,X (x)P n,Y (y) P n (x, y).(85)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, as it may not sum to one. Observe the comparative performance in (see Fig. 5 ). We notice in particular that the naive P IPWI n is outperformed by empirical measure uniformly over s, as by applying both reweightings simultaneously, the estimator does not satisfy either marginal constraint. On the other hand, under the maximum amount of target marginal corruption (ϵ = 0.5), the balancing-based estimator suffers an approximately half-order of magnitude in MSE. When s ≈ 1, the MSE of the balancing estimator decreases significantly. We hypothesize that this is because the data sources X and Y are nearly a function of one another, and if this function is estimable to high precision by a small amount of data, then a single marginal can identify the entire joint distribution via pushforward calculations. That being said, it is important to note that the quantities u j and v j in (15) also depend on s, so it is difficult to control the singular values without controlling the respective bases.As for the construction of the probability mass function and test function, let I m and 1 m×m denote the identity matrix and matrix of ones in R m×m . For any s ∈ (0, 1) and m ≥ 1, consider the probability mass matrix P given byP = 1 m 1 m 1 m×m +s I m - 1 m 1 m×m .The eigenvalues of the first matrix in the squared brackets are (1, 0, . . . , 0), as it is a rank 1 matrix for which 1 m is an eigenvector. The second matrix in the square brackets is the centering matrix (the projection matrix that subtracts the mean of a vector's components from the entire vector). Multiplied by s, it has eigenvalues (0, s, . . . , s) where 0 is associated to the eigenvector 1 m . Thus, the matrix in its entirety has eigenvalues (1/m, s/m, . . . , s/m), where the scaling factor ensures that P sums to one. The relation (13) holds for this choice of P and uniform marginals, with s 2 = . . . = s m = s. Thus, by tuning s, we may control the level of dependence between X and Y . Finally, because X and Y are finite, we can also specify the test function h via an m × m table indexed by i (meaning x i ) and j (meaning y j ). We let h(x i , y j ) = |Z ij | where the Z ij are independently drawn from a standard normal distribution. The resulting mean squared error is estimated with 200 seeds at n = 300. n and R (0) n , where neither marginal is set to uniform. Center: Q (1) n and R (1) n , which corresponds to the original CLIP loss. Right:Q (2)n and R (2) n , which correspond to two iterations of the balancing procedure within the loss. The blue bars are slightly non-uniform. [W f θ I (x)] j , which can then be evaluated using standard accuracy metrics on a 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 et al., 2013] benchmarks are given in Fig. 7 .Empirical Marginals in CLIP Balancing. To further clarify how the iterative balancing procedure is baked into the CLIP losses, recall from (83) that the objectives decompose into two terms, which depend on Q (k) n and R (k) n which differ only based on whether balancing to fit P Y or to fit P X is applied first, respectively. Thus, for any model parameterized by θ and any number of iterations k, there are four marginal distributions of interest: Q (k) θ,X , Q (k) θ,Y , R (k) θ,X , and R (k) θ,Y . Based on the order of iterations, we have that Q (1) θ,Y = R (2) θ,Y = P Y , and R (1) θ,X = Q (2) θ,X = P X . This is illustrated in Fig. 8 . We see that after only a few iterations, both marginal distributions converge to the uniform distribution.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).We hide dependence on θ for notational simplicity.Due to this technical consideration, we define P(k) n to be the empirical measure Pn when this condition is not satisfied, which we show occurs with low-probability. See Appx. D.4 for details.∞ p 2 ⋆ log 2 2n p ⋆ + m log (n + 1) 2-1{k} log 2mkn p ⋆ + (k -1) 2 p 2 ⋆ . 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: I have three main concerns with this work:

1/ 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.

2/ 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?

3/ 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.
questions: questions: I have the following questions for the authors:

- 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)?
- p4, example 1: What would the target marginals correspond to? What is \psi_n^(k) and P_n^(k) here? 
- p4, example 2: What would the target marginals correspond to? What is \psi_n^(k) and P_n^(k) here?
- Why do we need \tilde{\psi}_n^{(k)} in (12) and how does it relate to \psi_n^{(k)}?
- 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?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1.the experiments are somewhat limited in scope
2.adding more visualizations or intuitive explanations may be better for understanding the key finding of the paper.
3.will the assumptions limit the applicability of the findings?
questions: questions: 1. Can the findings be extended to other areas?
2. 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?
3. For Figure 2, Can the authors provide experiment results with more iterations?
4. Are there any other variance reduction techniques and how does data balancing compare to those techniques?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: questions: 1. In your work, your target marginals were uniform; how would your method respond to non-uniform marginals?
2. 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)?
3. Are there existing methods for variance reduction and data balancing? If so, why did you not include an empirical comparison to existing methods?
4. 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: 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?

Comments: title: Re: Rebuttals by authors
Response to os3n
Rebuttal Read
Thank you for your feedback / Addressing additional concerns
comment: I thank the authors for their response. My concern about the clarity of the material presented remains and I would have liked the authors to discuss the practical implications of their results.

However, I agree that, overall, the theory produced is non-trivial and appreciated the additional experiments conducted, which I found informative. I have raised my score as a result.
We are happy to hear that the response to the review met your expectations. We are a bit puzzled as it appears on our end that the score is the same as before (5: Borderline accept). Does it appear to be the same for you? Thank you once again.
I thank the authors for detailed responses and additional experiments. I've raised the score of contribution.
Thank you for taking the time to provide a detailed rebuttal. You have clarified the points I inquired about.
Dear reviewers,

As the discussion period draws to a close on August 13 (2 business days from now) we kindly request that you take a moment to review and acknowledge our responses to your questions and comments.

In response to the comments raised by Reviewers os3n and vTkj, we have added additional experiments including other tasks and other data balancing baselines. We have also addressed questions from Reviewers G8ZX, os3n, and vTkj about theoretical assumptions and the link between the formal statements in Section 3 and the examples in Sections 2 and 4. Please see the responses for details.

If any further concerns or questions arise that we could address, please do not hesitate to reach out. We appreciate your time and feedback.

Sincerely, 

The authors