Title: The Benefits of Balance: From Information Projections to Variance Reduction
Venue: NeurIPS 2024
Accepted: yes
Keywords: regularized optimal transport

self-supervised learning

variance reduction

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

The Benefits of Balance: From Information Projections to Variance Reduction

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.

1 Introduction

Deep neural networks have shown remarkable success at learning task-specific representations of data when provided supervision from massive amounts of labeled training examples. Recent trends, however, have shifted toward task-agnostic, universal representations that may be easily fine-tuned or 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 of independent and identically distributed (i.i.d.) training examples ( X 1 , Y 1 ) , . . . , ( X n , Y n ) ∼ P , where P is a joint probability measure on sample spaces X × Y (e.g. feature-label or image-caption pairs) and let P n = 1 n ∑ n i =1 δ ( X i ,Y i ) be the empirical distribution. For a model parameterized by θ ∈ R d with loss function h θ , a stochastic learning algorithm involves computing the minibatch loss

and backpropagating through it to produce a minibatch stochastic gradient estimate. The algorithm then proceeds with the stochastic gradient descent (SGD) or a variant thereof (e.g., Adam, SGD with momentum, etc). 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 are generated by comparison to all other elements of the minibatch. While introducing such interventions into the procedure has clearly proven useful practically, it remains conceptually unclear what exactly is being optimized by the learning algorithm.

In this work, we aim to gain 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 precise benefits over traditional learning methods. We show that such recipes often enjoy an unsuspected benefit: reducing the variance of the empirical minibatch objective.

Concretely, we formalize the model updates described above as two phrases. Let Z 1 , . . . , Z n be a minibatch containing data points of arbitrary type (e.g. unlabeled images). In the first phase, this original data source is mapped (possibly using a model parameterized by θ ) 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, and we derive ( X i , Y i ) by setting X i = Z i and letting Y i be the pseudo-label based on clustering the vector representations of the images. In CLIP, each Z i is an image-caption pair, and we derive ( X i , Y i ) by simply letting X i be the image and Y i be the caption. Note that Y i is not a label in the traditional sense in neither of these examples. In the second phase, we use the model to compute a probability distribution P n,θ over X × Y , and perform a stochastic gradient update for the objective

This 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 fixed target marginal distributions P X on X and P Y on Y , balancing refers to repeatedly applying the operations

in 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 /R X or its columns by P Y /R 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.

Returning to P n,θ in (2), we show in Sec. 2 that both self-labeling and contrastive approaches in SSL implicitly define P n,θ by the following steps: 1) constructing a method-specific 'initial' measure P (0) n on X × Y , then 2) applying k iterations of the operations (3) to generate a sequence P (0) n , . . . , P ( k ) n , and finally, 3) setting P n,θ := P ( k ) n . In other words, these methods embed a learnable balancing operation in their objectives. A natural question to consider is: if the marginals one uses accurately represent the ones of the true probability measure P governing the data, are balanced quantities 'better behaved' than their unbalanced counterparts? If so, in what way?

Inspired by this question, we fix the model parameter θ (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 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 fixed function of interest, we define the population parameter φ and its k -step balanced estimator φ ( k ) n by

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 detail the mathematical connection between balancing and the modern representation learning techniques mentioned above. In Sec. 3, we prove a new upper bound on the MSEof the balancing estimator φ ( k ) n . The bound decomposes into an O ( n -1 ) first-order variance term and an O ( k 6 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. The second-order term can be used to compute the asymptotic variance reduction for statistical efficiency comparisons. Our proof technique relies on a recursion decomposition for balancing-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.

2 Data Balancing in Practice

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 θ . See Fig. 1 for examples of different choices of the sample spaces X and Y .

Example 1: Self-Supervised Clustering. Balancing is used in discriminative clustering and 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

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 expressions

Directly optimizing ∑ x,y P (0) n ( x, y ) h ( x, y ) without any constraints would lead to collapse, so it is balanced before optimization. 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 be allotted, 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

Balanced Clusters

BalancedModalities

BalancedMetadata

Clusters

Modality1

"The first

example

Modality 2

of .."

Minibatch

Training Set

Counts of

"the"

"Apicture

of...

Self-Supervised Clustering

Contrastive Learning

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. To our knowledge, our interpretation of the CLIP objective as an implicit data balancing procedure is novel. 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

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:

While 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) as

The measure P 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 balancing an entire training set, as opposed to a particular minibatch. 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 (possibly captions for images) is synthesized using a list of keywords { y 1 , . . . , y l } so that

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 have

where ℓ θ ( 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 ImageNetCaptions 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.

3 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 construct

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. 2

P(k)

Py

h

72

p(k-1)

71

P(0

p(k-1)

n,X

n,Y

L²(Px)

O

h-μxh

L²(Py)

Qx=Px

P*

Qy=Py

'a

n

S=coSa

L²(Py)

To provide background, the scheme of alternating the operators (8) is often seen as an iterative algorithm to solve the problem

where Π( 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) are based on the alternating minimization approach of solving

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. [1991]) to analyze the variance reduction afforded by the solution to (9) directly. However, quantifying the variance reduction (in terms of properties of P ) using this approach is challenging, as there is no closed-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 ≡ k ( n ) →∞ (scaled appropriately against n ), we can determine the reduction for the solution of (9) for large n . This is the subject of Thm. 1.

From Information Projections to Orthogonal Projections. First, we will show that the variance reduction resulting from each nonlinear iteration of (8) is associated with a linear operator applied to h . Thus, instead of analyzing the alternating information projections 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 ) ,

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 notation Q ( h ) := E Q [ h ( X,Y )] , we have by linearity of expectation that

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 (the X -marginal of P ( k ) n ) 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

and all other terms are O ( k 6 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

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 (2 n/p ⋆ ) + m log ( n +1)] /p 2 ⋆ and k ≥ 1 , we have

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

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 } j m =1 of L 2 ( P X ) , a basis { β j } j m =1 of L 2 ( P Y ) , and real values { s j } m j =1 , that satisfy

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

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 (see the right panel of Fig. 2).

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 as Y being a deterministic function of X . 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 write

