Title: Optimal Algorithms for Augmented Testing of Discrete Distributions
Venue: NeurIPS 2024
Accepted: yes
Keywords: distribution testing

learning-augmented algorithms

data driven algorithm

hypothesis testing

hypothesis selection

distribution learning
Abstract: We consider the problem of hypothesis testing for discrete distributions. In the standard model, where we have sample access to an underlying distribution $p$, extensive research has established optimal bounds for uniformity testing,  identity testing (goodness of fit), and closeness testing (equivalence or two-sample testing). We explore these problems in a setting where a predicted data distribution, possibly derived from historical data or predictive machine learning models, is available. We demonstrate that such a predictor can indeed reduce the number of samples required for all three property testing tasks. The reduction in sample complexity depends directly on the predictor’s quality, measured by its total variation distance from $p$. A key advantage of our algorithms is their adaptability to the precision of the prediction. Specifically, our algorithms can self-adjust their sample complexity based on the accuracy of the available prediction, operating without any prior knowledge of the estimation’s accuracy (i.e. they are consistent). Additionally, we never use more samples than the standard approaches require, even if the predictions provide no meaningful information (i.e. they are also robust). We provide lower bounds to indicate that the improvements in sample complexity achieved by our algorithms are information-theoretically optimal. Furthermore, experimental results show that the performance of our algorithms on real data significantly exceeds our worst-case guarantees for sample complexity, demonstrating the practicality of our approach.

Optimal Algorithms for Augmented Testing of Discrete Distributions

Abstract

We consider the problem of hypothesis testing for discrete distributions. In the standard model, where we have sample access to an underlying distribution p , extensive research has established optimal bounds for uniformity testing, identity testing (goodness of fit), and closeness testing (equivalence or two-sample testing). We explore these problems in a setting where a predicted data distribution, possibly derived from historical data or predictive machine learning models, is available. We demonstrate that such a predictor can indeed reduce the number of samples required for all three property testing tasks. The reduction in sample complexity depends directly on the predictor's quality, measured by its total variation distance from p . A key advantage of our algorithms is their adaptability to the precision of the prediction. Specifically, our algorithms can self-adjust their sample complexity based on the accuracy of the available prediction, operating without any prior knowledge of the estimation's accuracy (i.e. they are consistent). Additionally, we never use more samples than the standard approaches require, even if the predictions provide no meaningful information (i.e. they are also robust). We provide lower bounds to indicate that the improvements in sample complexity achieved by our algorithms are information-theoretically optimal. Furthermore, experimental results show that the performance of our algorithms on real data significantly exceeds our worst-case guarantees for sample complexity, demonstrating the practicality of our approach.

1 Introduction

Property testing of distributions is a fundamental task that lies at the heart of many scientific endeavors: Given sample access to an underlying unknown distribution p , the goal is to infer whether p has a certain property or it is ϵ -far from any distribution that has the property (in some reasonable notion of distance, such as total variation distance) with as few samples as possible. Over the past century [66], this problem has been extensively explored in statistics, machine learning, and theoretical computer science. Indeed, distribution testing (also called hypothesis testing) is now a major pillar of modern learning theory and algorithmic statistics, with applications in learning mixtures of distributions such as Gaussians, Poisson Binomial Distributions and robust learning [40, 45, 36, 75, 34, 45]. The framework has also been extensively studied under privacy [56, 14, 27, 18, 26, 3, 72, 50] and low-memory constraints [8].

Ronitt Rubinfeld MIT

One of the most natural and well-studied questions in this framework is: given sample access to an unknown distribution p , can we determine whether p is equal to another distribution q , or ϵ -far from it (e.g. in total variation distance)? This problem has been studied under various assumptions about how we access q : It is called uniformity testing when q is a uniform distribution, identity testing (or 'goodness-of-fit') when a description of q is known, and closeness testing ('two-sample testing' or 'equivalence testing') when we only have sample access to q . The primary goal in solving these tasks is to design algorithms that use as few samples as possible. Optimal sample complexity bounds have been established for discrete distributions p and q over a domain of size n : Θ( √ n/ϵ 2 ) samples for uniformity testing [52, 15, 69, 38] and identity testing [76, 5, 39], and Θ ( n 2 / 3 /ϵ 4 / 3 + √ n/ϵ 2 ) samples for closeness testing [16, 30, 38, 39]. Other related versions of uniformity, identity, and closeness testing are presented in [63, 6, 19, 51, 46, 74, 12, 13, 68].

Given that the aforementioned results are tight and cannot be improved, any further progress requires equipping the algorithm with additional functionality. A natural approach is to leverage the fact that in numerous applications, the underlying distribution is not completely unknown; some prediction of the underlying distribution may be available or can be learned via a predictive machine learning model. For example, if distributions evolve slowly over time, earlier iterations can serve as approximations for later ones, e.g., network traffic data and search engine queries. Such estimations can often be learned from 'older' data by using it to train a predictor or regressor. In linguistics and text processing tasks that involve distributions over words, the length of a word can approximate its frequency, since longer words are known to be less frequent. Another example is when data is available at different 'scales.' For instance,demographic data on loan defaults at the national level could be informative for data from specific areas.

One challenge to using such information is that it rarely comes with a guarantee of precision. This fact leads to the information being deemed unreliable, as it may poorly predict the underlying distribution. For example, while the national loan default rate might be close to that of a typical area code, it could differ significantly in affluent areas. Thus a natural algorithmic question that arises is how to design algorithms that can exploit predictive information as much as possible without any prior assumptions about its accuracy. The goal is then to design an algorithm that solves the problem with as few samples as possible, given the quality of the prediction.

In this work, we study the fundamental problems of uniformity, identity, and closeness testing in the setting where a prediction of the underlying distribution is provided. This prediction can be formalized by assuming that the distribution testing algorithm has access to a predicted distribution ˆ p of p 1 . This is in addition to having sample access to p as in the standard model (without access to prediction). Our algorithms achieve the optimal reduction in the number of samples used compared to the standard case, where the improvement depends on the quality of the predictor ˆ p in terms of its total variation distance from p . Our algorithms can also self-adjust their sample complexity to the accuracy of ˆ p , minimizing sample complexity wherever feasible, without prior knowledge of ˆ p 's accuracy. Our approach ensures that the algorithm is resilient to inaccuracies in predictions and does not exceed the optimal sampling bound in the standard model, even when ˆ p significantly deviates from the actual p . Furthermore, our matching lower bounds demonstrate the optimality of our algorithms. Experimental results additionally confirm the practicality of our algorithms.

Measuring accuracy of predictor. We use the total variation distance between the prediction ˆ p and the unknown distribution p as our measure of predictor accuracy. Previous work often assumed a strong element-wise guarantees, where ˆ p i is within a constant multiplicative factor of p i for all domain elements i , a constraint that becomes limiting especially for small p i (see Section 1.3). This paper is the first to study a notion of average error between p and ˆ p , measuring via the TV distance. This metric was chosen for its prevalent use in statistical inference and its intuitive interpretation.

1.1 Our approach and problem formulation

Our approach to solving a distribution testing problem (uniformity, identity, or closeness testing) consists of two components: search and test . At a high level, search aims to guess ∥ p -ˆ p ∥ TV, and test performs the actual distribution testing using the guess of the accuracy provided by search as a suggested accuracy level.

More precisely, our augmented test component aims to evaluate whether p = q , while receiving ˆ p and a suggested accuracy level α (which may or may not reflect the true distance between p and ˆ p ). In addition to two possible outcomes of accept and reject in the standard setting 2 , our augmented test component may output inaccurate information when it determines that the ˆ p is not α -close to p . Our requirements for the augmented test component are the following: i ) If the test is conclusive, i.e., it chooses to output accept or reject , the answer should be correct regardless of ˆ p 's accuracy. ii ) If ˆ p is indeed α -close to p , the test component should not output inaccurate information .

We emphasize that the guess α is not guaranteed to be correct or may not even be a valid upper bound on the true TV distance. Thus, the algorithm is afforded a degree of flexibility to forego solving the problem when the distributions are not within α proximity (and can try again with a new guess).

Our search component aims to identify an appropriate accuracy level α such that the test component can test in a conclusive manner by returning accept or reject . Since the true value distance ∥ p -ˆ p ∥ TV is not known, we start by guessing a small α , corresponding to the most accurate ˆ p and the fewest samples needed, run the augmented test component with this α and ˆ p , and verify the conclusiveness of the testing. If inconclusive, our guess α is increased to a level that we can afford testing by doubling the sample size, and the search component proceeds with the next α . It continues until the desired accuracy is reached, and accept or reject is returned. Then, search halts with that result.

Clearly, if the accuracy level guess α is at least ∥ p -ˆ p ∥ TV, the test component has to output accept or reject with high probability. Thus, we show that it is unlikely that search proceeds when α ≫ ∥ p -ˆ p ∥ TV. Therefore, this method does not significantly increase the sample complexity, potentially increasing it by at most an O (log log( n/ϵ )) factor in expectation. The search component is applicable to all of the problems we study and we defer all details of the search component to Section B. The remainder of the main body focuses on designing the augmented testers, i.e. the test component.

Definition 1.1 (Augmented tester) . Suppose we are given four parameters, α ∈ [0 , 1] , ϵ ∈ (0 , 1) , δ ∈ (0 , 1) , n ∈ N , and two underlying distributions p and q , along with a prediction distribution ˆ p over [ n ] . Suppose A is an algorithm that receives all these parameters, and the description of ˆ p as input, and it has sample access to p and q . We say algorithm A is an ( α, ϵ, δ ) -augmented tester for closeness testing if the following holds for every p , q , and ˆ p over [ n ] :

If ˆ p and p are α -close in total variation distance, the algorithm outputs inaccurate information with a probability at most δ/ 2 .

If p = q , then the algorithm outputs reject with a probability at most δ/ 2 .

If p is ϵ -far from q , then the algorithm outputs accept with a probability at most δ/ 2 .

In this definition, if the description of q is known to the algorithm instead of having sample access, we say A is an ( α, ϵ, δ ) -augmented tester for identity testing. If q is a uniform distribution over [ n ] , then we say A is an ( α, ϵ, δ ) -augmented tester for uniformity testing.

To highlight the distinction between this definition and the standard definition, note that in the standard regime, no prediction ˆ p is available to the algorithm, and the algorithm lacks the option of outputting inaccurate information .

Remark 1. We assume that δ is a small constant, but, a standard amplification technique via Chernoff bounds can achieve an arbitrarily small confidence parameter δ with a O (log(1 /δ )) overhead.

1.2 Our results

Our theoretical results: In this paper we demonstrate that predictions can indeed reduce the number of samples needed to solve the three aforementioned testing problems. Our algorithms are parameterized by both ϵ and α . We provide tight sample complexity results (matching upper and lower bounds) for these problems. The sample complexity drops drastically compared to the standard case, depending on the total variation distance between q and ˆ p . Our algorithms are also robust: if the prediction error is high, our algorithm succeeds by using (asymptotically) the same number of samples as the standard setting without predictions.

̸

where d = ∥ q -ˆ p ∥ TV ( q is the known distribution for identity testing, or the uniform distribution for uniformity).

Remark 3. Note that d is not an input parameter to the algorithm. However, d is determined by ˆ p and q , which are known to the algorithm, allowing us to compute d .

For closeness testing, we prove the following.

Theorem 4 (Informal version of Theorem 9) . Augmented closeness testing for distributions over [ n ] , with parameters α , ϵ , and δ = 2 / 3 , requires Θ( n 2 / 3 α 1 / 3 /ϵ 4 / 3 + √ n/ϵ 2 ) samples.

It can be seen that e.g., for closeness testing, our non-trivial predictor improves over the best possible sampling bound in the standard model. Specifically, as long as α = o (1) , our bound in the augmented model improves over the prior work. At the same time, our algorithms are resilient: even if α = Θ(1) , our sampling bounds do not exceed those in the standard model. Note that all theorems are complemented by tight lower bounds. We highlight that all of our algorithms are also computationally efficient, running in polynomial time with respect to n and 1 /ϵ . The results are summarized in Table 1.

Our empirical Results: We empirically evaluate our augmented closeness testing algorithm on synthetic and real distributions and refer all details to Section E.

As a summary, our algorithm can indeed leverage predictions to obtain significantly improved sample complexity over the SOTA approach without predictions, as well as SOTA algorithms needing very accurate predictions [28].