Cor. 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

For simplicity, we only present the result for k even, i.e., σ 2 2 t .

Corollary 2. The variance reduction achieved by t + 1 iterations of the C Y C X operator can be quantified as

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.

4 Numerical Illustrations

We illustrate how data balancing manifests in the motivating examples mentioned in Sec. 2 with experiments with CLIP-type models. We focus here on zero-shot image classification tasks. Details on these experiments, and additional ones including linear probing and zero-shot retrieval, as well as an empirical investigation of the sensitivity to misspecified marginals, are all contained in Appx. E. Code to reproduce the data and experiments can be found at https://github.com/ronakdm/balancing.

0.2

0.4

0.6

0.8

CIFAR-10

m

=128

0.2

0.4

0.6

0.8

m

=512

0.2

0.4

0.6

0.8

m

=128

0.2

0.4

0.6

0.8

m

=512

0.10

0.15

0.20

0.25

0.30

m

=128

0.1

0.2

0.3

m

=512

0.0

0.2

0.4

CIFAR-100

0.2

0.4

0.1

0.2

0.3

0.0

0.1

0.2

0.3

0.01

0.02

0.03

0.01

0.02

0.03

0.04

0.05

Iterations

0.2

0.4

0.6

0.8

STL-10

Iterations

0.4

0.6

0.8

Iterations

0.2

0.4

0.6

0.8

Iterations

0.2

0.4

0.6

0.8

Iterations

0.1

0.2

0.3

Iterations

0.1

0.2

0.3

CLIP Text Embeddings

BERT Text Embeddings

GPT-2 Text Embeddings

No Balancing

CLIP Balancing (

k

=1

)

Multi-CLIP (

k

=2

)

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 ImageNetCaptions 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 )) , with

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.

Histogram of Metadata Categories in Pre-Training Dataset

Zero-Shot Accuracy of Models with Different Pre-Training Data

Evaluated on CIFAR-100

Evaluated on STL-10

Original Data

Rebalanced Data

Training Iterations

Training Iterations

0.9

10²

0.40

0.8

0.35

10

0.7

0.30

10

0.25

0.6

10

0.20

0.5

0.15

10

0.4

0.10

0.3

10

0.05

1000

2000

3000

4000

5000

1000

2000

3000

4000

5000

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 for solving entropy-regularized optimal transport (EOT). We refer to Peyr´ e 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.

5 Conclusion

We showed how several disparate techniques used towards the training of foundation models are instances of a data balancing algorithm, which has the unsuspected benefit of reducing the variance of learning objectives involving multiple sources of data. We proved 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.

Acknowledgements The authors are grateful to G. Ilharco, M. Wortsman, K. Pillutla, L. Schmidt, and J. Wellner for fruitful discussions related to this work. This work was supported by NSF DMS-2023166, CCF-2019844, DMS-2134012, PIMS 20240827-PRN01, NIH, and IARPA 202222072200003. Part of this work was done while L. Liu was with the University of Washington, and while R. Mehta and Z. Harchaoui were visiting the Simons Institute for the Theory of Computing.

Broader Impact While this paper is of a theoretical nature, the web-scale pre-training sets used to train foundation models can affect not only the biases of the models themselves but also the behavior of individuals who interact with them. In the case of representation learning, unrefined Internet data may lead to non-uniform performance among protected attributes such as gender, age, etc. For generative models, individuals of all ages may be influenced by harmful images or textual output. Studying the relationship between the balancing procedures considered in this paper and more holistic model evaluations presents a valuable direction for follow-up work.

References

Albertus and P. Berthet. Auxiliary information: The raking-ratio empirical process. Electronic Journal of Statistics , 13(1), 2019.

Asano, C. Rupprecht, and A. Vedaldi. Self-labelling via simultaneous clustering and representation learning. In ICLR , 2020.

Balestriero, M. Ibrahim, V. Sobal, A. Morcos, S. Shekhar, T. Goldstein, F. Bordes, A. Bardes, G. Mialon, Y. Tian, A. Schwarzschild, A. G. Wilson, J. Geiping, Q. Garrido, P. Fernandez, A. Bar, H. Pirsiavash, Y. LeCun, and M. Goldblum. A cookbook of self-supervised learning. arXiv preprint , 2023.

J. Bickel, Y . Ritov, and J. A. Wellner. Efficient estimation of linear functionals of a probability measure P with known marginal distributions. The Annals of Statistics , 1991.

Caron, I. Misra, J. Mairal, P. Goyal, P. Bojanowski, and A. Joulin. Unsupervised learning of visual features by contrasting cluster assignments. In NeurIPS , 2020.

Caron, H. Touvron, I. Misra, H. J´ egou, J. Mairal, P. Bojanowski, and A. Joulin. Emerging properties in self-supervised vision transformers. In ICCV , 2021.

Chen and K. He. Exploring simple Siamese representation learning. In CVPR , 2021.

Courty, R. Flamary, A. Habrard, and A. Rakotomamonjy. Joint distribution optimal transportation for domain adaptation. In NeurIPS , 2017.

M. Cover. Elements of Information Theory . John Wiley & Sons, 1999.

E. Deming and F. F. Stephan. On a least squares adjustment of a sampled frequency table when the expected marginal totals are known. Annals of Mathematical Statistics , 11, 1940.

Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. ImageNet: A large-scale hierarchical image database. In CVPR , 2009.

Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In ACL , 2019.

Everingham, L. Van Gool, C. K. I. Williams, J. Winn, and A. Zisserman. The PASCAL Visual Object Classes Challenge 2007 (VOC2007) Results, 2007.

Fang, G. Ilharco, M. Wortsman, Y. Wan, V. Shankar, A. Dave, and L. Schmidt. Data determines distributional robustness in contrastive language-image pre-training (CLIP). In ICML , 2013.