For distributions similar to our lower bound instances, our augmented algorithm achieves > 20x reduction in sample complexity to obtain comparable accuracy as the standard un-augmented algorithm, as shown in Figure 1. On real distributions curated from network traffic data, we see sample complexity reductions of up to 40% . Furthermore, our algorithm is empirically robust to noisy predictions, in contrast to prior state of the art approaches which assume very accurate, point-wise predictions (CRS'15 in Figure 1).

It is worth noting that our experiments on network traffic data reveal that the actual sample complexity is much lower than the anticipated worst-case sample complexity of our algorithm. In particular, this holds

1000

5750

10500

15250

20000

0

10

20

30

40

% Error

Closeness Tester Hard Instance

Closeness Tester

Augmented Tester

CRS '15

even when ˆ p is far from p in terms of total variation distance. Empirically, if ˆ p accurately reflects the high-probability elements in p , our algorithm can significantly reduce the sample complexity needed for testing by utilizing these heavy hitter 'hints' from ˆ p . This is validated by our results and depicted in Figures 5 and 7(b).

1.3 Related works

To the best of our knowledge, there have been only three prior works that studied any distribution property testing algorithms with predictions, each assuming a much stronger prediction model:

Distribution testing with perfect predictors [29]: this work studied distribution testing problems, including closeness, identity and uniformity, assuming query access to a perfect predictor, i.e., ˆ p = p . They show that, given a perfect predictor, it is possible to obtain highly efficient testers for a wide variety of problems with a small number of queries to the prediction. Unfortunately, the perfect-predictor assumption is often too strong to hold in practice, as demonstrated e.g., in [48] and in our experiments.

Distribution testing with (1 + ϵ ) -approximate predictors [28, 67]: these works relax the assumption used in the above paper, requiring only that ˆ p i = (1 ± ϵ/ 2) p i for all i and sufficiently small ϵ > 0 . However, this assumption is still quite restrictive, especially for low values of ϵ . Indeed in our experimental setting, such algorithms are not robust to prediction errors (see Section E).

Support estimation with c -approximate predictors [48]: this work focused on the single problem of support estimation, i.e., estimating the fraction of coordinates i such that p i > 0 . It further relaxes the assumption in [67] by allowing the predicted probabilities ˆ p i to be within a factor of c of the true probabilities p i , for any constant approximation factor c > 1 . This algorithm has been shown to work well in practice. However, the techniques presented in that paper seem to be applicable exclusively to support estimation. Furthermore, their result provably does not hold under the assumption that p and ˆ p are close in TV distance, as in this paper.

In summary, prior results required either highly restrictive assumptions, or were applicable to only a single problem. None of the previous algorithms worked under the TV distance assumption used in this paper which is arguably the most natural. Further exploration of measures such as L p -distances, KL-divergence, and Hellinger distance are interesting open questions.

We remark that we can mathematically show that these alternative oracles yield provably more power compared to ours (i.e., we make weaker assumptions about the predictor). We provide an alternative to our upper bound techniques for these alternate prediction models in Section G. We demonstrate how a variant of our algorithm, in conjunction with these stronger predictions, implies that uniformity testing, identity testing, and closeness testing in these models can be conducted using only O ( √ n/ϵ 2 ) samples. This low sample complexity effectively circumvents our lower bound for closeness testing, suggesting that these models provide stronger, and arguably less realistic, predictions.

Other related works: Other related works discussing algorithms with predictions and general distribution testing are discussed in Appendix A.

1.4 Notation and organization

Notation: We use [ n ] to denote the set { 1 , 2 , . . . , n } . All of our distributions will be over the domain [ n ] , and we assume n is always known. For a distribution p , we denote the probability of the i -th element by p i . For any subset of the domain S ⊆ [ n ] we denote the probability mass of S according to p by p ( S ) . We use p ⊗ s to denote the probability distribution of s i.i.d. samples drawn from p . We say p is a known distribution, if we have access to every p i . We say p is an unknown distribution if we have only sample access to p . We use Poi ( λ ) to denote a random variable from a Poisson distribution with mean λ . Similarly, Ber ( λ ) indicates a random variable from the Bernoulli distribution that is one with probability λ . Given two distributions p and q over a domain X , we use ∥ p -q ∥ TV := sup S ⊆X | p ( S ) -q ( S ) | to denote the total variation distance between p and q . We say p and q are ϵ -close ( ϵ -far), if the total variation distance between p and q is at most ϵ (larger than ϵ ).

Organization: We provide an overview of our theoretical results in Section 2. The upper bound on augmented closeness testing is presented in Section 3. The search algorithm is detailed in Section B.

Augmented uniformity and identity testing are discussed in Section C, where the upper bounds are presented in Section C.1, and the lower bounds are provided in Sections C.2 and C.3.

Augmented closeness testing is discussed in Section D, with the lower bounds provided in Section D.1. Our empirical evaluations are presented in Section E.

2 Overview of our proofs

Search algorithm: The search algorithm seeks to find the smallest value of α for which the problem is solvable via the augmented tester. It starts with the lowest value of α (most accurate prediction). Then, it iteratively increases the sample budget across rounds. In each round i , it selects an α i for the augmented tester, A , ensuring operation within the current sample budget. If A outputs accept or reject , the algorithm echoes this outcome. If inaccurate information is returned, the sample budget doubles for the next round. It is worth noting that this search scheme is applicable to a general distribution testing algorithm with polynomial sample complexity. The algorithm's pseudocode and its performance proof are in Section B.

Upper bound for augmented identity testing Let d represent the distance between ˆ p and q (the prediction and the known distribution). For establishing the upper bound, it is essential to assume d > α . If not, the prediction proves unhelpful, and we might as well resort to the standard tester.

̸

Our upper bound relies on a simple but fundamental observation regarding the total variation distance: this distance is the maximum discrepancy between the probability masses that two distributions assign to any subset of their domain. To prove that total variation between two distributions is small entails proving that the discrepancy across every domain subset is small. In contrast, to prove a large total variation distance, one only needs to identify a single subset with a large discrepancy as evidence of large total variation distance. We use the Scheffé set of ˆ p and q -characterized as the collection of elements x ∈ [ n ] where ˆ p ( x ) < q ( x ) , symbolized by S -as evidence of p 's divergence from either q or ˆ p . More precisely, it is known that S maximizing the discrepancy between the probability masses of ˆ p and q , implying: d := ∥ q -ˆ p ∥ TV = | q ( S ) -ˆ p ( S ) | . Next, we estimate the probability of set S according to p with reasonably high accuracy. Given q ( S ) and ˆ p ( S ) , then p ( S ) is either significantly different from q ( S ) , or it deviates from ˆ p ( S ) . In the first scenario, this discrepancy serves as evidence that p = q , allowing us to output reject . In the second scenario, the deviation confirms that p is α -distant from ˆ p , leading us to output inaccurate information . Further details can be found in Section C.1.

Lower bound for augmented uniformity testing We provide two lower bounds for augmented uniformity testing. One is purely based on a reduction to standard uniformity testing for the case where α ≥ d . (Recall that d was the total variation distance between q and ˆ p ). See Section C.2.

The other lower bound applies to the setting where α < d . One challenge of this problem was that it is hard to find two difficult distributions that the tester has to distinguish between; usually, for a pair of distributions, we could come up with one valid output that could serve them both. For example, for both uniform distribution, and the famous ϵ -far uniform distribution that assigns probabilities (1 ± ϵ ) /n to the elements, the algorithm may be able to output inaccurate information . Hence, we cannot draw lower bounds just by asking the algorithm to distinguish between two distributions.

For this reason, we provide three distributions that look similar when we draw too few samples. We formalize the similarity of these three distributions using a multivariate coupling argument. We show that these distributions are such that there is no possible answer that is valid for all three of them. Now, (similar to Le Cam's method), suppose we feed the algorithm with samples from one of these three distributions (each with probability 1/3). For any sample set, the algorithm outputs an answer (which may be randomized); however, this answer is considered wrong for at least one of the underlying distributions. This is due to the fact that there is no universally valid answer that is simultaneously correct for all three distributions. Hence, if the algorithm outputs a valid answer with high probability, it must be able to distinguish the underlying distributions to some degree. On the other hand, the indistinguishability result says it is impossible to tell these distributions apart. Thus, we reach a contradiction. And, the lower bound is concluded. See Section C.3 for further details.

Lower bound for augmented closeness testing: We provide two separate lower bounds for closeness testing based on the relationship between α and ϵ . Further details are available in Section D.1.

Our first lower bound, as detailed in Theorem 11, employs a reduction strategy from standard closeness testing to augmented closeness testing when α ≥ ϵ . This is achieved by taking instances used in standard closeness testing for distributions over [ n -1] and embedding them into the first n -1 domain elements of a new distribution p defined over [ n ] . The key to this strategy is to put the majority of the distribution's mass ( (1 -α ) mass) on its final element, and we set the prediction ˆ p to be a singleton distribution over the last element, which is α -close to p . Clearly, the prediction does not reveal any information about the first n -1 elements of p , implying that testing the closeness of p in the augmented setting is as challenging as in the standard setting, once the parameters are appropriately scaled.

Our second lower bound, outlined in Theorem 12, is more involved. In the standard setting, the lower bound for closeness testing is derived from the hard instances for uniformity testing with one crucial adjustment: adding new elements with large probability (approximately n -2 / 3 ) in the distributions. These large elements have identical probability masses in both p and q , indicating they do not contribute to the distance between the two distributions. However, their presence in the sample set confuses the algorithm: due to their high probabilities, their behavior in the sample set may misleadingly suggest non-uniformity, complicating the algorithm's task of determining the uniformity of the rest of the distribution. Therefore, the algorithm requires s ≈ n 2 / 3 samples to first identify these large elements before it can test the uniformity of the remaining distribution. Surprisingly, this requirement of s ≈ n 2 / 3 samples is significantly higher than the O ( √ n ) samples typically sufficient for testing uniformity.

The challenge in our case arises because ˆ p may disclose the large elements to the algorithm. To establish the lower bound, we set ˆ p to be the uniform distribution, we generate hard instances of p by adding as many large elements as possible, without altering ˆ p , by keeping the overall probability mass of the large elements limited to α . More precisely, p assigns approximately (1 -α ) /n probability mass to O ( n ) elements chosen at random, and assigns approximately n -2 / 3 probability mass to α · n 2 / 3 elements in the domain. Now, q has two scenarios. Half the time, q is identical to p . The other half, q retains the same large elements but deviates slightly from uniformity for the rest of the distribution. Specifically, q assigns probabilities (1 ± Θ( ϵ ))(1 -α ) /n to the randomly chosen O ( n ) elements, making it ϵ -far from p . It is not hard to show that testing closeness of p and q is a symmetric property (since permuting the domain elements does not affect our construction). By leveraging the wishful thinking theorem from [78], we demonstrate that these two scenarios are indistinguishable unless Ω( n 2 / 3 α 1 / 3 ) samples are drawn.

3 Upper bound for closeness testing

Our upper bound is based on a technique called flattening , which has been previously proposed by Diakonikolas and Kane [42]. This technique is instrumental in reducing the variance of the statistic used for closeness testing by reducing the ℓ 2 2 -norm of the input distributions. We adapt this technique for use in the augmented setting, aiming not only to flatten the distribution based on the samples received from it but also exploiting the prediction distribution ˆ p . We demonstrate that augmented flattening can significantly reduce the ℓ 2 2 -norm of p . In our algorithm, we initially check if the ℓ 2 2 -norm of p is reduced after flattening to the desired bound. If not, this indicates that the prediction was not sufficiently accurate, leading us to output inaccurate . Conversely, if the ℓ 2 2 -norm of p is small, we proceed with an efficient testing algorithm that requires fewer samples. We describe the standard flattening technique in Section 3.1. Our flattening technique presented in Section 3.2. Finally, we provide our algorithm in Section 3.2.1.

3.1 Background on flattening

Suppose we are given n parameters m 1 , m 2 , . . . , m n . One can create a randomized mapping F that assigns each i ∈ [ n ] to a pair ( i, j ) , where j is drawn uniformly at random from [ m i ] . Now, consider a given distribution p over [ n ] and a sample i drawn from p . This mapping induces a distribution over pairs ( i, j ) 's in D := { ( i, j ) | i ∈ [ n ] and j ∈ [ m i ] } . We denote this new distribution by p ( F ) satisfying the relation: p ( F ) ( i,j ) = p i /m i . The core idea of the above mapping is to divide

the probability of the i -th element into m i buckets . If the values of m i 's are large for elements in p with higher probabilities, then the resulting distribution p ( F ) will avoid having any elements with disproportionately large probabilities, thereby naturally lowering its ℓ 2 -norm. Diakonikolas and Kane [42] showed that if we draw Poi ( t ) samples from p , and set each m i to be the frequency of element i among these samples plus one, then we have:

where the expectation is taken over the randomness of the samples.

Connection to distribution testing: Chan et al. in [30] showed that one can test closeness of two distributions over [ n ] using O ( n max( ∥ p ∥ 2 , ∥ q ∥ 2 ) /ϵ 2 ) samples. Diakonikolas et al. [42] used this tester in combination of the flattening technique to map distributions to distributions over slightly larger domains with the goal of reducing their ℓ 2 -norms. They have shown that if we use the same mapping to reduce the ℓ 2 -norm of both p and q , the ℓ 2 -distance between the two distribution does not change the ℓ 1 distance between distributions.

Fact 3.1 ([42]) . For any flattening scheme F , the ℓ 1 -distance is preserved under F . That is, for every pair of distributions p and q , we have:

Therefore, to test p and q , we can draw samples to come up with the flattening F to reduce the ℓ 2 -norm of one of the underlying distributions. Then, we can test the closeness of p ( F ) and q ( F ) with O ( n max( ∥ ∥ p ( F ) ∥ ∥ 2 , ∥ ∥ q ( F ) ∥ ∥ 2 ) /ϵ 2 ) new samples. Later, Aliakbarpour et al. in [9] showed that O ( n min( ∥ p ∥ 2 , ∥ q ∥ 2 ) /ϵ 2 ) samples is sufficient.

The exact sample complexity is determined by balancing the samples needed to determine the flattening and the samples needed to test closeness of p ( F ) and q ( F ) . Note that flattened distributions have a larger domain size. Thus, one must also balance the gains obtained from the reduction in ℓ 2 norm with the increase in the domain.

3.2 Augmented flattening

We adapt the flattening argument in [42] to the augmented setting.

Lemma 3.2. Suppose we have an unknown distribution p = ( p 1 , p 2 , . . . , p n ) over [ n ] and Poi ( s f ) samples from p . Assume we are given a known distribution ˆ p that is α -close to p in TV distance. Then for every v ≤ 1 , there exists a flattening F which increases the domain size by 1 /v + s f in expectation and the expected ℓ 2 2 -norm of the p ( F ) is bounded by:

The expectation above is taken over the randomness of the samples from p .

Proof. Let ˆ s f := Poi ( s f ) be a Poisson random variable with mean s f . Let S f denote the multiset of ˆ s f samples from p . Let f i denotes the frequency of element i in S f . For every i , we define the number of buckets for element i as follows:

Let p ( F ) denote the flattened version of p where we split every element i into m i buckets. We show that the expected ℓ 2 -norm of p ( F ) is low. Define ∆ i to be p i -ˆ p i . Suppose A is a set of indices i ∈ [ n ] for which ∆ i ≥ ˆ p i . For every i ∈ A , p i is bounded from above by 2 ∆ i . On the other hand, for every i ∈ [ n ] \ A , p i is bounded by 2 ˆ p i .

As shown in [42], the expected value of p i / ( f i +1) taken over the randomness in the S f is bounded by 1 /s f :

The last inequality above is due to the fact that:

Therefore, we obtain:

3.2.1 The algorithm

In this section, we provide an algorithm for testing closeness of two distribution in the augmented setting. Our algorithm receives a suggested accuracy level α . Based on α , we draw s f (which depends on α ) samples from p and use them to flatten p . If the resulting distribution has sufficiently small ℓ 2 2 -norm we proceed to the testing phase. Otherwise, we declare that the accuracy level provided is not correct. The pseudocode of our algorithm is provided in Algorithm 1, and we prove its performance in Theorem 5.

Proof. We begin the proof by setting our parameters. Set s f := min ( n 2 / 3 α 1 / 3 /ϵ 4 / 3 , n ) . Suppose we draw ˆ s f = Poi ( s f ) samples from p . We flatten the distribution p according to the augmented flattening we described in Section 3.2 replacing ν with 1 /n . We denote the flattening by F = F ( m 1 , . . . , m n ) where m i 's are defined in Equation (1):

In the above definition, f i is the frequency of element i among the ˆ s f drawn samples.

Proof of correctness: Suppose we flatten both p and q according to F . The algorithm estimates the ℓ 2 2 of p ( F ) and q ( F ) . According to Fact F.1, these estimates are within a constant factor of their true values each with probability at least 1 -0 . 05 :

To prove the correctness of the algorithm, we show the three desired property of an augmented tester which are defined in Definition 1.1. First, assume ˆ p is α -close to p . We show the probability of outputting inaccurate information is small. Using Lemma 3.2, after applying F , the expected ℓ 2 2 -norm of p ( F ) is bounded by:

Using Markov's inequality, with probability at least 0.95, ℓ 2 2 -norm of p ( F ) is at most 20 times the above bound. Combined by the upper bound for L p in Equation (3) and the union bound, in this case L p is bounded by 30 · (2 α/s f +4 /n ) with probability at least 0.9. Hence, we the algorithm does not output inaccurate information in Line 10 with probability more than 0.1.

Second, we show that if p = q , the algorithm outputs reject with small probability. We may output reject in three cases: 1) in Line 5 2) in Line 13 and 3) in Line 15. Using Markov's inequality, a Poisson random variable is more than 10 times its expectation with probability at most 0.1. We show the other two cases are unlikely as long as L p and L q are accurate estimates of the ℓ 2 2 s of p ( F ) and q ( F ) (which happens with probability of 0.9). If p = q , p ( F ) and q ( F ) are identical. Therefore, in the algorithm L p and L q are the estimations of the ℓ 2 2 -norm of the same distribution. Note that we output reject in Line 13, only when L p < 3 L q . Using Equation (3), this event does not happen unless at least of the estimates are inaccurate (which has a probability at most 0.1). Furthermore if the estimated ℓ 2 2 -norm are accurate, the minimum ℓ 2 2 -norm of p ( F ) and q ( F ) are bounded by 90 · (2 α/s f +4 /n ) implying the tester would work correctly with probability at least 0.9. Using the union bound, the probability of outputting reject is bounded by 0.3.