Y. Gadre, G. Ilharco, A. Fang, J. Hayase, G. Smyrnis, T. Nguyen, R. Marten, M. Wortsman, D. Ghosh, J. Zhang, E. Orgad, R. Entezari, G. Daras, S. M. Pratt, V. Ramanujan, Y. Bitton, K. Marathe, S. Mussmann, R. Vencu, M. Cherti, R. Krishna, P. W. Koh, O. Saukh, A. Ratner, S. Song, H. Hajishirzi, A. Farhadi, R. Beaumont, S. Oh, A. Dimakis, J. Jitsev, Y. Carmon, V. Shankar, and L. Schmidt. DataComp: In search of the next generation of multimodal datasets. In NeurIPS , 2023.

Genevay, L. Chizat, F. Bach, M. Cuturi, and G. Peyr´ e. Sample complexity of Sinkhorn divergences. In AISTATS , 2019.

Gohberg, S. Goldberg, and M. Kaashoek. Classes of Linear Operators Vol. 1 . Springer, 1990.

J.-B. Grill, F. Strub, F. Altch´ e, C. Tallec, P. Richemond, E. Buchatskaya, C. Doersch, B. Avila Pires, Z. Guo, M. Gheshlaghi Azar, B. Piot, k. kavukcuoglu, R. Munos, and M. Valko. Bootstrap your own latent: A new approach to self-supervised learning. In NeurIPS , 2020.

Hodosh, P. Young, and J. Hockenmaier. Framing image description as a ranking task: Data, models and evaluation metrics. Journal of Artificial Intelligence Research , 2013.

A. Horn and C. R. Johnson. Matrix Analysis . Cambridge University Press, 2013.

T. Ireland and S. Kullback. Contingency tables with given marginals. Biometrika , 1968.

J. Johnston and C. J. Pattie. Entropy-maximizing and the iterative proportional fitting procedure. The Professional Geographer , 45, 1993.

Jones, V. Roulet, and Z. Harchaoui. Discriminative clustering with representation learning with any ratio of labeled to unlabeled data. Statistics and Computing , 2022.

Kingma and J. Ba. Adam: A method for stochastic optimization. In ICLR , 2015.

Klatt, C. Tameling, and A. Munk. Empirical regularized optimal transport: Statistical theory and applications. SIAM Journal on Mathematics of Data Science , 2020.

T.-Y. Lin, M. Maire, S. Belongie, L. Bourdev, R. Girshick, J. Hays, P. Perona, D. Ramanan, C. L. Zitnick, and P. Doll´ ar. Microsoft COCO: Common Objects in Context, 2015.

Maji, J. Kannala, E. Rahtu, M. Blaschko, and A. Vedaldi. Fine-Grained Visual Classification of Aircraft. Technical report, University of Oxford, 2013.

Mena and J. Niles-Weed. Statistical bounds for entropic optimal transport: Sample complexity and the central limit theorem. In NeurIPS , 2019.

Nutz. Introduction to Entropic Optimal Transport. Lecture notes, Columbia University , 2021.

Oquab, T. Darcet, T. Moutakanni, H. V. Vo, M. Szafraniec, V. Khalidov, P. Fernandez, D. HAZIZA, F. Massa, A. El-Nouby, M. Assran, N. Ballas, W. Galuba, R. Howes, P.-Y. Huang, S.-W. Li, I. Misra, M. Rabbat, V. Sharma, G. Synnaeve, H. Xu, H. Jegou, J. Mairal, P. Labatut, A. Joulin, and P. Bojanowski. DINOv2: Learning robust visual features without supervision. Transactions on Machine Learning Research , 2024.

Peng, Q. Bai, X. Xia, Z. Huang, K. Saenko, and B. Wang. Moment matching for multi-source domain adaptation. In ICCV , 2019.

Peyr´ e and M. Cuturi. Computational Optimal Transport: With Applications to Data Science. Foundations and Trends in Machine Learning , 11, 2019.

Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners, 2019.

Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In ICML , 2021.

Schuhmann, R. Beaumont, R. Vencu, C. W. Gordon, R. Wightman, M. Cherti, T. Coombes, A. Katta, C. Mullis, M. Wortsman, P. Schramowski, S. R. Kundurthy, K. Crowson, L. Schmidt, R. Kaczmarczyk, and J. Jitsev. LAION-5B: An open large-scale dataset for training next generation image-text models. In NeurIPS , 2022.

Shen, Y. Qu, W. Zhang, and Y. Yu. Wasserstein distance guided representation learning for domain adaptation. In AAAI , 2018.

Sinkhorn. Diagonal equivalence to matrices with prescribed row and column sums. American Mathematical Monthly , 74(4), 1967.

M. E. Thompson. Theory of Sample Surveys . Chapman & Hall, 2000.

Xu, S. Xie, X. Tan, P.-Y. Huang, R. Howes, V. Sharma, S.-W. Li, G. Ghosh, L. Zettlemoyer, and C. Feichtenhofer. Demystifying CLIP data. In ICLR , 2024.

Appendix

Table of Contents

A Notation

B Linear Operators and Variance Reduction

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 .

B.1 Singular Value Decomposition

Recall the conditional mean operators µ X and µ Y from Sec. 3,

with the corresponding debiasing (a.k.a. centering) operators defined by C X = I -µ X and C Y = I -µ Y .

Proposition 3. There exists a basis { α j } j m =1 of L 2 ( P X ) , a basis { β j } j m =1 of L 2 ( P Y ) , and real values { s j } m j =1 , which satisfy:

α 1 = 1 X , β 1 = 1 Y , s 1 = 1 and 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

Since µ Y : L 2 ( P X ) → L 2 ( P Y ) is a compact linear operator, by Gohberg et al. [1990, Section IV.1 Theorem 1.1] and Gohberg et al. [1990, Section IV.1 Corollary 1.2], we have that µ Y admits a singular value decomposition satisfying (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

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 ) ,