Third and last, we show that if p is ϵ -far from q ( F ) , the probability of outputting accept is small. We only output accept in Line 15. Similar to the second case we have discussed, as long as the ℓ 2 2 -norm are accurate, the tester does not make a mistake with probability more than 0.1. Therefore, the overall probability of making a mistake is bounded by 0.2.

Analysis of sample complexity: For the flattening step we have used ˆ s f ≤ 10 s f . After flattening, the new domain size is bounded from above by:

Using Fact F.1, estimation of ℓ 2 2 -norm requires O ( √ n ) samples. Using Fact F.2, the number of samples we use for the test in Line 15 is O ( n · b/ϵ 2 ) samples. Hence, the total sample complexity is:

In the last line we use that s f := min ( n 2 / 3 α 1 / 3 /ϵ 4 / 3 , n ) . Thus, the proof is complete.

Acknowledgments

Part of this work was conducted while the authors were visiting the Simons Institute for the Theory of Computing. M.A. was supported by NSF awards CNS-2120667, CNS-2120603, CCF-1934846, and BU's Hariri Institute for Computing. This work was initiated while M.A. was affiliated with Boston University and Northeastern University and was done in part while M.A. was a research fellow at the Simons Institute for the Theory of Computing. P.I. was supported by the NSF TRIPODS program (award DMS-2022448) and Simons Investigator Award. R.R. was supported by the NSF TRIPODS program (award DMS-2022448) and CCF-2310818.

References

Anders Aamand, Alexandr Andoni, Justin Y. Chen, Piotr Indyk, Shyam Narayanan, and Sandeep Silwal. Data structures for density estimation. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 1-18. PMLR, 23-29 Jul 2023.

Anders Aamand, Justin Y. Chen, Huy Lê Nguyen, Sandeep Silwal, and Ali Vakilian. Improved frequency estimation algorithms with and without predictions. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023 , 2023.

J. Acharya, C. Canonne, C. Freitag, and H. Tyagi. Test without trust: Optimal locally private distribution testing. In Proceedings of Machine Learning Research , volume 89 of Proceedings of Machine Learning Research , pages 2067-2076. PMLR, 2019.

J. Acharya, Z. Sun, and H. Zhang. Differentially private testing of identity and closeness of discrete distributions. In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018 , pages 6879-6891, 2018.

Jayadev Acharya, Constantinos Daskalakis, and Gautam Kamath. Optimal testing for properties of distributions. In Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada , pages 3591-3599, 2015.

Jayadev Acharya, Ashkan Jafarpour, Alon Orlitsky, and Ananda T. Suresh. Sublinear algorithms for outlier detection and generalized closeness testing. In IEEE ISIT , pages 3200-3204, 2014.

Maryam Aliakbarpour, Eric Blais, and Ronitt Rubinfeld. Learning and testing junta distributions. In Proceedings of the 29th Conference on Learning Theory, COLT 2016, New York, USA, June 23-26, 2016 , pages 19-46, 2016.

Maryam Aliakbarpour, Mark Bun, and Adam Smith. Hypothesis selection with memory constraints. To appear in NeurIPS 2023 , 2023.

Maryam Aliakbarpour, Ilias Diakonikolas, Daniel Kane, and Ronitt Rubinfeld. Private testing of distributions via sample permutations. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, 8-14 December 2019, Vancouver, BC, Canada , pages 10877-10888, 2019.

Maryam Aliakbarpour, Ilias Diakonikolas, and Ronitt Rubinfeld. Differentially private identity and equivalence testing of discrete distributions. In Proceedings of the 35th International Conference on Machine Learning (ICML) , volume 80, pages 169-178, 2018.

Maryam Aliakbarpour, Themis Gouleakis, John Peebles, Ronitt Rubinfeld, and Anak Yodpinyanee. Towards testing monotonicity of distributions over general posets. In Proceedings of the Thirty-Second Conference on Learning Theory, COLT , pages 34-82, 2019.

Maryam Aliakbarpour, Ravi Kumar, and Ronitt Rubinfeld. Testing mixtures of discrete distributions. In Alina Beygelzimer and Daniel Hsu, editors, Proceedings of the Thirty-Second Conference on Learning Theory , volume 99 of Proceedings of Machine Learning Research , pages 83-114, Phoenix, USA, 25-28 Jun 2019. PMLR.

Maryam Aliakbarpour and Sandeep Silwal. Testing Properties of Multiple Distributions with Few Samples. In Thomas Vidick, editor, 11th Innovations in Theoretical Computer Science Conference (ITCS 2020) , volume 151 of Leibniz International Proceedings in Informatics (LIPIcs) , pages 69:1-69:41, Dagstuhl, Germany, 2020. Schloss Dagstuhl-Leibniz-Zentrum fuer Informatik.

Kareem Amin, Matthew Joseph, and Jieming Mao. Pan-private uniformity testing. In Jacob D. Abernethy and Shivani Agarwal, editors, Conference on Learning Theory, COLT 2020, 9-12 July 2020, Virtual Event [Graz, Austria] , volume 125 of Proceedings of Machine Learning Research , pages 183-218. PMLR, 2020.

Tugkan Batu, Lance Fortnow, Ronitt Rubinfeld, Warren D. Smith, and Patrick White. Testing that distributions are close. In 41st Annual Symposium on Foundations of Computer Science, FOCS 2000, 12-14 November 2000, Redondo Beach, California, USA , pages 259-269, 2000.

Tugkan Batu, Lance Fortnow, Ronitt Rubinfeld, Warren D. Smith, and Patrick White. Testing closeness of discrete distributions. JACM , 60(1):4:1-4:25, 2013.

Tugkan Batu, Ravi Kumar, and Ronitt Rubinfeld. Sublinear algorithms for testing monotone and unimodal distributions. In Proceedings of the 36th Annual ACM Symposium on Theory of Computing, Chicago, IL, USA, June 13-16, 2004 , pages 381-390, 2004.

Thomas Berrett and Cristina Butucea. Locally private non-asymptotic testing of discrete distributions is faster using interactive mechanisms. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

Bhaswar B. Bhattacharya and Gregory Valiant. Testing closeness with unequal sized samples. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett, editors, Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada , pages 2611-2619, 2015.

Arnab Bhattacharyya, Eldar Fischer, Ronitt Rubinfeld, and Paul Valiant. Testing monotonicity of distributions over general partial orders. Electronic Colloquium on Computational Complexity (ECCC) , 17:27, 2010.

B. Cai, C. Daskalakis, and G. Kamath. Priv'it: Private and sample efficient identity testing. In International Conference on Machine Learning, ICML , pages 635-644, 2017.

Clément L. Canonne. Big data on the rise: Testing monotonicity of distributions. In ICALP , 2015.

Clément L. Canonne. Corrigendum: Are few bins enough: Testing histogram distributions. In Proceedings of the 42nd ACM SIGMOD-SIGACT-SIGAI Symposium on Principles of Database Systems, PODS 2023, Seattle, WA, USA, June 18-23, 2023 , page 359. ACM, 2023.

Clément L. Canonne, Ilias Diakonikolas, Daniel Kane, and Sihan Liu. Nearly-tight bounds for testing histogram distributions. In Advances in Neural Information Processing Systems 35 (NeurIPS) , 2022. To appear.

Clément L. Canonne, Ayush Jain, Gautam Kamath, and Jerry Li. The price of tolerance in distribution testing. In Po-Ling Loh and Maxim Raginsky, editors, Conference on Learning Theory, 2-5 July 2022, London, UK , volume 178 of Proceedings of Machine Learning Research , pages 573-624. PMLR, 2022.

Clément L. Canonne, Gautam Kamath, Audra McMillan, Adam D. Smith, and Jonathan R. Ullman. The structure of optimal private tests for simple hypotheses. In Moses Charikar and Edith Cohen, editors, Proceedings of the 51st Annual ACM SIGACT Symposium on Theory of Computing, STOC 2019, Phoenix, AZ, USA, June 23-26, 2019 , pages 310-321. ACM, 2019.

Clément L Canonne, Gautam Kamath, Audra McMillan, Jonathan Ullman, and Lydia Zakynthinou. Private identity testing for high-dimensional distributions. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 10099-10111. Curran Associates, Inc., 2020.

Clément L Canonne, Dana Ron, and Rocco A Servedio. Testing probability distributions using conditional samples. SIAM Journal on Computing , 44(3):540-616, 2015.

Clément L. Canonne and Ronitt Rubinfeld. Testing probability distributions underlying aggregated data. In Automata, Languages, and Programming - 41st International Colloquium, ICALP , pages 283-295, 2014.

Siu-on Chan, Ilias Diakonikolas, Paul Valiant, and Gregory Valiant. Optimal algorithms for testing closeness of discrete distributions. In SODA , pages 1193-1203, 2014.

Justin Y. Chen, Talya Eden, Piotr Indyk, Honghao Lin, Shyam Narayanan, Ronitt Rubinfeld, Sandeep Silwal, Tal Wagner, David P. Woodruff, and Michael Zhang. Triangle and four cycle counting with predictions in graph streams. In 10th International Conference on Learning Representations, ICLR , 2022.

Xi Chen, Rajesh Jayaram, Amit Levi, and Erik Waingarten. Learning and testing junta distributions with sub cube conditioning. In Mikhail Belkin and Samory Kpotufe, editors, Proceedings of Thirty Fourth Conference on Learning Theory , volume 134 of Proceedings of Machine Learning Research , pages 1060-1113. PMLR, 15-19 Aug 2021.

Edith Cohen, Ofir Geri, and Rasmus Pagh. Composable sketches for functions of frequencies: Beyond the worst case. In Proceedings of the 37th International Conference on Machine Learning , 2020.

Constantinos Daskalakis, Ilias Diakonikolas, and Rocco A. Servedio. Learning poisson binomial distributions. In Howard J. Karloff and Toniann Pitassi, editors, Proceedings of the 44th Symposium on Theory of Computing Conference, STOC 2012, New York, NY, USA, May 19 - 22, 2012 , pages 709-728. ACM, 2012.

Constantinos Daskalakis, Ilias Diakonikolas, Rocco A. Servedio, Gregory Valiant, and Paul Valiant. Testing k-modal distributions: Optimal algorithms via reductions. In Proceedings of the Twenty-fourth Annual ACM-SIAM Symposium on Discrete Algorithms , SODA '13, pages 1833-1852, Philadelphia, PA, USA, 2013. Society for Industrial and Applied Mathematics.

Constantinos Daskalakis and Gautam Kamath. Faster and sample near-optimal algorithms for proper learning mixtures of gaussians. In Conference on Learning Theory , pages 1183-1213. PMLR, 2014.

Ilias Diakonikolas, Themis Gouleakis, Daniel M. Kane, and Sankeerth Rao. Communication and memory efficient testing of discrete distributions. In Alina Beygelzimer and Daniel Hsu, editors, Proceedings of the Thirty-Second Conference on Learning Theory , volume 99 of Proceedings of Machine Learning Research , pages 1070-1106. PMLR, 25-28 Jun 2019.

Ilias Diakonikolas, Themis Gouleakis, J. Peebles, and Eric Price. Collision-based testers are optimal for uniformity and closeness. Chicago Journal of Theoretical Computer Science , 2019(1), MAY 2019.

Ilias Diakonikolas, Themis Gouleakis, John Peebles, and Eric Price. Sample-optimal identity testing with high probability. In 45th International Colloquium on Automata, Languages, and Programming, ICALP 2018, July 9-13, 2018, Prague, Czech Republic , pages 41:1-41:14, 2018.

Ilias Diakonikolas, Gautam Kamath, Daniel Kane, Jerry Li, Ankur Moitra, and Alistair Stewart. Robust estimators in high-dimensions without the computational intractability. SIAM J. Comput. , 48(2):742-864, 2019.

Ilias Diakonikolas, Daniel Kane, Vasilis Kontonis, Sihan Liu, and Nikos Zarifis. Efficient testable learning of halfspaces with adversarial label noise. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023 , 2023.

Ilias Diakonikolas and Daniel M. Kane. A new approach for testing properties of discrete distributions. In IEEE 57th Annual Symposium on Foundations of Computer Science, FOCS 2016, 9-11 October 2016, Hyatt Regency, New Brunswick, New Jersey, USA , pages 685-694, 2016.

Ilias Diakonikolas, Daniel M. Kane, and Vladimir Nikishkin. Optimal algorithms and lower bounds for testing closeness of structured distributions. In FOCS , pages 1183-1202, 2015.

Ilias Diakonikolas, Daniel M. Kane, and Vladimir Nikishkin. Testing identity of structured distributions. In SODA , pages 1841-1854, 2015.

Ilias Diakonikolas, Daniel M. Kane, and Alistair Stewart. Properly learning poisson binomial distributions in almost polynomial time. In Vitaly Feldman, Alexander Rakhlin, and Ohad Shamir, editors, Proceedings of the 29th Conference on Learning Theory, COLT 2016, New York, USA, June 23-26, 2016 , volume 49 of JMLR Workshop and Conference Proceedings , pages 850-878. JMLR.org, 2016.

Ilias Diakonikolas, Daniel M. Kane, and Alistair Stewart. Sharp bounds for generalized uniformity testing. In NeurIPS , pages 6204-6213, 2018.

Elbert Du, Franklyn Wang, and Michael Mitzenmacher. Putting the 'learning" into learningaugmented algorithms for frequency estimation. In Proceedings of the 38th International Conference on Machine Learning , pages 2860-2869, 2021.

Talya Eden, Piotr Indyk, Shyam Narayanan, Ronitt Rubinfeld, Sandeep Silwal, and Tal Wagner. Learning-based support estimation in sublinear time. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 . OpenReview.net, 2021.

M. Gaboardi, H. W. Lim, R. M. Rogers, and S. P. Vadhan. Differentially private chi-squared hypothesis testing: Goodness of fit and independence testing. In International Conference on Machine Learning, ICML , pages 2111-2120, 2016.

M. Gaboardi and R. Rogers. Local private hypothesis testing: Chi-square tests. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018 , pages 1612-1621, 2018.

Oded Goldreich. The uniform distribution is complete with respect to testing identity to a fixed distribution. In Computational Complexity and Property Testing - On the Interplay Between Randomness and Computation , pages 152-172. 2020.

Oded Goldreich and Dana Ron. On Testing Expansion in Bounded-Degree Graphs , pages 68-75. Springer Berlin Heidelberg, Berlin, Heidelberg, 2011.

Aravind Gollakota, Adam R. Klivans, and Pravesh K. Kothari. A moment-matching approach to testable learning and a new characterization of rademacher complexity. In Proceedings of the 55th Annual ACM Symposium on Theory of Computing , STOC 2023, pages 1657-1670, New York, NY, USA, 2023. Association for Computing Machinery.

Aravind Gollakota, Adam R. Klivans, Konstantinos Stavropoulos, and Arsen Vasilyan. Testerlearners for halfspaces: Universal algorithms. In Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023 , 2023.

Aravind Gollakota, Adam R. Klivans, Konstantinos Stavropoulos, and Arsen Vasilyan. An efficient tester-learner for halfspaces. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024 , 2024.

Sivakanth Gopi, Gautam Kamath, Janardhan Kulkarni, Aleksandar Nikolov, Zhiwei Steven Wu, and Huanyu Zhang. Locally private hypothesis selection. In Jacob D. Abernethy and Shivani Agarwal, editors, Conference on Learning Theory, COLT 2020, 9-12 July 2020, Virtual Event [Graz, Austria] , volume 125 of Proceedings of Machine Learning Research , pages 1785-1816. PMLR, 2020.

Chen-Yu Hsu, Piotr Indyk, Dina Katabi, and Ali Vakilian. Learning-based frequency estimation algorithms. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019 . OpenReview.net, 2019.

Piotr Indyk, Reut Levi, and Ronitt Rubinfeld. Approximating and testing k -histogram distributions in sub-linear time. In PODS , pages 15-22, 2012.

Piotr Indyk, Ali Vakilian, and Yang Yuan. Learning-based low-rank approximations. In Advances in Neural Information Processing Systems , pages 7400-7410, 2019.

Tanqiu Jiang, Yi Li, Honghao Lin, Yisong Ruan, and David P. Woodruff. Learning-augmented data stream algorithms. In International Conference on Learning Representations , 2020.

Adam R. Klivans, Konstantinos Stavropoulos, and Arsen Vasilyan. Learning intersections of halfspaces with distribution shift: Improved algorithms and SQ lower bounds. In The Thirty Seventh Annual Conference on Learning Theory, June 30 - July 3, 2023, Edmonton, Canada , Proceedings of Machine Learning Research. PMLR, 2024.

Adam R. Klivans, Konstantinos Stavropoulos, and Arsen Vasilyan. Testable learning with distribution shift. In The Thirty Seventh Annual Conference on Learning Theory, June 30 - July 3, 2023, Edmonton, Canada , Proceedings of Machine Learning Research. PMLR, 2024.

Reut Levi, Dana Ron, and Ronitt Rubinfeld. Testing properties of collections of distributions. Theory of Computing , 9(8):295-347, 2013.

Yi Li, Honghao Lin, Simin Liu, Ali Vakilian, and David Woodruff. Learning the positions in countsketch. In 11th International Conference on Learning Representations, ICLR , 2023.

Shyam Narayanan. Private high-dimensional hypothesis testing. In Conference on Learning Theory, 2-5 July 2022, London, UK , volume 178 of Proceedings of Machine Learning Research , pages 3979-4027. PMLR, 2022.

Jerzy Neyman and Egon S. Pearson. On the problem of the most efficient tests of statistical hypotheses. Philosophical Transactions of the Royal Society of London. Series A, Containing Papers of a Mathematical or Physical Character , 231(694-706):289-337, 1933.

Krzysztof Onak and Xiaorui Sun. Probability-revealing samples. In International Conference on Artificial Intelligence and Statistics , volume 84 of Proceedings of Machine Learning Research , pages 2018-2026. PMLR, 2018.

Aadil Oufkir, Omar Fawzi, Nicolas Flammarion, and Aurélien Garivier. Sequential algorithms for testing closeness of distributions. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 11655-11664, 2021.

Liam Paninski. A coincidence-based test for uniformity given very sparsely sampled discrete data. IEEE Trans. Inf. Theory , 54(10):4750-4755, 2008.

Sampriti Roy and Yadu Vasudev. Testing properties of distributions in the streaming model. In 34th International Symposium on Algorithms and Computation, ISAAC 2023, December 3-6, 2023, Kyoto, Japan , volume 283 of LIPIcs , pages 56:1-56:17, 2023.

Ronitt Rubinfeld and Arsen Vasilyan. Testing distributional assumptions of learning algorithms. In Proceedings of the 55th Annual ACM Symposium on Theory of Computing , STOC 2023, pages 1643-1656, New York, NY, USA, 2023. Association for Computing Machinery.

O. Sheffet. Locally private hypothesis testing. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018 , pages 4612-4621, 2018.

Sandeep Silwal, Sara Ahmadian, Andrew Nystrom, Andrew McCallum, Deepak Ramachandran, and Seyed Mehran Kazemi. Kwikbucks: Correlation clustering with cheap-weak and expensivestrong signals. In The Eleventh International Conference on Learning Representations , 2023.

Alistair Stewart, Ilias Diakonikolas, and Clément L. Canonne. Testing for families of distributions via the Fourier transform. In NeurIPS , pages 10084-10095, 2018.

Ananda Theertha Suresh, Alon Orlitsky, Jayadev Acharya, and Ashkan Jafarpour. Near-optimalsample estimators for spherical gaussian mixtures. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems , volume 27, 2014.

Gregory Valiant and Paul Valiant. An automatic inequality prover and instance optimal identity testing. SICOMP , 46(1):429-455, 2017.

Gregory Valiant and Paul Valiant. Estimating the unseen: Improved estimators for entropy and other properties. J. ACM , 64(6):37:1-37:41, 2017.

Paul Valiant. Testing symmetric properties of distributions. SIAM J. Comput. , 40(6):1927-1968, 2011.

A Other related works

General property testing of distributions: Testing properties of distributions has been extensively studied over the past few decades. Distribution testing under computational constraints has also been explored in [37, 70]. Hypothesis testing (and hypothesis selection) have received significant interest within the privacy community in machine learning [49, 21, 10, 4, 50, 72, 26, 27, 65]. Examples of other distributional properties that have been examined include testing monotonicity [20, 22, 11], testing histograms [24, 23], testing junta-ness [7, 32], and testing under structural properties [17, 58, 35, 44, 43].