This proves that s 1 ≤ 1 . For equality, notice that µ X µ Y 1 X = µ X 1 Y = 1 X , completing the proof.

B.2 Proof of Main Results

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 compare

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 that E P [ h ] = 0 , let

Then, we have that

Proof. By orthogonality, we have that

Because P ( h ) = 0 , it holds by the tower property of conditional expectation that P X ( µ X h ) = 0 , which implies that

For the second term, observe that P X ( C X h ) = 0 , so it holds by the tower property that P Y ( µ Y C X h ) = 0 , so

Next, we compute the term in the square by applying Prop. 3:

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

We have the following.

Lemma 5. For all t ≥ 0 , it holds that

Proof. Fix any j ∈ [ m ] , and use Prop. 3 to write

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 ) j m =1 and ( v j ) j m =1 by

Then, we may quantify the variance reduction achieved by t +1 iterations of the C Y C X operator as

as t →∞ .

Proof. Apply Lem. 4 ( t +1) -times so that

as by Lem. 5, we have that v τ,j = 0 for τ > 0 . Next, we unroll the definition of u τ,j so that

Substitute u 0 ,j = u j and v 0 ,j = v j to complete the proof.

Wealso 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 ) j m =1 by

Then, we may quantify the variance reduction achieved by t +1 iterations of the C X C Y operator as

as t →∞ .

In order to have full monotonicity, we also need that σ 2 0 ≥ σ 2 1 . This follows by orthogonality, as

Thus, we can combine Prop. 7 and (18) to fully quantify the relationship between σ 2 0 and σ 2 k for k odd.

C From Information Projections to Data Balancing

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 form

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 Balancing as Information Projections

C.1.1 Projection in KL-Divergence

Proposition 8. Assume that P X ≪ R X and P Y ≪ R Y , and define

Then, it holds that and

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 . Write

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 when Q 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 .

C.1.2 Projection in Reverse KL-Divergence

Proposition 9. Assume that P Y ≪ R X and P Y ≪ R Y , and define

Then, it holds that and

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 minimizing -∑ x,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

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 from R 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 -Divergence

Let 1 denote the function that is identically equal to 1 . Consider the following optimization problem, which is the subject of the subsequent lemmas:

where

Lemma 10. Assume that P X ≪ R X , and define The problem (25) is feasible, and its solution can be written as

for any f ∈ L 2 ( R ) , where the linear operator C R X is specified by

Proof. First, we establish feasibility by letting

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 that ∑ r j =1 α j = 1 . Then,

indicating that ∑ r j =1 α j f j ( x, y ) ∈ A X and A X is an affine subset of L 2 ( R ) . Define

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 of

we 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 ( I -µ R X )( 1 -f ) = C R X ( 1 -f ) , because C R X = I -µ R X . The claim is proved.

and let ξ ⋆ be the solution of problem (25) . Then,

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 solving

Notice 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.

for any y ′ ∈ Y . Thus, the likelihood ratio of Q ⋆ with respect to R is a marginal reweighting. Accordingly,

completing the proof.

Proposition 12. Assume that P X ≪ R X and P Y ≪ R Y . Define

Then, it holds that

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.

C.2 Proof of Main Results

We may now control the errors of the ratio of marginals using the projection interpretation established in the previous sections. Recall the event S as defined in Tab. 1. The following result, the monotonicity of the marginal violation terms in terms of KL, will be useful in the bound.

Proposition 13. [Nutz, 2021, Proposition 6.10] Under the event S , it holds that

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 that

In addition, we have that

Moreover, when KL( P n,X ∥ P X ) ≤ p 2 ⋆ / 2 , we have

Proof. We first show that P ( 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 ,

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 ,

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

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 that

where the last inequality follows by the monotonicity of Sinkhorn iterations given in Prop. 13. We apply the lower bounds to write

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 thus

Hence,

Now, for k even, set k = 2 t for t ≥ 0 . We have that

Invoke Prop. 13 once again to achieve

which completes the proof.

D Statistical Analysis of Balancing Estimators

This section contains the proof of the main result, namely Thm. 1. We first consolidate notation and then give 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 by

so that and

Recalling in addition that C k = C X for k odd and C k = C Y for k even. The event

is used for purely technical reasons in many results.

Proof Outline. We first establish that the recursion formula

holds in Prop. 15, where

The quantity V ( k -1) n ( C k h ) describes an error term that accumulates for each iteration of balancing, which explains why k must be scaled appropriately against n to ensure the error does not accumulate too fast. Applying the recursion repeatedly to the balanced sequence ( P ( k ) n ) k ≥ 1 and unrolling the recursion, we see that when k is odd,

Additionally, 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 as ∑ k ℓ =1 V ( ℓ -1) n ( h ℓ,k ) . Because our original goal is to upper bound the mean squared error, we use the expansion above to write

Regarding the first term, we have that E | 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.

D.1 Recursion of Estimation Error

We first recall that the sequence ( P ( k ) n ) k ≥ 1 can be computed with the following formula:

Prop. 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 ).

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 (34) and V ( k ) n defined in (36) , it holds that and

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

Arguing analogously for k even and subtracting P ( h ) on both sides, we have that

We 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.

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,

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.

D.2 Technical Tools & Intermediate Results

Having established the backbone of the argument, we collect in this subsection some useful tools that are used in the remainder of the proofs.

The following result follows from the method of types in information theory and will be helpful in deriving the dependence of the higher-order term on n .

Theorem 16. [Cover, 1999, Theorem 11.2.1] Let ν 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 that

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 pattern

and so on. This will be useful when bounding h ℓ,k uniformly.

Lemma 17. For any k ≥ 1 and ℓ ∈ { 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 ),

The red terms and blue terms cancel out to zero. This leaves

wherein we combine the red terms and re-index the blue terms to get

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 .

D.3 Analysis of Higher-Order Term

Returning to the outline at the start of this section, we may now bound the higher-order remainder term in (37), namely

depends on controlling the quantity V ( k -1) n in the summation, which we recall for convenience:

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.

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 by

for k ≥ 1 . We also enumerate the sample spaces as X = { x 1 , . . . , x m } and Y = { y 1 , . . . , y m } , and define the function

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 notation

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 that

whereas for k even, it holds that

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 h 1 jk ) = 0 . Indeed,

̸

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 h 1 jk ) = 0 , which implies the desired result. The proof for k even follows symmetrically.

The higher-order term in (37), can be bounded using Prop. 20.

Proposition 20. For any k ≥ 1 , the following holds under the event S :

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,

Using the statement above, we have that

The bound above holds for ℓ even as well. Then, using the (39) from Prop. 15 again, we have that for ℓ ≥ 2 ,

which implies that

by Lem. 18 and M 1 ≤ B 1 and M ℓ ≤ B 2 for ℓ ≥ 2 . Summing these bounds, we have that

because |X| = m . We sum up the last term:

completing the proof.

D.4 Proof of Main Results

We can now show the main result of this section: the bound on the mean squared error of the rebalanced estimator. Recall the event

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

so we decompose on the event S to write

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 (37) to write

for

By definition, we have that E 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 ) ≤ 2 m (1 -p ⋆ ) n . Moreover, for any δ ∈ (0 , 1) , we have

̸

Proof. Define F X := { Supp ( P n,X ) = Supp ( P X ) } and F 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 = ∪ j m =1 F j , which implies by the union bound that

̸

Similarly, we have that P ( F Y ) ≤ m (1 -p ⋆ ) n and thus P ( S c ) ≤ 2 m (1 -p ⋆ ) n , which gives the first claim.

To control the expectation, consider any δ > 0 , and define the event

By Hoeffding's inequality, it holds that P ( E δ ) ≥ 1 -δ . Furthermore, we get

In order to bound the terms appearing in (48), we introduce the events E δ 1 , E δ 2 , and E δ 3 , defined by

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 (49) . For any δ > 0 , assuming that n ≥ 2[log 2 (2 /δ ) + m log( n +1)] /p 2 ⋆ , we have that

Proof. The following computations are done under the event S . First, apply Prop. 20 to write

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 that P ( 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:

so that by plugging into (50),

and in turn,

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 get

and the bounds based on E δ 2 which give

By plugging into (50),

In turn,

Combining together both (56) and (51) and using that [log 2 (2 /δ ) + 2 m log( n +1)] ≥ 1 , we have that

the result as desired.

Lemma 23 (Cross term bound) . Let T 1 and T 2 be defined as in (49) . For any δ > 0 , assuming that n ≥ 2[log 2 (2 /δ ) + m log( n +1)] /p 2 ⋆ , we have that

Proof. The following computations are done under the event S . First, apply Prop. 20 to write

Wedecompose 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:

so that by plugging into (57),

and in turn,

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)] /p 2 ⋆ apply (33) from Prop. 14 to get

and the bounds based on E δ 2 ∩ E δ 2 ∩ E δ 3 which give

By plugging into (57),

In turn,

Combining together both (59) and (58) and using that [log 2 (2 /δ ) + 2 m log( n +1)] ≥ 1 , we have that

the result as desired.

We now combine the previous results to prove Thm. 24.

Theorem 24. For a sequence of rebalanced distributions ( ˜ P n ( k ) ) k ≥ 1 , there exists an absolute constant C > 0 such that when n ≥ C [log 2 (2 n/p ⋆ ) + m log ( n +1)] /p 2 ⋆ ,

where

Proof. We apply the decomposition (47), and subsequently handle the second term using bounds on the terms in (48). 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 that

when n ≥ ˜ C [log 2 (2 n/p ⋆ ) + m log ( n +1)] /p 2 ⋆ . This then implies that there is an absolute constant C 3 such that

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 that

which is the claimed result.

While not shown in the main text, similar techniques to those used above can also control the bias of ˜ P n ( k ) ( h ) as in Thm. 25. Interestingly, this bias is of order O ( n -2 ) which confirms the intuition that even thought ˜ P n ( 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 constant C > 0 such that when n ≥ C [log 2 (2 n/p ⋆ ) + m log ( n +1)] /p 2 ⋆ ,

where B is as defined in Thm. 24.

Proof. First, apply the decomposition (47) so that

By using the argument of Prop. 21, we have that

Then, by the recursion formula Equation (37), we have that

Because G (0) n ( C 1 . . . C k h ) has zero mean, it follows that

We have by Hoeffding's inequality that P ( E δ 3 ) ≥ 1 -δ , and that by Lem. 18 that G (0) n ( C 1 . . . C k h ) ≤ 4 k √ n ∥ h ∥ ∞ universally. As a result, applying Prop. 21 once again,

Using a similar argument to Lem. 22, we have that under S\ ( E δ 1 ∩E δ 2 ) (which occurs with probability no more than 2 δ ),

and that under S ∩ E δ 1 ∩ E δ 2 (which occurs with probability at least 1 -2 δ ),

Applying the decomposition | E 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.

D.5 Misspecified Marginal Distributions

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 ( ˆ P X,ϵ , ˆ P Y,ϵ ) which satisfy the following structure.

Assumption 1. There exist fixed probability mass functions ˆ P X and ˆ P Y for some ϵ ∈ [0 , 1) ,

We also have the existence of the positive quantity

Given the existence of ˆ p ⋆ > 0 , we may also define

To be precise, the iterations of balancing follow ˆ P (0) n = P n and

We start by deriving a result similar to Prop. 15. Since ϵ < 1 , the (possibly misspecified) target marginals ˆ P X,ϵ ( x ) > 0 and ˆ P Y,ϵ ( y ) > 0 for all x ∈ X and y ∈ Y . Define

and

The format of this section will be to derive results analogous 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 ⋆,ϵ ) .

D.5.1 Intermediate Results

Proposition 26. Let ( ˆ P ( k ) n ) k ≥ 1 be a sequence computed according to (62) . Define

These iterations are well-defined under the event S , and for G ( k ) n defined in (43) , it holds that