Connections to tolerant testing: Tolerant testing asks us estimate the true TV distance between a known distribution and another which we only have sample access to, up to a small additive error. Readers familiar with strong lower bounds in tolerant testing (see e.g., [78, 77, 25] where it's shown that Ω( n/ log( n )) samples are needed in the case where the known distribution is uniform, compared to only O ( √ n ) samples needed for uniformity testing) might find our results surprising. It is incorrect to conclude that our algorithm in Section B, i.e. the search component (which is adaptive to the distance between ˆ p and p ), can perform tolerant testing between p and ˆ p . In reality, the algorithm finds an α which allows the testing component to terminate (either with accept or reject). Thus, while we know that the α found by our algorithm is never larger than ∥ p -ˆ p ∥ TV, it could in fact be much lower than ∥ p -ˆ p ∥ TV, meaning the α found by the search algorithm is not a good estimate for ∥ p -ˆ p ∥ TV.

Testable learning: Our framework bears some resemblance to testable learning as introduced in [71]. In this framework, the focus is on designing learning algorithms that can check whether the required underlying assumptions hold. If the assumptions do not hold, the algorithm may forego solving the problem. However, if it chooses to solve the problem, it must do so accurately, regardless of whether the assumptions hold. This is similar to our notion of testing with a suggested accuracy level, where the algorithm can either forego solving the problem if the assumption does not hold or solve it accurately regardless. Some examples of results in this framework are presented in [53, 54, 41, 55, 62, 61].

Algorithms with predictions: Recently, there has been a burgeoning interesting in augmenting classical algorithm design with learned information. Most relevant to us are works which study learning-augmented algorithms under sublinear constraints, such as memory or sample complexity [57, 59, 60, 33, 47, 48, 31, 64, 73, 2]. We refer to the interested reader to the website https: //algorithms-with-predictions.github.io/ for an up-to-date collection of literature on learning-augmented algorithms.

B Searching for the appropriate accuracy level

Early on, we defined the augmented tester for testing identity, uniformity, and closeness. We generalize this concept for other properties of distributions. More formally, a property P is a set of distributions, and we say a distribution has property P iff it is in P . The goal is to distinguish whether an unknown distribution p is in P , or it is far from all distributions that have the property. Similar to Definition 1.1, we define the notion of an augmented tester for P as follows:

Definition B.1. Suppose we are given three parameters, α ∈ [0 , 1] , ϵ ∈ (0 , 1) , δ ∈ (0 , 1) . Assume A is an algorithm that receives all these parameters, and the description of a known distribution ˆ p as input, and it has sample access to an unknown underlying distribution p . We say algorithm A is an ( α, ϵ, δ ) -augmented tester for property P for every p , q , and ˆ p over [ n ] :

If ˆ p and p are α -close in total variation distance, the algorithm outputs inaccurate information with a probability at most δ/ 2 .

If p ∈ P , then the algorithm outputs reject with a probability at most δ/ 2 .

If p is ϵ -far from every member of P , then the algorithm outputs accept with a probability at most δ/ 2 .

With this definition in mind, suppose we have an augmented tester for a property P . Our goal is to find an appropriate α such that the augmented tester solves the problem: it outputs accept or reject , but also, it does not use too many samples. In this section, we introduce a search algorithm that seeks to find this appropriate α . It is important to note that the α identified by our algorithm may not necessarily match the true accuracy level of the prediction, i.e., ∥ p -ˆ p ∥ TV. Instead, it corresponds to the minimum number of samples that the augmented tester can solve the problem.

Our search algorithm runs in rounds, each set by a sample budget that increases as the process progresses. In round i , the algorithm determines an appropriate value for α i and invokes the augmented tester, namely A , with this chosen α i . The value of α i is selected such that the augmented tester operates within the sample budget allocated for that round. If the tester outputs accept or reject , the search algorithm replicates this response. However, if the tester returns inaccurate information , we double the sample budget and proceed to the next round. The pseudocode for this procedure is provided in Algorithm 2. In the following theorem, we prove its performance.

Algorithm 2 An augmented property tester without knowledge of α

Theorem 6. Fix a parameter δ < 1 / 2 . Suppose A is an ( α, ϵ, δ ) -augmented tester for property P that receives a suggested accuracy level α and confidence parameter δ . Let f : [0 , 1] × [0 , 1] → N be a non-decreasing function for which A uses f ( α, δ ) 3 samples when it is invoked with parameter α and aims for the confidence parameter δ . Algorithm 2 is a ( ϵ, δ ′ ) -augmented tester, without knowledge of α , for property P where

In addition, if α ∗ is the true accuracy level, then Algorithm 2 uses O ( f ( α ∗ )) samples in expectation.

Proof. First, we define the notation we use in this proof. Let s min and s max represent the smallest and largest sample sizes used by A , respectively. s min may be one or higher, and s max is the sample size when the prediction did not make any improvements (the sample complexity of a standard tester). Let t := log 2 ( s max /s min ) . The algorithms runs in t +1 rounds i ∈ { 0 , 1 , . . . , t } . In round i < t , our sample budget is 2 i · s min . We run the algorithm with a parameter α i ensuring the sample complexity of A remains at most 2 i -1 · s min samples. In cases where multiple α values satisfy this criterion, we select the largest. We use (abuse in fact) the inverse function notation, defining α i as α i = f -1 ( 2 i · s min ) . If A finds an answer ( accept or reject ), we return the answer; Otherwise, we proceed to the next round. In round i = t , where we have s max ≈ 2 i · s min samples, we run the standard tester and return its answer.

Now, we focus on proof of correctness. Note that we (may) call A t times and the standard tester one time. Each of these tester works with probability at least 1 -δ . Thus, by the union bound, we can assume that they return the correct answer with probability at least 1 -δ ′ where δ ′ := δ/ ( t +1) . As we have noted in Definition 1.1, A is resilient to inaccuracies, implying that even if the suggested accuracy level is not valid, if the tester does not output a false accept or reject with probability more than δ/ 2 . The same statement is correct for the standard tester. Now, our algorithm here replicates the

output that is produced by one of the testers. Thus, if they all of them outputs the correct answer, our algorithm outputs the correct answer as well. And, this event happens with probability at least 1 -δ ′ .

Next, we focus on the analysis of the sample complexity of this algorithm. Let α ∗ be the true prediction accuracy, and let i ∗ be the first round where α i ≥ α ∗ . Recall that we assume that α i is the largest α such that f ( α i ) ≤ 2 i · s min . In addition, we assume that f is non-decreasing. Thus, we have 4 :

If the algorithm ends at round i , we have used the following number of samples:

On the other hand, the probability that algorithms end at round i > i ∗ cannot be too high. If the algorithm ends at round i or later, it means that A in rounds i ∗ , i ∗ +1 , . . . , i -1 must have returned inaccurate information . However, given our assumption, the true accuracy is not worse than α i , which makes outputting inaccurate information wrong. And, this event does not happen in each round with probability more than δ/ 2 for each round independently. Let I end indicate a random variable that indicates the index of the round for which the algorithms ends. Thus, we have for every i ≥ i ∗ :

Now, we are ready to bound the expected value of the number of samples we use. Let S denote the number of samples we use. Then, we have:

Next, we write the second sum in terms of a new variable k := i -i ∗ -1 . Then, we swap the order of the summations on k and j :

C Identity and uniformity testing

In this section, we focus on the problem of identity testing, which involves testing the equality between a known distribution q and an unknown distribution p . Specifically, our goal is to determine whether p = q or if they are ϵ -far from each other. In an augmented setting, we are provided with another known distribution ˆ p , predicted to represent p , along with a suggested level of accuracy α . Surprisingly, our findings indicate that the sample complexity for this problem is influenced by a new parameter: the total variation distance between q and ˆ p , denoted by d . In particular, we have the following theorem:

Theorem 7. Fix two arbitrary parameters ϵ ∈ (0 , 1] and α ∈ [0 , 1] . Algorithm 3 is an ( α, ϵ, δ = 0 . 1) -augmented tester for identity that uses the following number of samples s , where:

and d refers to ∥ q -ˆ p ∥ TV. For any α ∈ [0 , 1] , and ϵ ∈ [0 , 1 / 2] , we show the same number of samples is necessary for any ( α, ϵ, δ = 2 / 3) -augmented identity tester. In fact, the lower bound holds even when q is a uniform distribution over [ n ] .

The proof of this theorem follows from Proposition C.1, Proposition C.2, and Proposition C.3.

Remark 8. A particular instance of this problem is uniformity testing , where q is a uniform distribution over [ n ] . In Theorem 7, our upper bound applies to any arbitrary q . Furthermore, our tight lower bound is based on a hard instance, where q is a uniform distribution. These results establish optimal upper and lower bounds for both identity testing and uniformity testing simultaneously.

C.1 Upper bound for identity testing

Our upper bound relies on a fundamental observation regarding the total variation distance: this distance is the maximum discrepancy between the probability masses that two distributions assign to any subset of their domain. Demonstrating a small total variation distance between two distributions entails proving that the discrepancy across every domain subset is minimal. In contrast, to prove a large total variation distance, one only needs to identify a single subset with a significant discrepancy. We employ the Scheffé set of ˆ p and q -defined as the set of elements x ∈ [ n ] where ˆ p ( x ) < q ( x ) , denoted by S -to serve as a witness for the farness of p from either q or ˆ p . It is known that S maximizing the discrepancy between the probability masses of ˆ p and q , implying:

In scenarios where the distance d > α , accurately estimating the probability of S according to p with an accuracy of ( d -α ) / 4 provides a basis for distinguishing the farness from either q or ˆ p . If the probability masses assigned to S by q and p differ by more than ( d -α ) / 4 , it clearly indicates that p is not identical to q , leading to outputting reject . Conversely, if the estimated probability of p ( S ) is close to q ( S ) , then the discrepancy between p ( S ) and ˆ p ( S ) must exceed α , leading to outputting inaccurate information . The pseudocode of our algorithm presented in Algorithm 3. We prove the performance of our algorithm in the following proposition:

Proposition C.1. Fix two arbitrary parameters ϵ ∈ (0 , 1] and α ∈ [0 , 1] . Algorithm 3 is an ( α, ϵ, δ = 0 . 1) -augmented tester for identity that uses the following number of samples s , where:

and d refers to ∥ q -ˆ p ∥ TV.

Proof. The proof of theorem is trivial in the setting where d ≤ α and 1 ( d -α ) 2 > √ n ϵ 2 . In these cases, we use the standard tester for identity testing (e.g., [76, 30, 5]) with δ ′ = δ/ 2 = 0 . 05 . Clearly, a wrong answer was produced with probability less than δ/ 2 . It is known that these tester only use O ( √ n/ϵ 2 ) samples.

Now, suppose d > α . Simple application of Chernoff's bound shows that one can estimate the probability of the Scheffé set of q and ˆ p , S , up to error ( d -α ) / 4 with probability 1 -0 . 95 using O (1 / ( d -α ) 2 ) samples. We refer to this estimate as σ , and we have with probability 0.95:

Hence, we do not output reject with probability more than 0 . 05 = δ/ 2 . On the other hand, assume ∥ ˆ p -p ∥ TV ≤ α , then we show it is unlikely for us to output inaccurate information .

(by triangle inequality and Eq. (6))

Note that if p = q , then

Therefore, we obtain that with probability 0.95:

Hence, we do not output inaccurate information with probability more than 0 . 05 = δ/ 2 . Hence, Algorithm 3 is an ( α, ϵ, δ = 0 . 1) -agumented tester for identity. The sample complexity of the algorithm, in the case where α < d and 1 ( d -α ) 2 ≤ √ n ϵ 2 is O (1 / ( d -α ) 2 ) . Thus, the proof is complete.

C.2 Lower bound for uniformity testing when α ≥ d

In this section, we focus on the case where α ≥ d . We observed for this problem the number of samples depends on an additional parameter: the distance between q and ˆ p . If this distance is at most α then prediction does not help the algorithm. The algorithm would still be required to draw Ω( n/ϵ 2 ) samples (which is the sufficient amount of sample for the standard case). Our proof is based on a reduction from standard identity testing to the augmented version of this problem which establishes the desired lower bound.

Proposition C.2. Suppose we are given two known distributions q and ˆ p , and an unknown distribution p over [ n ] . Let d := ∥ q -ˆ p ∥ TV. For any α in [0 , 1] , if α ≥ d , then any ( α, ϵ, δ = 2 / 3) -augmented tester requires Ω( √ n/ϵ 2 ) samples.

Proof. We prove that standard identity testing can be reduced to augmented version of this problem if α ≥ d . For the standard tester, we are given a known distribution q and an unknown distribution p . Let ˆ p be any arbitrary distribution that is in distance d from q . Consider A as an ( α, ϵ, δ = 2 / 3) -augmented tester. The procedure for the standard tester is straightforward: Run A on p , ˆ p , and q . If it returns accept , we also return accept ; if it returns reject or inaccurate information , we return reject .

q

d

ˆ

p

α

ϵ

accept

reject

reject

inaccurate information

We show that the augmented tester distinguishes between the cases where p = q and ∥ p -q ∥ TV > ϵ . See Figure 2. More precisely, Note that if p = q , then p is within α distance of ˆ p . Thus, the only valid answer of A in this case is accept . In this case, according to Definition 1.1, A returns reject and inaccurate information each with a probability of at most δ/ 2 = 1 / 6 . Therefore, by the union bound, we return accept with a probability of at least 2 / 3 . When p and q are ϵ -far from each other, A returns accept with a probability of at most δ/ 2 = 1 / 6 . Consequently, we output the correct answer with a probability greater than 2/3.

This reduction indicates that A must use at least as many samples as required for identity testing. Considering the existing lower bound for uniformity testing (where q is a uniform distribution over [ n ] ), augmented identity testing necessitates O ( √ n/ϵ 2 ) samples [69].

C.3 Lower bound for uniformity testing when α < d

In this section, we consider the lower bound for the uniformity testing problem in the setting where α < d . On the other hand, if d -α is not too small, the required lower bound is only Ω(1 / ( d -α ) 2 ) . Otherwise, Ω( √ n/ϵ 2 ) samples is needed (as it is required for the standard tester).

At a high level, our proof consists of the following steps. First, we construct three distributions that look similar when we draw too few samples. We formalize the similarity of these three distributions using a multivariate coupling argument. Next, we describe valid answers for each of the distributions. The main message of this part is that there is no possible answer that is valid for all three distributions. Now, (similar to Le Cam's method), suppose we feed the algorithm with samples from one of these three distributions (each with probability 1/3). For any sample set, the algorithm outputs an answer (which may be randomized), however, this answer is considered as wrong for at least one of the underlying distributions. This is due to the fact that there is no universally valid answer that is simultaneously correct for all three distributions. Hence, if the algorithm outputs a valid answer with high probability, it must be able to distinguish the underlying distributions to some degree. On the other hand, the indistinguishably result says, it is impossible to tell these distributions apart. Thus, we reach a contradiction. And, the lower bound is concluded. More formally, we have the following proposition:

Proposition C.3. Suppose we are given two known distributions q and ˆ p , and an unknown distribution p over [ n ] . Let d := ∥ q -ˆ p ∥ TV. For any α ∈ [0 , 1] , ϵ ∈ [0 , 1 / 2] , and d ∈ [0 , 1 / 2] , if d > α any ( α, ϵ, δ = 2 / 3) -augmented algorithm for testing identity of p and q with prediction ˆ p requires s = Ω ( min ( 1 ( d -α ) 2 , √ n ϵ 2 )) samples.

Proof. Without loss of generality, assume d -α ≤ 0 . 4 . Otherwise, we are required to establish a lower bound of Ω(1) , which is necessary for any non-trivial testing problem. 5 As we have discussed it early, we prove this proposition in the following steps:

Construction of distributions: We assume q = U n is a uniform distribution over [ n ] . Without loss of generality assume n is even. Otherwise, we can set the probability of one of the elements in all the distributions to zero. We define the prediction distribution as follows for every i ∈ [ n ] :

It is not hard to assert that ∥ ˆ p -U n ∥ TV is d satisfying the assumption we had in the statement of the proposition. Next, we construct two distributions p · and p ⋄ .

Suppose we have a random vector Z in {-1 , +1 } n/ 2 where each coordinate Z i is one with probability 1/2 independently. Now, we define the following distributions over [ n ] :

It is not hard to see that the probabilities of each distributions sum up to one by observing the probability of two consecutive elements is 2 /n .

Indistinguishably of the distributions Fix the number of samples s , and let S | p denote a random variable that is a sample set of size s drawn from p . We show that S | U n , S | p ⋄ , and S | p · are three random variables with small total variation distances between each pair. These distances are so small that it is practically impossible for the algorithm to tell them apart. We formalized this by providing a multivariate coupling between these random variables. We extend Le Cam's method for three random variables, and show that no algorithm with low error probability exists unless s is large, establishing the desired lower bound for the problem.

Lemma C.4. Suppose we are given the following parameters ϵ, d, α, n , and s . Assume d -α ≤ 0 . 4 . Let U n denote the uniform distribution over [ n ] , and let p · and p ⋄ be the distributions defined in Equation (7) and Equation (8) . Let S | U n , S | p ⋄ , and S | p · be three random variables that are sample sets of size s from each of these distributions. If s ≤ min ( 0 . 00005 ( d -α ) 2 , 0 . 004 √ n ϵ 2 ) , then there exists a distribution over triples of three sample sets of size s (that is ([ n ] s ) 3 ), which we call a multivariate coupling C , between S | U n , S | p ⋄ , and S | p · such that:

The marginals of C correspond to the three probability distributions U ⊗ s n , p ·⊗ s , and p ⋄⊗ s . More precisely, for a sample ( S 1 , S 2 , S 3 ) drawn from C , and every S ∈ [ n ] s , we have:

Pr ( S 1 ,S 2 ,S 3 ) ∼C [ S 1 = S 2 = S 3 ] ≥ 0 . 98 .

For the proof, see Section C.3.1. This lemma states that when the number of samples is too small, the sample sets looks very similar in all three cases: p = U n , p = p · , and p = p ⋄ .

q

=

U

n

d

ˆ

p

α

ϵ

p

⋄

p

•

cannot output

inaccurate information

cannot output

accept

cannot output

reject

Valid answers for each distributions: We focus on the valid output of the algorithm in these three cases using the Definition 1.1. According to this definition, we have:

If p = q , then reject is not a valid answer.

If ∥ p -q ∥ TV > ϵ , then accept is not a valid answer.

If ∥ p -ˆ p ∥ TV ≤ α then inaccurate information is not considered as a valid answer.

We denote the set of valid output for each case by V | p . An accurate augmented tester with confidence parameter δ , has to output an answer in V | p with probability at least 1 -δ .

p = U p = U n p = U : In this case we have p = q = U . Clearly, in this case, reject is not a valid answer:

p = p · p = p · p = p · : The total variation distance between q = U n and p · is ϵ > ϵ . Thus, accept is not considered a valid output:

p = p ⋄ p = p ⋄ p = p ⋄ : It is not hard to see that ∥ p ⋄ -ˆ p ∥ TV = α . Thus, inaccurate information is not a valid answer:

The diagram of the invalid outputs is shown in Figure 3 for an instance of these distributions. Let us explain what these valid sets are implying. The intersection of the three valid sets is empty, because among all possible outputs each set lacks at least one of them :

Therefore, there is no possible output that is considered valid on all the three distributions 6 . And, since these distributions are indistinguishable with high probability, no algorithm should perform well under all three possibilities.

Deriving a contradiction. Our proof is via contradiction. Suppose there exists an ( α, ϵ, δ ) -augmented tester, called A , for uniformity testing that uses s samples.

Let I represent the internal coin tosses of the algorithm. Since the probability of A outputting an invalid answer is bounded by δ , we have:

We can write the above equation for p · and p ⋄ and sum of them up. Thus, we obtain:

For the last line above, note that when S 1 = S 2 = S 3 , the terms inside the expectations are the sum of probabilities of all possible outputs for A . Hence, the sum of these probabilities is one.

Recall that earlier, we assumed that δ = 0 . 3 , and using the properties of the coupling, we know that S 1 = S 2 = S 3 with probability at least 0.98. Thus, we get:

Hence, we reach a contradiction. Thus, no such algorithm like A exists.

C.3.1 Multivariate coupling between the three hard distributions

In this section, we prove the lemma on coupling that we used to prove our lower bound. This lemma implies indistinguishability between S | U n , S | p · , and S | p ⋄ . For a definition and some basics on coupling, see Section F.3.

Lemma C.4. Suppose we are given the following parameters ϵ, d, α, n , and s . Assume d -α ≤ 0 . 4 . Let U n denote the uniform distribution over [ n ] , and let p · and p ⋄ be the distributions defined in Equation (7) and Equation (8) . Let S | U n , S | p ⋄ , and S | p · be three random variables that are sample sets of size s from each of these distributions. If s ≤ min ( 0 . 00005 ( d -α ) 2 , 0 . 004 √ n ϵ 2 ) , then there exists a distribution over triples of three sample sets of size s (that is ([ n ] s ) 3 ), which we call a multivariate coupling C , between S | U n , S | p ⋄ , and S | p · such that:

The marginals of C correspond to the three probability distributions U ⊗ s n , p ·⊗ s , and p ⋄⊗ s . More precisely, for a sample ( S 1 , S 2 , S 3 ) drawn from C , and every S ∈ [ n ] s , we have:

Pr ( S 1 ,S 2 ,S 3 ) ∼C [ S 1 = S 2 = S 3 ] ≥ 0 . 98 .

Proof. We start by bounding the total variation distance between two pairs of distributions: ∥ U ⊗ s n -p ·⊗ s ∥ TV, and ∥ U ⊗ s n -p ⋄⊗ s ∥ TV.

In [69], the author has shown that the total variation distance between S | U n and S | p · is bounded by:

The last inequality above comes from our assumption that s ≤ 0 . 004 √ n ϵ 2 .

Next, we focus on the total variation distance between S | p ⋄ and S | U n . Note that p ⋄ is a distribution that has some bias towards the even elements, but it is uniform on both sets of odd and even elements. Therefore, to bound the total variation distance between U ⊗ s n and p ⋄⊗ s , we use a similar argument for showing a lower bound of Ω(1 / ∆) 2 to distinguish whether a coin is fair or it has a bias of (1 + ∆) / 2 . Below is our formal argument.

It is not hard to see that one can use the following process to generate a sample from p ⋄ . First, we draw a sample X ∼ Ber (1 / 2 + ( d -α )) from the Bernoulli distribution with a success probability 1 / 2 + ( d -α ) . If X = 1 , we pick an even element (uniformly); Otherwise, we pick an odd element. Similarly, we can draw samples from a uniform distribution as follows: draw X ′ ∼ Ber (1 / 2) from the Bernoulli distribution with a success probability 1 / 2 . If X ′ = 1 , we pick an even element (uniformly); Otherwise, we pick an odd element. One can view this process of generating a sample in [ n ] from a binary variable ( X or X ′ ) as a channel . The data processing inequality says that the distance (more precisely any f -divergences) between two random variables does not increases after they pass through a channel. Hence, we use this fact and bound the total variation distance between p ⋄⊗ s and U ⊗ s n :

Since we had this assumption where s is at most 0 . 00005 ( d -α ) 2 , we get.

So far, we have shown that ∥ U ⊗ s n -p ·⊗ s ∥ TV and ∥ U ⊗ s n -p ⋄⊗ s ∥ TV are at most 0.01. Using the coupling lemma (See Fact F.3), there exist two maximal couplings C · and C ⋄ between the following pairs of distributions: ( U ⊗ s n , p ·⊗ s ) and ( U ⊗ s n , p ⋄⊗ s ) . The properties of these maximal couplings are:

The marginals of the couplings are equal to the two pairs of probability distributions. That is:

Pr ( S 1 ,S 2 ) ∼C · [ S 1 = S 2 ] ≥ 0 . 99 , and Pr ( S 1 ,S 2 ) ∼C ⋄ [ S 1 = S 2 ] ≥ 0 . 99 .

For every S ∈ [ n ] m , let C ⋄ 2 | S be a probability distribution over [ n ] m . The probability of S ′ according to C ⋄ 2 | S is:

With these definitions in mind, we create our multivariate coupling C by essentially joining the two couplings C · and C ⋄ . The following randomized procedure produces a random sample ( S 1 , S 2 , S 3 ) from C :

Let ( S 1 , S 2 ) be a sample from C · .

Let S 3 be a sample from C ⋄ 2 | S 1 .

Output ( S 1 , S 2 , S 3 ) as a sample from C .

Given the definition of C · , it is clear that the first two marginals of C are exactly U ⊗ s n and p ·⊗ s . To see that the third marginal of C is equal to p ⋄ , we have the following for every S ′ ∈ [ n ] m :

Now, we show that S 1 , S 2 , and S 3 are equal with high probability. This stems from the fact that C · and C ⋄ were maximal couplings, and with high probability, the pairs drawn from them are equal. More formally, via the union bound, we have:

̸

̸

̸

̸

̸

̸

D Augmented closeness testing

In this section, we focus on the problem of closeness testing in the augmented setting. More formally, we have the following theorem.

Theorem 9. For every α and ϵ in (0 , 1) , there exists an algorithm for augmented closeness testing that uses Θ ( n 2 / 3 α 1 / 3 ϵ 4 / 3 + √ n ϵ 2 ) samples and succeeds with probability at least 2 / 3 . In addition, any algorithm for this task is required to use Ω ( n 2 / 3 α 1 / 3 ϵ 4 / 3 + √ n ϵ 2 ) samples.

For the proof of the upper bound see section 3. And, for the proof of lower bound see section D.1.

D.1 Lower bound for closeness testing

In this section, we prove a lower bound for augmented closeness testing where a known distribution ˆ p is provided satisfying ∥ p -ˆ p ∥ TV = α . We have the following theorem.

Theorem 10. Suppose we are given α and ϵ in (0 , 1) . Any ( α, ϵ, δ := 11 / 24) -augmented tester for closeness testing of distributions over [ n ] uses Ω ( n 2 / 3 α 1 / 3 ϵ 4 / 3 + √ n ϵ 2 ) samples. The lower bound holds even when α is provided to the algorithm.

Proof. The Ω( √ n/ϵ 2 ) term comes from the existing lower bound for uniformity testing since we can reduce uniformity testing to augmented closeness testing as follows: Suppose we wish to test whether a distribution q over [ n ] is uniform or ϵ -far from uniform. Suppose A is an ( α, ϵ, 11 / 24) -augmented closeness tester. Set p to be the uniform distribution over [ n ] . Let ˆ p be any distribution that is α -far from uniform. For instance, ˆ p can be a distribution that is (1 + α ) /n on half of the elements and (1 -α ) /n on the other half. Since A is an augmented tester, it can be used to distinguish whether p = q or ∥ p -q ∥ TV > ϵ with probability 11 / 24 .

Upon receiving ˆ p and samples from p and q , if A returns inaccurate information or reject , we output reject ; otherwise, we output accept . Given Definition 1.1, if q is uniform, we will not output inaccurate information or reject with probability more than 11 / 24 . And, if q is ϵ -far from the uniform distribution (equivalently p ), we will not output accept with probability more than 11 / 48 . Hence, we have a tester with confidence parameter 11 / 48 for uniformity testing. Thus, s must be Ω( √ n/ϵ 2 ) , which is the number of samples required for uniformity testing [69].

To prove the lower bound of Ω ( n 2 / 3 α 1 / 3 /ϵ 4 / 3 ) samples, we only need to focus on the case where α ≥ ( √ nϵ 2 ) -1 . Otherwise, n 2 / 3 α 1 / 3 /ϵ 4 / 3 = O ( √ n/ϵ 2 ) , making the lower bound trivial due to the Ω( √ n/ϵ 2 ) lower bound shown earlier. We consider the following cases depending on ϵ :

Case 1: ϵ ≤ ( √ nϵ 2 ) -1 . Therefore, in this case, the interesting regime of parameter for α is that α ≥ ϵ . We prove the lower bound for this case in Theorem 11 in Section D.1.1.

Case 2: ϵ > ( √ nϵ 2 ) -1 . In this case, in addition to considering the case where α ≥ ϵ , we also need to study the case where α is in ( ( √ nϵ 2 ) -1 , ϵ ) . We prove the lower bound for this case in Theorem 12 in Section D.1.2.

This completes the proof.

D.1.1 Lower bound for α ≥ ϵ

Below, we provide a lower bound for augmented closeness testing when α ≥ ϵ . At a high level, Theorem 11 achieves this lower bound by creating a challenging instance. This involves embedding the difficult instances of the standard closeness testing problem into another distribution, spanning a domain of [ n -1] elements, while concentrating the majority of the distribution's mass on the last

element. The prediction, being a singleton distribution centered on the last element, does not facilitate solving the standard closeness testing problem, which is embedded within the first [ n -1] elements.

Theorem 11. Suppose we are given α and ϵ in (0 , 1) . Assume we have two unknown distributions p and q and a known distribution ˆ p over the domain [ n ] . If α ≥ ϵ , then any ( α, ϵ, δ := 11 / 24) -augmented tester for testing closeness of p and q requires the following number of samples:

Proof. Our proof is based on a reduction from standard closeness testing to augmented closeness testing. Assume there exists an algorithm, A , that upon receiving ˆ p and s samples from p and q , can distinguish whether p = q or ∥ p -q ∥ TV ≥ ϵ with probability 11/24. We show that this algorithm can then be used to test closeness of two distributions, ˜ p and ˜ q over [ n -1] with proximity parameter ϵ ′ := ϵ/α .

Suppose we have sample access to ˜ p and ˜ q . We construct p and q based on ˜ p and ˜ q as follows:

Note that given a sample from ˜ p (respectively ˜ q ), we can generate a sample from p (respectively q ) by outputting that sample with probability α and n otherwise. Also, if ˜ p is equal to ˜ q , then p and q are equal. And, if ∥ ˜ p -˜ q ∥ TV ≥ ϵ ′ , then ∥ p -q ∥ TV ≥ ϵ due to the following:

Let ˆ p be the singleton distribution on n (i.e., ˆ p n = 1 ). In this case, ∥ p -ˆ p ∥ TV equals α . We feed A with ˆ p and samples from p and q . Now, if A outputs p = q , we declare ˜ p = ˜ q . Otherwise, we declare ˜ p is ϵ ′ -far from ˜ q . Given our construction, this answer is true with probability 11/24.

Note that in this process, in expectation, we draw αs samples from ˜ p and ˜ q . Thus, using Markov's inequality, the probability of drawing more than 100 αs samples from each of ˜ p and ˜ q is at most 0 . 01 . Hence, using A , there exists an algorithm for testing closeness of ˜ p and ˜ q with probability 1 -(11 / 24 + 0 . 01) > 0 . 53 that uses 100 αs samples from ˜ p and ˜ q . Using the existing lower bound for the closeness testing problem (see Theorem 4 in [69] and Proposition 4.1 in [30]), we have:

Thus, by dividing by 100 α , we get the following bound as desired:

D.1.2 Lower bound for α in ( c · ( √ nϵ 2 ) -1 , ϵ )

We extend the existing lower bound for closeness testing to the augmented setting. In the standard setting, the lower bound for closeness testing is derived from the hard instances for uniformity testing with one key adjustment: introducing elements with large probability ( ≈ n -2 / 3 ) in the distributions. These large elements share identical probability masses in both p and q , indicating that they do not contribute to the distance between the two distributions. However, their presence in the sample set confuses the algorithm: Because of their large probabilities, their behaviour in the sample set may be deceivingly hint 'not uniform', making it difficult for the algorithm to decide whether the rest of the distributions are uniform or not. Therefore, the algorithm needs s ≈ n 2 / 3 samples to first identify these large elements, and then it can test the uniformity on the rest of the distributions. The surprising part about this result is that s ≈ n 2 / 3 is much larger than the √ n samples which suffice for testing uniformity.

The challenge in our case is that ˆ p may give away the large elements to the algorithm. Hence, to prove the lower bound, we create hard instances by adding as many large elements as possible, without changing ˆ p by limiting the overall probability mass of the large elements to α .

Here is our instance: we assume ˆ p is the uniform distribution; p assigns ≈ (1 -α ) /n probability mass to O ( n ) elements chosen at random, and assigns ≈ n -2 / 3 probability mass to α · n 2 / 3 many elements in the domain. Now, q has two options. Half the time, q is equal to p . The other half, q shares the same large elements, but it is not quite uniform on the rest of the distributions. In particular, q assigns probabilities (1 ± Θ( ϵ ))(1 -α ) /n . to the chosen O ( n ) elements randomly, making it ϵ -far from p . 7 Using this construction, we use a result of Valiant [78] to show that these two cases are indistinguishable unless we draw Ω( n 2 / 3 α 1 / 3 ) samples. More precisely, we have the following theorem:

Theorem 12. Suppose we are given the following parameters: a positive integer n , ϵ ∈ (0 , 1 / 6) , and α ∈ ( ( √ nϵ 2 ) -1 , ϵ ) (with this implicit assumption that ( √ nϵ 2 ) -1 is less than ϵ ). There exists a distribution ˆ p and a family of pairs of distribution p and q such that any augmented algorithm for closeness testing must use

samples.

Proof. Let k be an integer that is at most

for a sufficiently small absolute constant c k which we determine later. Assume we have k samples from p and q that are available to the algorithm. We show that there are two (families of) distributions that any augmented tester has to distinguish them with high probability. However, they are information-theoretically indistinguishable using only k samples.

Setting up parameters: First, we focus on determining a series of parameters which we use in our proof. We define ℓ as follows:

̸

It is not hard to see that ℓ is an integer. For our proof, we require that ℓ and n are either both even or both odd. Here, the role of the indicator variable is to reduce ℓ by one in case they are not. Also, ℓ is positive, since due to the range of α and ϵ , we have:

Given this definition, there exists an absolute constant c ℓ < 0 . 5 , such that ℓ = c ℓ · n 2 / 3 α 4 / 3 /ϵ 4 / 3 . We use this identity as the definition of ℓ throughout this proof. Given the assumption about the range of α , we also have:

Let L denote a set of ℓ random elements in [ n ] to be our large elements. L := [ n ] \ L denotes the set of small elements. Since n and ℓ are either both even or both odd, L has even number of elements. Now, partition L to two random subset of the same size L 1 and L 2 . Set ϵ ′ := 6 ϵ < 1 . Set

ˆ

p

=

U

p

+

p

-

L

L

1

/n

λ

small

λ

large

λ

large

(1 +

ϵ

′

)

·

λ

small

(1

-

ϵ

′

)

·

λ

small

Constructing the distributions: We set ˆ p to be a uniform distribution over [ n ] . That is, ˆ p i = 1 /n for every i ∈ [ n ] . We construct two distributions p + and p -as follows. For every i ∈ [ n ] , the probability mass of i according to these distributions are as follows:

For a visual representation of these distributions, refer to Figure 4. It is not hard to verify that the probability masses sum up to one since | L | = ℓ , and ∣ ∣ L ∣ ∣ = n -ℓ . Also, these probabilities are non-negative due to Equation (13) and the fact that α < ϵ ≤ 1 / 6 :

Also, given these definitions, p + is α -close to ˆ p , and p + and p -are ϵ -far from each other:

Proof of indistinguishability: Now, consider two pairs of distributions ( p = p + , q = p + ) and ( p = p + , q = p -) . Given the distances above any algorithm for augmented closeness testing should output accept on ( p + , p + ) , and reject on ( p + , p -) . To establish our lower bound, we show these pairs of distributions are indistinguishable with high probability. We use the moment based indistinguishably result for symmetric properties 8 of pairs of distributions in [78].

Before stating their theorem, we need to define the notion of the ( k, k ) -based moments of a pair distributions ( p, q ) as follows:

Definition D.1 ( ( k, k ) -based moments) . For every positive integer k the ( k, k ) -based moments of a pair of distributions p and q are the following values:

for every non-negative integers a and b .

Fact D.2 (Adapted from [78]) . Given an integer k , suppose we have four distributions, p (1) , p (2) , q (1) , and q (2) over [ n ] where the maximum probability that each of them assigns to an element is less than 1 / (1000 k ) . If

then no algorithm can distinguish the pair ( p (1) , q (1) ) from ( p (2) , q (2) ) with a success probability more than 13 / 24 upon receiving Poi( k ) samples from each distribution in the pair.

First, we verify that all the domain elements of p + and p -have probability masses at most 1 / (1000 k ) according to both p + and p -. Recall that ϵ < 1 . Observe that it suffices to show that 1 /n < 1 / (2000 k ) and α/ℓ < 1 / (2000 k ) . The first inequality holds due to the fact that α is in ( √ nϵ 2 ) -1 , ϵ ) :

And, the second inequality comes from the definitions of ℓ , k , and c k :

For the rest of this proof, we focus on showing Equation (14) for the pairs ( p + , p + ) and ( p + , p -) . For simplicity in our equations, we use the following notation:

Using the Definition D.1, it is not hard to see that:

Therefore, we get:

symmetric property since permuting the labels of the elements does not change the distance between p and q ; Furthermore, ˆ p is identical for every permutation of domain elements.

On the other hand, we have:

Using the two equations above, we can bound the right hand side of Equation (14):

Note that for b = 0 or 1 , the term in the sum is zero. Hence, we only need to iterate over all b ≥ 2 and a ≥ 0 . Let T denote √ k/λ large · λ small. Then, we can write the above equation in the following form:

To bound these terms, we have the following lemma:

Lemma D.3. For x ∈ R , we have

and for x ∈ [0 , 1 . 33] , we have:

Proof. We start by splitting the terms corresponding to odd and even i 's in the sum, and obtain the following:

The last equality above is due to the Taylor expansion of e y = 1 + y/ 1! + y 2 / 2! + . . . about y = 0 . Next, we have:

Now, if 0 < y ≤ 1 . 79 , then the following inequality holds: y < e y -1 < y + y 2 . Thus, for x ∈ (0 , 1 . 33) , we have:

Now, assume T · (1 + ϵ ′ ) ≤ 1 . 33 . We continue bounding the right hand side of Equation (14) via the above fact:

We can bound T from above given the trivial bounds we have that λ small ≤ 1 /n and λ large ≥ α/ℓ by their definitions in Equation (15):

The second to last inequality above is due to the range we assumed for α : it is greater than ( √ nϵ 2 ) -1 . For the last inequality, we assumed that α < ϵ . Moreover, using c k ≤ (5080320 · c ℓ ) -1 , we have √ c ℓ · c k < 6 . Therefore, we obtain:

For the last inequality, recall that early on we set k to be at most c k ( n 2 / 3 α 1 / 3 /ϵ 4 / 3 ) and c k is defined to be at most (5080320 · c ℓ ) -1 . Therefore, the requirements in Fact D.2 hold as desired in our setting. Thus, the proof is complete.

E Empirical Evaluation

We evaluate our algorithm empirically on real and synthetic data on the closeness testing problem. Recall that given samples to two unknown distributions p and q over the domain [ n ] , we wish to test if p = q or if ∥ p -q ∥ TV ≥ ϵ . They are referred to as the ( p, p ) case or the ( p, q ) case respectively. Our experiments test the performance of our algorithm and baselines.

Datasets.

Hard Instance : This is a simplified version of the hard instance described in Section D.1 for Theorem 12. A conceptually similar hard instance is also used in the classical closeness testing setting without any learned augmentation, as detailed in Section D.1.2

Let [ n ] be the domain. In the instance used in the experiments, p and q agree on the first m := ⌊ n 2 / 3 ⌋ elements. Both distributions have probability mass 1 / (2 m ) on these domain elements. p

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

p

p

TV

5

10

15

20

25

30

35

% Error

Augmented Tester,

5

10 3

Samples

Augmented Tester

100

0

100

200

300

400

500

600

Estimator Value

0

5

10

15

20

Frequency

Augmented Tester,

10 4

Samples

p

=

q

p

q

TV

1/2

(a) Distribution of

Z

(augmented)

500

0

500

1000

1500

Estimator Value

0

5

10

15

20

25

Frequency

Un-augmented Tester,

10 4

Samples

p

=

q

p

q

TV

1/2

(b) Distribution of

Z

(un-augmented)

also has probability mass 2 /n on domain elements between n/ 2 and 3 n/ 4 and q has probability mass 2 /n on domain elements between 3 n/ 4 and n . All other probability masses are 0 . We can check that ∥ p -q ∥ TV = 1 / 2 . The first m domain elements represent 'heavy hitters' which conspire to fool any testing algorithm and the rest of the domain are items which contribute to all the TV distance but are rarely sampled. Indeed, we built upon this intuition to prove a formal lower bound in Theorem 12.

Weconsider hints of the form ˆ p = (1 -β ) · p + β · Unif ( n ) , where Unif ( n ) is the uniform distribution on [ n ] elements and α ∈ [0 , 1] is an interpolation parameter. α = 0 corresponds to a perfect hint and the quality of the hint degrades as α increase. When α = 1 , we receive ˆ p = Unif ( n ) , i.e. a hint which not require any learned information to instantiate. In our experiments, we set n , the domain size, to n = 10 6 and our goal is to determine if we are sampling from ( p, p ) or ( p, q ) .

IP : The data is internet traffic data collected at a backbone link of a Tier1 ISP in a data center 9 . Using the preprocessing in [1], we construct the empirical distribution over source IP addresses in consecutive chunks of time. Each distribution represents approximately 170 ms of empirical network traffic data. The support size of each distribution is ≈ 45 , 000 while the total domain size is > 2 . 5 · 10 6 . Distributions curated from network traffic closer in time are closer in TV distance. This dataset has been used in other distribution testing literature [1, 48].

100

200

300

400

500

600

0

10

20

30

40

% Error

IP Traffic Data

Closeness Tester

Augmented Tester

CRS '15

(a) Error vs samples

100

200

300

400

500

600

0

5

10

15

20

25

30

35

% Error

IP Traffic Data

Closeness Tester

p

1

p

2

(b) Comparison between two predictors

We set p to be distribution for the first chunk of time and we let q be the distribution for the last chunk of time. Empirically, we have ∥ p -q ∥ TV ≈ 0 . 72 . The hint is for p , denoted as ˆ p , is the second chunk of time. Unlike the previous case, the hint is relatively far from p in terms of TV distance ( ∥ p -ˆ p ∥ TV ≈ 0 . 59 ). Nevertheless, they share similar 'heavy' items (IP addresses with large probability mass), since p and ˆ p represent distributions of network traffic which are close in time.

Baselines. We compare our main algorithm with the following baselines.

CRS '15 . This is the state of the art algorithm with access to almost-perfect predictions from [28]. More precisely, they assume predictions to both distributions p and q , denoted as ˆ p and ˆ q . Furthermore, they require ˆ p ( i ) = (1 ± ε/ 100) p ( i ) and ˆ q ( i ) = (1 ± ε/ 100) q ( i ) for all i ∈ [ n ] . Note that we only assume access to (much weaker) predictions for only p .

Closeness Tester . This is the state of the art algorithm of [42] without any predictor access. We refer to this algorithm as the standard or the un-augmented tester as well.

Error Measurement. At a high level, both our algorithm (denoted as 'Augmented Tester' and the prior SOTA un-augmented approach, (abbreviated as 'Closeness Tester') compute an estimator Z based on the samples requested. Then, both algorithms threshold the value of Z to determine which case we are in (either p = q or ∥ p -q ∥ TV ≥ ε . Note that Z is a random variable in both cases (since it depends on random samples). Ideally, the distribution of Z is well-separated in the two different hypotheses.

Thus in our experiments, given a sample budget, we compute the empirical distribution of Z for both algorithms under both hypothesis cases across 100 trials. For a fixed algorithm, our error measurement computes the dissimilarity of the two empirical histograms: if the histograms are 'well-separated' then we know the algorithm can adequately distinguish the two hypotheses. On the other hand, if the histograms of Z in the two cases are highly overlapping, then it is clear that the algorithm cannot distinguish well. Thus, given the two empirical histograms, we calculate the best empirical threshold which best separates them. This is done for both algorithms, ours and the classical Closeness Tester. Then we measure the fraction of data points in the histogram which are misclassified as the error. Calculating the best empirical threshold also makes both of the algorithms more practical as the theoretical thresholds are only stated asymptotically with un-optimized constant factors, making the theoretical values impractical. Finally, note that it is trivial to get 50% error.

CRS '15 is a fundamentally different algorithm. Rather than calculating an estimator, it queries the hints ˆ p and ˆ q on the samples received and check if they are consistent with each other. If all samples i satisfy ˆ p ( i ) = (1 ± ε/ 10)ˆ q ( i ) , the algorithm declares we are in the p = q case.

Note that as implemented, our algorithm and the standard Closeness Tester do not require the knowledge of ε ( ε is not needed to calculate the estimator Z ). On the other hand CRS '15 crucially

requires the knowledge of ε . Thus, we make this baseline even stronger by providing it with the exact knowledge of ε (which the other two algorithms do not have access to).

Results.

Hard Instance . Our results are shown in Figures 1 and 5. We now describe the plots. Figure 1 displays the performance of all algorithms as a function of the number of samples seen. Note that the y axis measures the fraction of times we can distinguish the two hypotheses across 100 trials. In Figure 1, our algorithm has access to ˆ p = p and CRS '15 has access to ˆ p = p . Furthermore, it has access to a hint ˆ q which satisfies ∥ ˆ q -q ∥ TV < 0 . 002 . ˆ q hat is created by removing a small fraction of the mass on a small fraction of the heavy domain elements between 1 and m .

̸

We see that the performance of both our Augmented Tester and the standard Closeness Tester improves as the sample complexity increases. On the other hand, CRS '15 essentially always obtains the trivial guarantee of 50% error even when it is fed a large number of samples. This is because with decent probability, we sample an element where ˆ q ( i ) = 0 and ˆ p ( i ) = 0 . Thus in the p = q case, it is very likely to output the incorrect answer.

Figure 1 also demonstrates that our augmented tester achieves up to > 20x reduction in error than the standard Closeness Tester. It also requires > 2x fewer samples to achieve 0% error empirically.

̸

Our improvements are due to the fact that the empirical estimators of Z in the two cases of p = q and p = q are very well-separated for the Augmented Tester. In fact, as shown in Figures 6(a) and 6(b), if we fix the number of samples to 10 4 , the empirical distributions of Z in the two cases do not overlap at all for the augmented algorithm. On the other hand, thee is a significant overlap for the un-augmented tester (the standard Closeness Tester). This implies that the Closeness Tester baseline cannot distinguish the two cases as well as the augmented algorithm.

IP Data . Our results are shown in Figure 7(a). Again we see the same qualitative behavior: CRS '15 is quite brittle to prediction errors: in the case ( p, p ) case, we feed CRS '15 with the perfect prediction p for the first distribution, and ˆ p for the second. In this case, it typically outputted the incorrect answer. Furthermore, our augmented algorithm has a better sample vs error trade-off than the un-augmented approach. Indeed, we observed that while the augmented tester only required ∼ 500 samples to obtain 0% error in 100 trials, the standard tester required > 700 samples, implying our augmented approach is > 40% more efficient.

Robustness to Prediction Error. Our experiments also demonstrate the robustness of our augmented approach to the quality of the predictions. Our theory crisply quantifies how the theoretical sample complexity is affected by a predictor which satisfies ∥ p -ˆ p ∥ TV = α , and correspondingly in practice, we observe that our algorithm is able to take advantage of predictors of varying quality.

Indeed, for the hard instance, we observe that the augmented approach still obtains better error than the standard tester, even if ∥ p -ˆ p ∥ TV is high. This is demonstrated in Figure 5 where we monotonically increase β (as described above). The plot shows that while the error also increases monotonically, it is only when ∥ p -ˆ p ∥ TV ≥ . 7 that the quality is comparable to the standard closeness tester.

Indeed, even at high values of β , ˆ p is still informative of the heavy-hitter structure of p . Therefore, the augmented algorithm can still reliably use ˆ p to perform a suitable flattening.

The phenomenon repeats itself in for the IP data. In Figure 7(b), we compare the performance of using two different predictors. ˆ p 1 is the same predictor as Figure 7(a) whereas ˆ p 2 is the IP distribution that is in between p and q (it is formed using the chunk of time that is exactly in between the first and the last). We observed that now ∥ p -ˆ p 2 ∥ TV ≈ 0 . 69 which is much higher than ∥ p -ˆ p 1 ∥ TV. Nevertheless, it is still the case that many heavy-hitters between p and ˆ p 2 are shared. Thus even though the TV distribution between p and the hint ˆ p 2 maybe high, our algorithm can still exploit the shared heavy-hitter structure to obtain improved sample complexity over the un-augmented approach. This demonstrates the practical versatility of our algorithm beyond the precise theoretical bounds.

Algorithm 4 ℓ 2 2 -norm estimation

F Background and existing results

F.1 Estimating the ℓ 2 -norm

In this section, we show that one can estimate the ℓ 2 -norm of a distribution via the number of collisions among O ( √ n ) samples. This result was implicitly known from previous work including [52]. Here, we include a formal statement and a proof for the sake of completeness.

Fact F.1. [Adapted from [52]] Suppose we have a parameter δ ∈ (0 , 1) and an arbitrary distribution p over [ n ] for which we have sample access to. Algorithm 4 receives n and δ and s = O ( √ n · log(1 /δ )) samples from p and outputs an estimate L of the the ℓ 2 2 -norm of p such that with probability 1 -δ , we have:

Proof. Suppose we have s samples drawn from p : x 1 , x 2 , . . . , x s . Let X denote the number of collisions among these samples. That is, the number of pairs of equal samples X := ∑ i<j ✶ x i = x j . In [52], they have shown that E [ X ] = ( s 2 ) · ∥ p ∥ 2 2 and Var [ X ] is at most 2 ( E [ X ]) 3 / 2 . Hence, using Chebyshev's inequality, we obtain:

The second to the last inequality holds since the ℓ 2 -norm of any discrete distribution over [ n ] is at most 1 / √ n . And, the last inequality holds for s ≥ 160 √ n. The above inequality implies that each L i in the algorithm is within the desired bound with probability at least 0 . 9 .

To find an accurate estimate with a probability 1 -δ , we use standard amplification technique: We repeat this process O (log(1 /δ )) times and take the median of these estimates. Using the Chernoff bound, one can show that median is accurate with probability at least 1 -δ .

F.2 Standard closeness tester

In this section, we present the result of [30] for robust ℓ 2 -distance estimation between two unknown distributions p and q . Their results implies an ℓ 1 -tester for closeness testing of p and q as it has been used in [42]. The important feature of this tester is its adaptive sample complexity to max ( ∥ p ∥ 2 2 , ∥ q ∥ 2 2 ) which combined with the flattening technique would achieve the optimal sample complexities for several distribution testing problems. Later [9] have shown that one can modify the this closeness tester such that the sample complexity only deepens on min ( ∥ p ∥ 2 2 , ∥ q ∥ 2 2 ) .

Fact F.2 (Adapted from [30, 42, 9]) . Suppose we have two unknown distributions p and q over [ n ] . Let b ≥ min ( ∥ p ∥ 2 2 , ∥ q ∥ 2 2 ) . Algorithm 5 can distinguish whether p = q or ϵ -far from each other with probability 1 -δ using O ( n √ b/ϵ 2 ) samples.

11: Let X i and Y i denote the number of occurrences of element i in the samples drawn from distributions p and q respectively.

Algorithm 5 ℓ 2 2 -norm estimation 1: procedure STANDARD CLOSENESS TESTER(sample access to p , n , δ ) 2: L p ← ESTIMATEℓ 2 2 ( p, n, δ/ 3) 3: L q ← ESTIMATEℓ 2 2 ( q, n, δ/ 3) 4: if L p /L q ̸∈ [1 / 3 , 3] then 5: return reject 6: t ← O (log(1 /δ ) 7: m ← O ( n · √ b ϵ 2 ) 8: accept-count ← 0 9: for i = 1 , . . . , t do 10: Draw m samples from p and q . 12: Z ← √∑ n i =1 ( X i -Y i ) 2 -X i -Y i m 13: if Z ≤ ϵ 2 √ n then 14: accept-count ← accept-count +1 15: if accept-count ≥ t/ 2 then 16: return accept 17: return reject

F.3 Useful facts

Coupling: Suppose we have two probability distributions p and q over a finite domain Ω . A coupling between p and q is a joint distribution over Ω 2 where the two marginals of C are p and q . Below is the fundamental coupling lemma:

Fact F.3. (Coupling Lemma) For every two probability distributions over a finite domain Ω , we have:

For any coupling C of p and q , we have:

̸

There exists a maximal coupling C ∗ (also known as optimal coupling) for which the above inequality holds:

̸

G Augmented flattening for guarantees other than total variation distance

Below, we provide a simple lemma that indicates how one can use ˆ p to flatten a distribution based on various guarantees of the prediction. Suppose our target ℓ 2 2 -norm which we desire is v . In the following lemma, we propose a flattening F to decrease the ℓ 2 2 to v up to some multiplicative factor that depends on the quality of the estimates.

Lemma G.1. Suppose we have an unknown distribution p = ( p 1 , p 2 , . . . , p n ) over [ n ] . Assume for every i ∈ [ n ] we are give an estimate ˆ p i . Then for every v ≤ 1 , there exists a flattening such that it increases the domain size by 1 /v and the ℓ 2 2 -norm of the p ( F ) is bounded by:

Proof. For every i ∈ [ n ] , we set m i as follows:

It is not hard to see:

In addition, the new domain size is bounded as follows:

Using ∑ n i =1 ˆ p i = 1 , we conclude the statement of the lemma.

In the following, we discuss a few cases for which we have some guarantee on the accuracy of the estimates and how it affects the ℓ 2 -norm reduction. Note that our desired case here is to not blow up the domain size by more than a constant factor and bring down the ℓ 2 -norm to O (1 / √ n ) .

Exact values: Assume ˆ p i = p i . By setting v = 1 /n , the ℓ 2 -norm reduces to Θ(1 / √ n ) , and the new domain size is Θ( n ) .

Multiplicative upper bounds: Suppose for every i , we are given an estimate ˆ p i such that ˆ p i ≥ αp i for a parameter α ≥ 1 . Then , we have ∑ n i =1 p 2 i / ˆ p i is at most 1 /α . Also, let β denote ∑ n i =1 ˆ p i . Thus, the sample complexity of the tester in [30], is:

Note that the last inequality is achieved by minimizing over v and setting v = β/n . Moreover, this bound gives us the optimal sample complexity for constant α and β .

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: All theoretical claimed are backed by full proofs.

2. Limitations

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

Answer: [Yes]

Justification: We prove tight lower bounds for our setting and discuss alternate prediction models.

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: All theorems have full, formal proofs.

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: A Jupyter notebook is included in the submission which shows has the full code.

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: We include a link to the CAIDA dataset and a Jupyter notebook is attached to the submission.

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: All experimental details are given in Section E.

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: See e.g. Figures 6(a), 6(b).

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: [NA]

Justification: We use sample complexity as a measure of efficiency, not computational time (although all of our algorithms are very efficient).

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: The paper conforms to the NeurIPS Code of Ethics.

10. Broader Impacts

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

Answer: [NA]

Justification: There is no societal impact of the work performed.

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: [NA]

Justification: The paper has no language models, image generators, or scraped datasets and poses no such risks.

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: The datasets used are fully described in Section E.

13. New Assets

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

Answer: [NA]

Justification: No new assets introduced.

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: The paper does not involve crowdsourcing nor research with human subjects.

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: No human subjects used.

Peer reviews:
comments: SUMMARY:
This paper consider the problem of hypothesis testing for discrete distributions, including identity testing, closeness testing and uniformity testing. The authors investigated these testing problems under the setting where a predicted data distribution is available. Utilizing this additional information, the authors propose "augmented tester", which can reduce the sample complexity when the given predicted distribution is accurate (when it is inaccurate, the proposed algorithm is still robust). Lower bounds are also provided for justifying the optimality of the proposed algorithm.

STRENGTHS:
The idea of augmenting hypothesis testing using additional information is interesting. The proposed algorithm can reduce the sample complexity when the given prediction is accurate enough, which is satisfactory. Lower bounds are also provided, making this paper comprehensive.

WEAKNESSES:
A possible weakness is that, the structure of this paper is a little bit weird. You first give your theorems in the introduction when your proposed algorithm is not mentioned yet. Then you provide the proofs in section 2, and state the upper bound again in section 3. The algorithm has not been mentioned until the last.

QUESTIONS:
I am curious about the intuition behind your algorithm, i.e., how to utilize the additional information.  It will be good if the authors could provide more explanation of the proposed algorithm.

LIMITATIONS:
no limitations are stated.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the problem of property testing (specifically uniformity, identity, and closeness testing) in the context of the learning-augmented algorithms framework. They show how a prediction about the underlying distribution could be harnessed to provably reduce the number of samples required when the prediction is of high quality. They give efficient algorithms and provide experimental evaluation with code.\n\nSTRENGTHS:\nThe problem is well-motivated and the paper is well written.\n\nWEAKNESSES:\nI did not check the proofs in the appendix in detail but the intuition given for the approaches make sense.\n\nQUESTIONS:\n- In the paragraph starting on Line 177: given \"perfect predictors\", why does one even need additional samples from the distribution? From my limited understanding, that paper talks about having additional access to other kinds of queries on the underlying distribution (apart from just drawing IID samples) and \"samples\" may refer to those other kinds of queries (for instance, one can query the probability mass of an element and get its true mass in $p$). I'm not sure how that directly compares with your setting where \"samples\" mean IID samples only.\n- How do you obtain an **upper bound** expression for the final big-O containment for the equation block between Lines 378-379? To be precise, we have $s_f \\leq = \\min\\{\\frac{n^{2/3} \\alpha^{1/3}}{\\varepsilon^{4/3}}, n\\} \\leq \\frac{n^{2/3} \\alpha^{1/3}}{\\varepsilon^{4/3}}$ so $\\frac{n}{\\varepsilon^2} \\cdot \\sqrt{\\frac{\\alpha}{s_f}} \\geq \\frac{n^{2/3} \\alpha^{1/3}}{\\varepsilon^{4/3}}$, which is a **lower bound** right? Where did I mess up?\n\nPossible typos:\n- Line 266: \"...strategy is put the...\" should be \"...strategy is **to** put the...\"?\n- First line of the equation between Lines 328-329: Should the last equality be $\\leq$?\n- Line 375: Should it be $\\leq 100 s_f$ instead?\n- Line 376: Should it be $\\leq 102 n$ instead?\n\nLIMITATIONS:\nNil", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: SUMMARY:
The paper looks at the problem of hypothesis testing for discrete distributions where a predicted data distribution is available. The paper gives algorithms (Algorithm 1; closeness testing, Algorithm 3; identity and uniform testing) which either reduces the number of samples required for testing, or do no worse than standard approaches. Lower bounds on samples are given, and experiments are done to validate the performance on real data.

STRENGTHS:
The paper is extremely thorough and detailed, giving an overview of the important results (including experiments), as well as a roadmap of the intuition required behind the proofs for the upper bound and lower bounds.

The appendix also includes motivation and explanation for the numerous proofs involved.

The main contributions of needing significantly fewer sample sizes for hypothesis testing compared to CRS'15 or [41] is also significant.

WEAKNESSES:
I am not sure whether this is an appropriate venue mostly due to the heavy technical details in the paper, and with a lot of key details in the appendix ; but to be fair, the conference format makes it difficult to present detailed work like this.

Unfortunately, this is not mostly in my area of speciality, so I am unable to give meaningful comments about potential weaknesses (or strengths).

QUESTIONS:
These questions may be straightforward to those familiar with the relevant literature, but these are the questions I have:
 - Are there certain types of data where making weak assumptions about the predictor (compared to related works) might give bad results (or similar results to existing methods)?
 - Is there a reason why experiments for identity testing and uniformity testing are not given?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the task of identity/closeness testing when the tester is augmented with predictions of the unknown distribution involved apriori. In particular, in identity testing, in addition to sample access to the unknown distribution $p$, the algorithm is also given some $\\hat p$ that may or may not satisfy the guarantee $ \\text{TV}(p, \\hat p) \\leq \\alpha$. If the guarantee does not hold, the algorithm is allowed to answer **inaccurate**. Otherwise, the algorithm is asked to perform the standard testing task: return YES if $p$ equals to some known distribution $q$, and NO if $p$ is at least $\\epsilon$-far from $q$ in total variation distance. For closeness testing, though both $p$ and $q$ are unknown, the authors assume that such prior prediction $\\hat p$ is only available for one of the distributions. \n\nFor identity testing, the sample complexity goes through a phase transition depending on the relative size of the prediction error tolerance parameter $\\alpha$ and $d = \\text{TV}(q, \\hat p )$, the TV distance between the prediction distribution and the known distribution. When $d<\\alpha$, the sample complexity is the same as the usual identity testing. When $d > \\alpha$, the sample complexity is $\\min\\left( \\frac{1}{ (d - \\alpha)^2 }, \\frac{\\sqrt{n}}{\\epsilon^2} \\right)$. Intuitively, this is saying that the prediction could potentially help a lot if it predicts that the test is in the soundness case (as $\\hat p$ is indeed far from $q$) but not so much if it predicts the test is in the completeness case. This intuition is clear from their upper bound approach. They leverages the Scheffé set $S$ (the set that realizes the maximum discrepancy between $\\hat p$ and $q$), and test for discrepancy between $p(S)$, $q(S)$ and $\\hat p(S)$. This will either invalidate the prediction or lead to rejection. Furthermore, since this is simply a 1-dimensional bias estimation problem, they could avoid any dependency on the doamin size (when $d$ is sufficiently separated from $\\alpha$). The above strategy is not so helpful when there is no significant discrepancy between $p(S)$, $q(S)$ and $\\hat p(S)$, and therefore their algorithm fall back to the standard identity tester when the prediction says that the test is in the completeness case.\n\nFor closeness testing, they show that the sample complexity is given by $\\sqrt{n} \\alpha^{1/3} / \\epsilon^{4/3} + \\sqrt{n} / \\epsilon^2$. Here the Scheffé set strategy no longer applies as $q$ is also unknown to us. However, the authors show that the prior $\\hat p$ is still very useful in an important closeness testing sub-routine (commonly referred as **flattening**). At a high level, when the unknown distributions have large $\\ell_2$ norm, standard collision-based test statistics may have large variance. Flattening techniques could then be used to transform the distributions to reduce their $\\ell_2$ norms, and hence also the variance of the test statistics. The authors show that such a routine can be implemented in a more sample-efficient way if the algorithm is equipped with a prediction $\\hat p$.\n\nSTRENGTHS:\nIdentity/closeness testings are fundamental problems in many different areas. In the standard setting, the sample complexity scales polynomially with respect to the domain size, making the tests prohibitively expensive for distributions with large supports. The authors demonstrated that the sample complexity may be significantly improved if the tester is augmented with a prediction of the unknown distribution. The surprising part is that the prediction need not to be an accurate one. While the tester draws significantly fewer samples if the prediction is accurate, even if the prediction is not, the tester is guaranteed not to be misguided. In particular, it can identify inaccuracies in the prediction, and simply fall back to the standard testing approach in that case. Lastly, the bounds are optimal up to constant factors and technically solid.\n\nWEAKNESSES:\nThe authors mention that the potential application is when the unknown distribution evolves over time. In that case, the ``unknown'' distribution is not completely unknown to us as we may extract information from past data. It will be more compelling if the authors could provide more detailed mathematical setup of this. For example, it will be interesting to see analysis of the behavior of augmented testers when we face a sequence of tests where the unknown distribution may go through random distribution shifts from one task to another.\n\nQUESTIONS:\nHave the authors considered situation when we have prior knowledge about both of the unknown distributions in closeness testing? Then it seems like the Scheffé set may again become useful, and could lead to huge sample complexity improvements.\n\nComment: The setup bears some similarities to the notion of testable learning, where the algorithm is given some prior distribution assumption that may or may not hold. In particular, the algorithm is also given the ability to output ``reject'' when this prior knowledge is false.\n\nLIMITATIONS:\nYes, they have.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}