and

Furthermore,

Proof. 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 that

This proves the uncentered recursion formula given in (64). We then show the centered version.

Next, we bound the additional error term. By the Cauchy-Schwarz inequality,

as µ X is an orthogonal projection in L 2 ( P ) . Using convexity of f -divergences, we have that

This achieves the desired result.

Using similar ideas, we then prove an analog of Lem. 19.

Lemma 27. For k odd, it holds that

whereas for k even, it holds that

Proof. We give the proof for k odd. We claim that we need only show that P ( C X h 1 jk ) = 0 . This would show that

where P ( C X h 1 jk ) = 0 is employed in the last step. Now the result follows from (65) in Prop. 26 and the definition of ˆ G ( k ) n ( h ) . To prove the claim, as in Lem. 19, write

̸

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 h 1 jk ) = 0 , which implies the desired result. The proof for k even follows symmetrically.

For the remainder of the argument, we see that (66) can be unrolled so that

where we use that G (0) n = ˆ G (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

In addition, we have that

Moreover, when KL( P n,X ∥ P X ) ≤ ˆ p 2 ⋆,ϵ 8 and ϵ ≤ 1 -exp ( -ˆ p 2 ⋆,ϵ 8 ) , we have

Proof. 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 ,

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 ,

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

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 that

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 ∥ ˆ P X,ϵ ) = KL( P n,X ∥ ˆ P X,ϵ ) , which may not decay to zero as n →∞ . Because ϵ < 1 , write

We can then apply the lower bounds

Finally, combining the arguments above, we have that

where the last step invoked the assumption that

This means that

Hence,

Now, for k even, set k = 2 t for t ≥ 0 . We have that

Invoke Prop. 13 once again to achieve

Proceeding with similar steps, define the quantities

We must now establish an analog of Prop. 20.

Proposition 29. For any k ≥ 1 , the following holds under the event S :

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 ℓ ,

The bound above holds for ℓ even as well. Then, using (64) from Prop. 26 along with the triangle inequality, we have that for ℓ ≥ 2 ,

which implies that

by Lem. 18 and ˆ M 1 ≤ ˆ B 1 and ˆ M ℓ ≤ ˆ B 2 for ℓ ≥ 2 . The bound above holds trivially for ℓ = 1 . Summing these bounds over ℓ and j , we have that

because |X| = m . We sum up the last term:

which completes the proof.

D.5.2 Mean Squared Error Bound

Ultimately, we wish to construct an upper bound for

as the method returns P n ( h ) when S is not satisfied. The first term will be controlled by intermediate tools developed above. The second term that includes S c is no different from the one analyzed in Prop. 21. We handle the second term first. Recall from Prop. 21 that for any δ ∈ (0 , 1) ,

Repeat the argument from the proof of Thm. 24: because 2[log 2 (2 /δ ) + m log( n +1)] ≥ log( m/δ ) and -log(1 -p ⋆ ) ≥ p ⋆ ≥ p 2 ⋆ , we have that

This in turn implies that m (1 -p ⋆ ) n ≤ δ , and gives as a condition on the sample size n . Further in the analysis, we will set δ = (ˆ p ⋆,ϵ /n ) 4 , so right-hand side of (74) can then be upper bounded further, resulting in

a higher-order term compared to other components of the bound.

Next, we must control the left-hand side of (73). We perform the decomposition based on (68):

for

Recall the events E δ 1 and E δ 2 and E δ 3 from Appx. D.4. To perform this computation efficiently, we will split the bounds on each term into two components. In particular, we will show that

Under the event S ∩ E δ 1 ∩ E δ 2 : ∣ ∣ ∣ ˆ T 2 ∣ ∣ ∣ ≤ T 2 + E 2 ,

Under the event S\ ( E δ 1 ∩ E δ 2 ) : ∣ ∣ ∣ ˆ T 2 ∣ ∣ ∣ ≤ T c 2 + E c 2 ,

Under the event S ∩ E δ 3 : | T 1 | ≤ T 1 ,

Under the event S\E δ 3 : | T 1 | ≤ T c 1 ,

where any term denoted with ' E ' will represent all error terms that include ϵ and will be written in bigO notation. There are no errors for the bounds on T 1 , as this term does not depend on the misspecified marginals. The idea is that for the ' T 2 ' terms we may reuse the bounds derived in Appx. D.4 by simply replacing p ⋆ with ˆ p ⋆,ϵ . This is due to the fact that the dependence of the analogous terms from Appx. D.4 depend on p ⋆ only through Prop. 14; similarly, the corresponding terms in this section depend on ˆ p ⋆,ϵ through Prop. 28. We return to the terms in (76) and (77).

Decomposing on E δ 3 will result in a bound of the form

Decomposing on E δ 1 ∩ E δ 2 will result in a bounds of the form

Finally, decomposing on E δ 1 ∩ E δ 2 ∩ E δ 3 will result in a bound of the form

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 δ = ˆ p 4 ⋆,ϵ /n 4 , we will achieve a similar result to Thm. 24, i.e., that

It remains to quantify the ˜ O 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 ( -ˆ p 2 ⋆,ϵ 8 ) . Then, it holds that

Proof. The following computations are done under the event S . First, apply Prop. 29 to write

We decompose on the event E 1 ∩ E 2 .

Bound | T 2 | under the event S\ ( E δ 1 ∩E δ 2 ) . In this case, we apply (69) from Prop. 28 to get ˆ B 1 ≤ n and ˆ B 2 ≤ 1 / ˆ p 2 ⋆,ϵ , along with the universal bounds from Lem. 18:

so that by plugging into (81),

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 (70) from Prop. 28 to get

The bounds based on E δ 2 give

By plugging into (81), we can reuse the steps in the bound from (55) (for all terms without ϵ ) to write so that

We compute E 2 by using that

which gives

We now make the corresponding argument for the term T 1 .

Lemma 31. For δ = (ˆ p ⋆,ϵ /n ) 4 , it holds that

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:

so that T c 1 = 4 k ∥ h ∥ ∞

Bound | T 1 | under the event S ∩ E δ 3 . Now, we may use the definition of the event E δ 3 to achieve

Knowing that E c 2 = 0 , we simplify (79) and (79) to read

We now combine the bounds from the previous two lemmas to compute (79) and (80) to state the main result.

Theorem 32. Let Asm. 1 be true with error ϵ ∈ [0 , 1) . For a sequence of rebalanced distributions ( ˆ P ( k ) n ) k ≥ 1 , there exists an absolute constant C > 0 such that when n ≥ C [log 2 (2 n/ ˆ p ⋆,ϵ ) + m log ( n +1)] / min { p ⋆ , ˆ p ⋆,ϵ } 2 , we have that

E Experimental Details

We provide the full experimental details of the experimental results from Sec. 4. We report additional evaluations on downstream tasks with linear probing and zero-shot retrieval. Finally, we give illustrations of the sensitivity to misspecified marginals, and of the convergence to the given marginals.

E.1 Datasets

Pre-Training Data. The pre-training data was taken from the public ImageNet-Captions dataset [Fang et al., 2013]. We subset the dataset by selecting the 250 classes that were most frequent in the dataset, resulting in 174,594 images and associated Flickr captions. The exact images used and their associated captions are given in the code supplement.

Evaluation Data. We perform zero-shot classification (as described in Sec. 4), zero-shot retrieval, and linear probing with various image classification and image-caption datasets. We used the default class captions (for classification) and default linear probing parameters from the CLIP Benchmark repo. The datasets (test splits) used were:

CIFAR-10: 10,000 colored natural images labeled with one of 10 classes.

CIFAR-100: 10,000 colored natural images labeled with one of 100 classes.

STL-10: 80,000 colored natural images labeled with one of 10 classes.

MS-COCO: 41,000 colored natural images with associated captions.

Flickr8k: 8,000 colored natural images with associated captions.

Rendered SST2: 1,821 images of typed natural language with sentiment label (2 classes).

VOC2007: 4,952 colored natural images labeled with one of 20 classes.

FGVC Aircraft: 34,000 colored natural images labeled with one of 102 classes.

Evaluation scripts using the various embedding models (described below) are provided.

E.2 Model Specification and Hyperparameters

Architecture and Implementation. The models considered CLIP models [Radford et al., 2021], and are specified by pairs of encoders ( f θ , g θ ) , representing images and text, respectively. The encoders decompose into f θ = f head θ ◦ f base θ (similarly for g θ ) where f base θ denotes a base image encoder and f head θ denotes a trainable head model. The head models are feed-forward networks with two hidden layers, 256 hidden units, and 128-dimensional output representations. Their input dimensions may be 512 or 768, depending on whether a CLIP model or BERT/GPT-2 model is used as the base. For the image base/foundation models, we use the open-source OpenCLIP implementation of the ViT-B/32 model with the laion2b s34b b79k model tag. For the text encoder, we use the encoder of the variant of the ViT-B/32 with tag datacomp xl s13b b90k . For the other text encoders the Huggingface implementations of GPT-2 and BERT were used.

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 of varying losses as described in Sec. 4.

E.3 Compute Environment

Experiments were run on a CPU/GPU workstation with 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 repositories were used for zero-shot evaluation.

E.4 CLIP and Multi-CLIP

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 optimizing

This 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.

E.5 Metadata Curation

We considered in the metadata curation example from Sec. 2 how to use balancing to adjust the entire pre-training set, in the spirit of Xu 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.

E.6 Additional Experiments

In this section, we provide 1) a synthetic data example that helps elucidate the role of the spectral decomposition introduced in Sec. 3, and 2) additional evaluations on downstream tasks such as zero-shot retrieval and linear probing. For the latter, we maintain the experimental settings as used in the zero-shot classification example from Sec. 4 (Fig. 3). That is, we train variants of CLIP models (see Sec. 2) on the ImageNet-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 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 be

where ˆ P X and ˆ P Y 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 as

0.0

0.2

0.4

0.6

0.8

Leading Singular Value

s

10

4

10

3

10

2

Mean Squared Error

Marginal Fitting Techniques under Misspecifiation

=

0.25 (IPWI)

=

0.125 (IPWI)

=

0.0 (IPWI)

=

0.25 (Bal.)

=

0.125 (Bal.)

=

0.0 (Bal.)

Emp. Measure

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 by

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 .

Zero-Shot Retrieval. In this evaluation, we assess the ability of the learned representations to match queries from one modality to their counterparts in another modality. We are given a test sets X test = { x 1 , . . . , x M } of images and Y test = { y 1 , . . . , y N } of texts in natural language. We are also given a matrix of annotations A ∈ { 0 , 1 } M × N where A ij = 1 if and only if y j is a 'relevant' caption

0.0

0.1

0.2

0.3

MSCOCO (Image)

CLIP Text Embeddings

0.000

0.025

0.050

0.075

0.100

BERT Text Embeddings

0.00100

0.00125

0.00150

0.00175

0.00200

GPT-2 Text Embeddings

0.0

0.1

0.2

0.3

MSCOCO (Text)

0.00

0.05

0.10

0.002

0.004

0.006

0.008

0.0

0.2

0.4

Flickr-8k (Image)

0.00

0.05

0.10

0.15

0.0045

0.0050

0.0055

0.0060

0.0065

Training Steps

0.0

0.2

0.4

Flickr-8k (Text)

Training Steps

0.00

0.05

0.10

0.15

0.20

Training Steps

0.005

0.010

0.015

0.020

Zeroshot Retrieval Performance (Recall @ 5)

No Balancing

CLIP Balancing (

k

=1

)

Multi-CLIP (

k

=2

)

0.60

0.62

0.64

Rendered SST-2

CLIP Text Embeddings

BERT Text Embeddings

GPT-2 Text Embeddings

0.2

0.4

0.6

0.8

VOC2007

Training Steps

0.2

0.3

0.4

FGVC Aircraft

Training Steps

Training Steps

No Balancing

CLIP Balancing (

k

=1

)

Multi-CLIP (

k

=2

)

for image x i (and vice versa). Given a particular query y ∈ Y test, we define the topK neighborhood of y as

i.e. the images in the test set that have the closest embeddings under the given model. Then, we may define the average recall at K for image retrieval metric as

In words, the metric evaluates the retrieval system's ability to detect relevant items in the dataset, in this case by comparing the closeness of the image-text representations. We can analogously define the average recall at K for text retrieval metric by swapping the role of x and y above. The results for both retrieval metrics on the MS-COCO [Lin et al., 2015] and Flickr8k [Hodosh et al., 2013] benchmarks are given in Fig. 6.

Linear Probe. Here, we evaluate the quality of the model's encoders by fine-tuning a single linear layer on top of the learned representations for a classification task. In the case of linear probing via image classification, we use only the image encoder f θ I . We are given a training set { ( x 1 , c 1 ) , . . . , ( x N , c N ) } of image-label pairs, where each c i ∈ { 1 , . . . , C } . We fix the model

ROXx

ROXx

RrO

n,X

RCo)

n,Y

C

y

C

y

C

y

parameter θ I and solve the regularized multinomial cross entropy (MCE) objective

where λ > 0 is a regularization parameter, ∥·∥ F denotes the Frobenius norm on R C × r and LogSoftmax : R C → R C is given by LogSoftmax( z ) = z -log ∑ C j =1 exp( z j ) . This results in a classifier

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 (82) 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.

F NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Theoretical claims, the focus of this paper, are supported with proofs.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We clarify that the setting studied has some dissimilarities with practice in Sec. 2.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate 'Limitations' section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: This is done for all theoretical statements.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Code is provided to reproduce the main results and an extensive experimental details section is written.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: This is written in the public repo provided in Sec. 4.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Yes, we even give a list of the specific images of ImageNet used to train the multimodal models.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Our evaluation metrics are shown with all seeds and their mean plotted in the corresponding figures.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer 'Yes' if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Please see Appx. E.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: NA

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: The work is primarily theoretical and retrospective.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

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).

11. 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.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide a list of notebooks and scripts to illustrate our method and connect it to existing software repositories such as OpenCLIP.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: NA

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: NA

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe paper has theoretical contributions that include the derivation of non-asymptotic bounds that quantify the variance reduction achieved through their data balancing technique. The authors also present empirical studies that demonstrate the effectiveness of their proposed balancing technique. The authors discuss the utility of data balancing across different tasks, such as image-caption pair matching and self-supervised clustering, identifying the utility of their approach. Their approach has the potential for adoption in various domains, including in the training of foundation models.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n1. In your work, your target marginals were uniform; how would your method respond to non-uniform marginals?\n2. Your target marginals were accurately specified (as uniform). What if the target marginals of the two distributions were less accurately specified (i.e., the true underlying distributions are not well-known)? How do you think that this would influence the empirical results of your technique (e.g., zero-shot average per-class recall)?\n3. Are there existing methods for variance reduction and data balancing? If so, why did you not include an empirical comparison to existing methods?\n4. You mention that the zero-shot evaluation metrics are difficult to produce intervals for (i.e., you are missing error bars). Why is this the case?\n\nLIMITATIONS:\nI 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?", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. the paper provides a new perspective on the benefits of data balancing\n2. provide different examples of data balancing in practice and prove a non-asymptotic bound on the MSE of balanced estimators.\n3. The findings may have implications for improving SSL models\n\nWEAKNESSES:\n1.the experiments are somewhat limited in scope\n2.adding more visualizations or intuitive explanations may be better for understanding the key finding of the paper.\n3.will the assumptions limit the applicability of the findings?\n\nQUESTIONS:\n1. Can the findings be extended to other areas?\n2. Can the authors provide more details on how the assumptions, such as the spectral gap condition, hold or when will they not hold in practical scenarios? Are there specific types of data or models where these assumptions are more likely to be satisfied?\n3. For Figure 2, Can the authors provide experiment results with more iterations?\n4. Are there any other variance reduction techniques and how does data balancing compare to those techniques?\n\nLIMITATIONS:\nIn 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis paper attempts to shed light on SSL training and the role of data balancing. The paper formalises the problem and develops extensive theory. The main results is pretty cool and insightful in the sense that the upper bound on the MSE shows that data balancing has a variance reduction effect. The topic is of interest to the community and the work is focussing on a poorly understood paradigm that is becoming dominant.\n\nWEAKNESSES:\nI have three main concerns with this work:\n\n1/ The theory is *very* extensive. The Appendix contains several pages of proofs that are difficult to parse and come on top of the formalism presented in the main paper. It seems like the main body could be simplified and made more to the point to convey the main gist of the contribution and make it more accessible.\n\n2/ It is unclear how the data balancing examples in Section 2 map to the formalism introduced in Section 3. For example, what would (4) look like for example 1 and example 2?\n\n3/ It is unclear what the experiments bring to the table and how they provide evidence to the main result. Making the link more explicit and explaining what are the key take aways from these results would help the reader.\n\nQUESTIONS:\nI have the following questions for the authors:\n\n- Line 54: What do you mean by \"X and Y are forms of the data that are related to, but distinct from, the form of Z\"\" given that Z is equal to (X,Y)?\n- p4, example 1: What would the target marginals correspond to? What is \\psi_n^(k) and P_n^(k) here? \n- p4, example 2: What would the target marginals correspond to? What is \\psi_n^(k) and P_n^(k) here?\n- Why do we need \\tilde{\\psi}_n^{(k)} in (12) and how does it relate to \\psi_n^{(k)}?\n- How does (15) relate to the clip example introduced earlier in the paper and why is this a valid and sensible simplification to study?\n- Does the main result have implications in practice in terms of design of algorithm?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-08-05"}