Optimal Algorithms for Augmented Testing of Discrete Distributions

Authors: maryam  aliakbarpour
piotr  indyk
ronitt  rubinfeld
sandeep  silwal
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 07:37:58

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.

Main paper: 
 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] .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 p1 . 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 pi 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.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.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.Theorem 2 (Informal version of Theorem 7) . Augmented uniformity and identity testing for distributions over [n], with parameters α, ϵ, and δ = 2/3, require the following number of samples:s =      Θ √ n ϵ 2 if d ≤ α Θ min 1 (d-α) 2 , √ n ϵ 2 if d > α, 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 .Table 1 : Optimal sample complexity bounds in the standard model versus the augmented bounds. α is the suggested accuracy level for the L 1 -distance between p and p. d denotes the total variation distance between the known distribution q and p.Standard sample complexity Augmented sample complexity (this paper)Closeness Θ n 2/3 ϵ 4/3 + √ n ϵ 2 [30, 69] Θ n 2/3 α 1/3 ϵ 4/3 + √ n ϵ 2 Identity/Uniformity Θ √ n ϵ 2 Θ √ n ϵ 2 , if d ≤ α [69, 5] Θ min 1 (d-α) 2 , √ n ϵ 2 , if d > α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 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).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 pi = (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 pi 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.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 closeness testing is discussed in Section D, with the lower bounds provided in Section D.1. Our empirical evaluations are presented in Section E. 
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. 
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.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:E p (F ) 2 2 ≤ 1 t ,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.). For any flattening scheme F , the ℓ 1 -distance is preserved under F . That is, for every pair of distributions p and q, we have:∥p -q∥ 1 = ∥p (F ) -q (F ) ∥ 1 .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.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:E p (F ) 2 2 ≤ 2α s f + 4 • ν .The expectation above is taken over the randomness of the samples from p.Proof. Let ŝf := Poi(s f ) be a Poisson random variable with mean s f . Let S f denote the multiset of ŝ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:EQUATIONLet 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 -pi . Suppose A is a set of indices i ∈ [n] for which ∆ i ≥ pi . 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 pi . p (F ) 2 2 = i∈[n] j∈[mi] p (F ) (i, j) 2 = i∈[n] p 2 i m i = i∈A 2 ∆ i • p i m i + i∈[n]\A (2 pi ) 2 m i ≤ i∈A 2 ∆ i • p i f i + 1 + i∈[n]\A (2 pi ) 2 pi /ν ≤ i∈A 2 ∆ i • p i f i + 1 + 4 ν.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 :E F p i f i + 1 ≤ 1 s fTherefore, we obtain:E F p (F ) 2 2 ≤ i∈A 2 ∆ i s f + 4 ν ≤ 2 α s f + 4 ν .The last inequality above is due to the fact that:i∈A ∆ i = p(A) -p(A) ≤ ∥p -p∥ TV = α .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.Algorithm 1 An augmented tester for testing closeness of p and q with a suggested value α 1: procedure AUGMENTED-CLOSENESS-TESTER(n, α, ϵ, δ 1 = 0.1, δ 2 = 2/3, p, sample access to p and q)2:s f ← min n 2/3 α 1/3 /ϵ 4/3 , n 3: ŝf ← Poi(s f ) 4:if ŝf > 10 s f then 5:return reject.6:F ← F (m 1 , m 2 , . . . , m n )where m i are define in Equation 2.7:L p ← ESTIMATE-ℓ 2 2 p (F ) , n i=1 m i , 0.05 8: L q ← ESTIMATE-ℓ 2 2 q (F ) , n i=1 m i , 0.05 9: if L p > 30 2α s f + 4 n then 10:return inaccurate information.11:else 12: if L q > 90 2α s f + 4 n then 13:return reject.14:else 15: return CLOSENESS-TESTER ϵ, δ = 0.1, b = 90 2α s f + 4 n , p (F ) , q (F ) .Theorem 5. For any given parameters n, α, ϵ, and any two unknown distributions p and q and a known predicted distribution p for p, Algorithm 1 is an (α, ϵ, δ = 0.3)-augmented tester for testing closeness of p and q which uses O(n 2/3 α 1/3 /ϵ 4/3 ) samples.Proof. We begin the proof by setting our parameters. Set s f := min n 2/3 α 1/3 /ϵ 4/3 , n . Suppose we draw ŝ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):EQUATIONIn the above definition, f i is the frequency of element i among the ŝ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:EQUATION)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:E p (F ) 2 ≤ 2α s f + 4 n .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 < 3L 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.For the flattening step we have used ŝf ≤ 10s f . After flattening, the new domain size is bounded from above by:n i=1 m i ≤ n i=1 n • pi + f i + 1 = n + ŝf + n ≤ 12n. 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:# samples = O s f + √ n + n • b ϵ 2 = O s f + √ n + n • 2α/s f + 4/n ϵ 2 = O √ n ϵ 2 + n 2/3 α 1/3 ϵ 4/3 .In the last line we use that s f := min n 2/3 α 1/3 /ϵ 4/3 , n . Thus, the proof is complete.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.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 α1: procedure AUGMENTED-TESTER-WITH-SEARCH(ϵ, δ ′ , description of function f, algorithm A) 2:s min ← min α∈[0,1] f (α, 1/3) 3:s max ← max α∈[0,1] f (α, 1/3) 4: t ← log 2 smax smin 5: δ ← δ ′ t+1 6:for i = 0, 1, 2, . . . , t -1 do 7:α i ← f -1 2 i • s min 8:O ← the output of A with suggested accuracy level α i and confidence parameter δ 9:if O ∈ {accept, reject} then 10:return O.Run the standard tester (without any prediction) with confidence parameter δ and return the answer. 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δ ′ := δ • 1 + log max α∈[0,1] f (α) min α∈[0,1] f (α)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 . IfA 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 have4 :2 i * -1 • s min < f (α * ) ≤ f (α i * ) ≤ 2 i * • s minIf the algorithm ends at round i, we have used the following number of samples:# samples for the first i rounds ≤ i j=0EQUATIONOn 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 * :EQUATIONNow, 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:E[ [ [ S ] ] ] = t i=0 E[ [ [ S | I end = i] ] ] • Pr[ [ [ I end = i] ] ] ≤ 4 f (α * ) • t i=0 2 i-i * • Pr[ [ [ I end = i] ] ](by Eq. ( 4))≤ 4 f (α * ) •   i * i=0 Pr[ [ [ I end = i] ] ] + t i=i * +1   1 + i-i * -1 j=0 2 j   • Pr[ [ [ I end = i] ] ]   = 4 f (α * ) •   t i=0 Pr[ [ [ I end = i] ] ] + t i=i * +1 i-i * -1 j=0 2 j • Pr[ [ [ I end = i] ] ]  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:E[ [ [ S ] ] ] ≤ 4 f (α * ) •   1 + t-i * -1 k=0 k j=0 2 j • Pr[ [ [ I end = k + i * + 1] ] ]   ≤ 4 f (α * ) •   1 + t-i * -1 j=0 2 j •   t-i * -1 k=j Pr[ [ [ I end = k + i * + 1] ] ]     ≤ 4 f (α * ) •   1 + t-i * -1 j=0 2 j • Pr[ [ [ I end ≥ j + i * + 1] ] ]   ≤ 4 f (α * ) •   1 + t-i * -1 j=0 2 j • δ 2 j+1  (by Eq. ( 5))≤ 6f (α * ) . (assuming δ < 1/2)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:s =      O √ n ϵ 2 if d ≤ α O min 1 (d-α) 2 , √ n ϵ 2 if d > α ,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.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:d := ∥q -p∥ TV = |q(S) -p(S)| .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:s =      O √ n ϵ 2 if d ≤ α O min 1 (d-α) 2 , √ n ϵ 2 if d > α ,and d refers to ∥q -p∥ TV .Algorithm 3 An augmented tester for testing identity of p and q with a suggested value α 1: procedure AUGMENTED-IDENTITY-TESTER(n, α, ϵ, δ = 0.1, q, p, sample access to p)2: d ← ∥q -p∥ TV 3: if d ≤ α or 1 (d-α) 2 > √ n ϵ 2 then 4:return the standard testers answer.S ← Scheffé set of p and q ▷ x ∈ [n] is in the Scheffé set S iff p(x) < q(x).6:Draw m = O 1 (d-α) 2 samples from p 7:σ ← fraction of samples that are in S 8:if |U n (S) -σ| > d-α 4 then 9:return reject. return inaccurate information.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:EQUATIONNote that if p = q, then|σ -q(S)| = |σ -p(S)| ≤ d -α 4 .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))Therefore, we obtain that with probability 0.95:|σ -q(S)| ≥ 3 (d -α) 4 .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 requiress = Ω 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]:pi :=        1 + 2d n i is even; 1 -2d n i is odd.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]:EQUATION)EQUATIONIt 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 |Un , 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. , 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 |Un , 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 (S1,S2,S3)∼C [ [ [ S 1 = S ] ] ] = Pr S |Un ∼U ⊗s n S |Un = S , Pr (S1,S2,S3)∼C [ [ [ S 2 = S ] ] ] = Pr S |p • ∼p • ⊗s S |p • = S , Pr (S1,S2,S3)∼C [ [ [ S 3 = S ] ] ] = Pr S |p ⋄ ∼p ⋄ ⊗s S |p ⋄ = S . • Pr (S1,S2,S3)∼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 ⋄ . 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 -δ.1. p = U n p = U n p = U n :In this case we have p = q = U n . Clearly, in this case, reject is not a valid answer:reject ̸ ∈ V |Un . 2. 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:accept ̸ ∈ V |p • .3. p = p ⋄ p = p ⋄ p = p ⋄ : It is not hard to see that ∥p ⋄ -p∥ TV = α. Thus, inaccurate information is not a valid answer:inaccurate information ̸ ∈ V |p ⋄ .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 :V |Un ∩ V |p • ∩ V |p ⋄ = ∅ .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.3 • (1 -δ) ≤ Pr S |Un ∼Un A(S |Un ) ∈ V |Un + Pr S |p • ∼p • A(S |p • ) ∈ V |p • + Pr S |p ⋄ ∼p ⋄ A(S |p ⋄ ) ∈ V |p ⋄Let I represent the internal coin tosses of the algorithm. Since the probability of A outputting an invalid answer is bounded by δ, we have:δ ≥ Pr I, S |Un ∼U ⊗s n A(S |Un ) ̸ ∈ V |Un = E S |Un ∼U ⊗s n Pr I A(S |Un ) ̸ ∈ V |Un = E S |Un ∼U ⊗s n Pr I A(S |Un ) = reject = E (S1,S2,S3)∼C [ [ [ Pr I [ [ [ A(S 1 ) = reject] ] ]] ] ] ≥ E (S1,S2,S3)∼C [ [ [ Pr I [ [ [ A(S 1 ) = reject] ] ] | S 1 = S 2 = S 3 ] ] ] • Pr (S1,S2,S3)∼C [ [ [ S 1 = S 2 = S 3 ] ] ] .(by law of total expectation)We can write the above equation for p • and p ⋄ and sum of them up. Thus, we obtain:δ ≥ 1 3 • Pr (S1,S2,S3)∼C [ [ [ S 1 = S 2 = S 3 ] ] ] • E (S1,S2,S3)∼C [ [ [ Pr I [ [ [ A(S 1 ) = reject] ] ] + Pr I [ [ [ A(S 2 ) = accept] ] ] + Pr I [ [ [ A(S 3 ) = inaccurate information] ] ] | S 1 = S 2 = S 3 ] ] ] = 1 3 • Pr (S1,S2,S3)∼C [ [ [ S 1 = S 2 = S 3 ] ] ] .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:0.3 = δ ≥ 1 3 • Pr (S1,S2,S3)∼C [ [ [ S 1 = S 2 = S 3 ] ] ] ≥ 0.98 3 > 0.32 .Hence, we reach a contradiction. Thus, no such algorithm like A exists.In this section, we prove the lemma on coupling that we used to prove our lower bound. This lemma implies indistinguishability between S |Un , S |p , 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 |Un , 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 (S1,S2,S3)∼C [ [ [ S 1 = S ] ] ] = Pr S |Un ∼U ⊗s n S |Un = S , Pr (S1,S2,S3)∼C [ [ [ S 2 = S ] ] ] = Pr S |p • ∼p • ⊗s S |p • = S , Pr (S1,S2,S3)∼C [ [ [ S 3 = S ] ] ] = Pr S |p ⋄ ∼p ⋄ ⊗s S |p ⋄ = S . • Pr (S1,S2,S3)∼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 np •⊗s ∥ TV , and ∥U ⊗s n -p ⋄⊗s ∥ TV . In [69] , the author has shown that the total variation distance between S |Un and S |p • is bounded by:EQUATION)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 |Un . 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 :∥U ⊗s n -p ⋄⊗s ∥ TV ≤ 1 2• KL U ⊗s n ∥p ⋄⊗s (by Pinsker's inequality)≤ s 2 • KL (U n ∥p ⋄ ) (since samples are drawn i.i.d.) ≤ s 2 • KL (Ber(1/2)∥Ber(1/2 + (d -α))) .(by data processing inequality)≤ √ 2 s • (d -α) (when 2 (d -α) ≤ 0.8)Since we had this assumption where s is at most 0.00005 (d-α) 2 , we get.EQUATIONSo 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:EQUATIONPr (S1,S2)∼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: C ⋄ 2|S (S ′ ) := Pr (S1,S2)∼C ⋄ [ [ [ S 2 = S ′ |S 1 = S ] ] ]. 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:1. Let (S 1 , S 2 ) be a sample from C • . 2. Let S 3 be a sample from C ⋄ 2|S1 .3. 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 everyS ′ ∈ [n] m : Pr[ [ [ S 3 = S ′ ] ] ] = S1∈[n] m Pr (S ′ 1 ,S ′ 2 )∼C • [ [ [ S ′ 1 = S 1 ] ] ] • C ⋄ 2|S1 (S ′ ) = S1∈[n] m Pr (S ′ 1 ,S ′ 2 )∼C • [ [ [ S ′ 1 = S 1 ] ] ] • Pr (S ′′ 1 ,S ′′ 2 )∼C ⋄ [ [ [ S ′′ 2 = S ′ |S ′′ 1 = S 1 ] ] ] = S1∈[n] m Pr (S ′ 1 ,S ′ 2 )∼C ⋄ [ [ [ S ′ 1 = S 1 ] ] ] • Pr (S ′′ 1 ,S ′′ 2 )∼C ⋄ [ [ [ S ′′ 2 = S ′ |S ′′ 1 = S 1 ] ] ] = S1∈[n] m C ⋄ (S 1 , S ′ ) = C ⋄ 2 = p ⋄ (S ′ ) .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:Pr (S1,S2,S3)∼C [ [ [ S 1 , S 2 , S 3 ] ] ] ≥ 1 -Pr (S1,S2,S3)∼C [ [ [ S 1 ̸ = S 2 ] ] ] -Pr (S1,S2,S3)∼C [ [ [ S 1 ̸ = S 3 ] ] ] = Pr (S1,S2,S3)∼C [ [ [ S 1 = S 3 ] ] ] -Pr (S1,S2)∼C • [ [ [ S 1 ̸ = S 2 ] ] ] = S∈[n] m Pr (S1,S2,S3)∼C [ [ [ S 1 = S and S 3 = S ] ] ] -Pr (S1,S2)∼C • [ [ [ S 1 ̸ = S 2 ] ] ] = S∈[n] m C ⋄ (S, S) -Pr (S1,S2)∼C • [ [ [ S 1 ̸ = S 2 ] ] ](With a similar argument as above)= Pr (S1,S2)∼C ⋄ [ [ [ S 1 = S 2 ] ] ] -Pr (S1,S2)∼C • [ [ [ S 1 ̸ = S 2 ] ] ] ≥ 0.99 -(1 -0.99) = 0.98 .(by Eq. 9 and Eq. 10)Hence, the proof is complete.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 ϵ 2samples 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.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 ϵ 2samples. 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.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:s = Ω n 2/3 α 1/3 ϵ 4/3 + √ n α ϵ 2 .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:p i = α • pi ∀i ∈ [n -1] , p n = 1 -α , q i = α • qi ∀i ∈ [n -1] , q n = 1 -α .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:∥p -q∥ TV = 1 2 i∈[n-1] |p i -q i | = α 2 i∈[n-1] |p i -qi | = α • ∥p -q∥ TV .Let p be the singleton distribution on n (i.e., pn = 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:100 α s = Ω n 2/3 ϵ ′4/3 + √ n ϵ ′2 = Ω n 2/3 α 4/3 ϵ 4/3 + √ n α 2 ϵ 2 .Thus, by dividing by 100α, we get the following bound as desired:s = Ω n 2/3 α 1/3 ϵ 4/3 + √ n α ϵ 2 . 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Ω n 2/3 α 1/3 ϵ 4/3samples.Proof. Let k be an integer that is at mostEQUATIONfor 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:EQUATIONIt 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:n 2/3 α 4/3 2 ϵ 4/3 > n 2/3 2 ϵ 4/3 • 1 √ n ϵ 2 4/3 = 1 2 ϵ 4 ≥ 6 4 2 .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:EQUATION)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. SetEQUATION. Constructing the distributions: We set p to be a uniform distribution over [n] . That is, pi = 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:p = U p + p - L L 1/n λsmall λlarge λlarge (1 + ϵ ′ ) • λsmall (1 -ϵ ′ ) • λsmallp + i :=      1 n + α ℓ i ∈ L 1 n - α n -ℓ i ∈ L p - i :=                  1 n + α ℓ i ∈ L (1 + ϵ ′ ) • 1 n - α n -ℓ i ∈ L 1 (1 -ϵ ′ ) • 1 n - α n -ℓ i ∈ L 2For 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:1 n - α n -ℓ ≥ 1 n - α n/2 > 0.Also, given these definitions, p + is α-close to p, and p + and p -are ϵ-far from each other:∥p + -p∥ TV = 1 2   i∈L α ℓ + i∈L α n -ℓ   = α , ∥p + -p -∥ TV = 1 2 i∈L ϵ ′ • 1 n - α n -ℓ = ϵ ′ 2 • n -ℓ n -α > ϵ ′ 2 • 1 2 - 1 6 ≥ ϵ .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:M (k) a,b (p, q) := k a+b i∈[n]p a i q b i 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 a+b≥2 M (k)a,b p (1) , q (1) -Ma,b p (2) , q (2)a 2 ! b 2 ! 1 + max M (k)a,b p (1) , q (1) , Ma,b p (2) , q (2)EQUATIONthen 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 , ϵ :k n = c k n 2/3 α 1/3 n ϵ 4/3 = c k • α ( √ nϵ 2 ) 2 1/3 < c k • α ϵ 2 1/3 < c k • ϵ < c k ≤ 1 2000And, the second inequality comes from the definitions of ℓ, k, and c k :α ℓ = α ϵ 4/3 c ℓ n 2/3 α 4/3 = ϵ 4/3 c ℓ n 2/3 α 1/3 = c k c ℓ k ≤ 1 2000 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:EQUATIONUsing the Definition D.1, it is not hard to see that:M (k) a,b p + , p + = k a+b   i∈L λ large a+b + i∈L λ small a+b   = k a+b ℓ • λ large a+b + (n -ℓ) • λ small a+b M (k) a,b p + , p -= k a+b   i∈L λ large a+b + i∈L1 (1 + ϵ ′ ) b • λ small a+b + i∈L2 (1 -ϵ ′ ) b • λ small a+b   = k a+b ℓ • λ large a+b + (n -ℓ) • λ small a+b • (1 + ϵ ′ ) b + (1 -ϵ ′ ) b 2 .Therefore, we get:M (k) a,b p + , p + -M (k) a,b p + , p -= (n -ℓ) • λ small a+b • (1 + ϵ ′ ) b + (1 -ϵ ′ ) b 2 -1 .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:1 + max M (k)a,b p (1) , q (1) , Ma,b p (2) , q (2) ≥ k a+b • ℓ • λ large a+b .Using the two equations above, we can bound the right hand side of Equation ( 14):RHS of Eq. ( 14)≤ n -ℓ √ ℓ a+b≥2 k (a+b)/2 • λ small a+b a 2 ! b 2 ! λ large (a+b)/2 • (1 + ϵ ′ ) b + (1 -ϵ ′ ) b 2 -1Note 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: RHS of Eq. ( 14)≤ n -ℓ √ ℓ •   a≥0 T a a 2 !   •   1 2 b≥2 (T • (1 + ϵ ′ )) b b 2 ! + 1 2 b≥2 (T • (1 -ϵ ′ )) b b 2 ! - b≥2 (T ) b b 2 !  To bound these terms, we have the following lemma:Lemma D.3. For x ∈ R, we have ∞ i=0 x i i 2 ! = (1 + x)e x 2 ,and for x ∈ [0, 1.33], we have:x 2 + x 3 ≤ ∞ i=2 x i i 2 ! < x 2 + x 3 + 3 x 4 .Proof. We start by splitting the terms corresponding to odd and even i's in the sum, and obtain the following:∞ i=0 x i i 2 ! = ∞ j=0 x 2j j! + x 2j+1 j! = (1 + x) ∞ j=0 x 2j j! = (1 + x) • e x 2 .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:∞ i=2 x i i 2 ! = ∞ i=0 x i i 2 ! -x -1 = (1 + x) • e x 2 -1 .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:x 2 +x 3 = (1+x)•x 2 ≤ (1+x)• e x 2 -1 < (1+x)•(x 2 +x 4 ) = x 2 +x 3 +(1+x) x 4 < x 2 +x 3 +3 x 4 .Now, assume T • (1 + ϵ ′ ) ≤ 1.33. We continue bounding the right hand side of Equation ( 14) via the above fact: RHS of Eq. ( 14)≤ n -ℓ √ ℓ •   a≥0 T a a 2 !   •   1 2 b≥2 (T • (1 + ϵ ′ )) b b 2 ! + 1 2 b≥2 (T • (1 -ϵ ′ )) b b 2 ! - b≥2 (T ) b b 2 !   < n √ ℓ • (1 + T ) • e T 2 • 1 2 T 2 • (1 + ϵ ′ ) 2 + T 3 • (1 + ϵ ′ ) 3 + 3 T 4 • (1 + ϵ ′ ) 4 + 1 2 T 2 • (1 -ϵ ′ ) 2 + T 3 • (1 -ϵ ′ ) 3 + 3 T 4 • (1 -ϵ ′ ) 4 -T 2 -T 3 = n √ ℓ • 14 • (1 + ϵ ′ ) 2 + (1 -ϵ ′ ) 2 2 -1 • T 2 + (1 + ϵ ′ ) 3 + (1 -ϵ ′ ) 3 2 -1 • T 3 + 3 (1 + ϵ ′ ) 4 + 3 (1 -ϵ ′ ) 4 2 • T 4 ≤ n √ ℓ • 14 • ϵ ′2 T 2 + 3ϵ ′2 • T 3 + 24 T 4 . (using ϵ ′ ≤ 1)We can bound T from above given the trivial bounds we have that λ small ≤ 1/n and λ large ≥ α/ℓ by their definitions in Equation ( 15):EQUATIONThe 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: RHS of Eq. ( 14EQUATION))≤ 14 • 28 • 6 2 • n √ ℓ • ϵ 2 • k • ℓ α • 1n 2 (since ϵ ′ := 6 ϵ and T ≤ k • ℓ/α/n in middle of Eq. ( 16))= 14112 • ϵ 2 k α n • √ ℓ ≤ (14112 • c ℓ ) • ϵ 4/3 k n 2/3 α 1/3 (using ℓ := c ℓ n 2/3 α 4/3 /ϵ 4/3 ) ≤ 14112 • c ℓ • c k ≤ 1 360 .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.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. Figure 6 : Additional figures for the 'Hard Instance' dataset also has probability mass 2/n on domain elements between n/2 and 3n/4 and q has probability mass 2/n on domain elements between 3n/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. We consider 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 170ms 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] . 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∥ T V < 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. p1 is the same predictor as Figure 7 (a) whereas p2 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 -p2 ∥ TV ≈ 0.69 which is much higher than ∥p -p1 ∥ TV . Nevertheless, it is still the case that many heavy-hitters between p and p2 are shared. Thus even though the TV distribution between p and the hint p2 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. for i = 1, . . . , O(log(1/δ)) do 3:s ← Θ ( √ n • log(1/δ))4:x 1 , x 2 , . . . , x s ← Draw s samples from p. L i ← X/ s 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:∥p∥ 2 2 ≤ L ≤ 3 ∥p∥ 2 2 .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 1 xi=xj .In [52] , they have shown that 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.E[ [ [ X ] ] ] = s 2 • ∥p∥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 -δ.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∥ Fact F.2 (Adapted from [30, 42, 9] ). Suppose we have two unknown distributions p and q over [n]. 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).1. Exact values: Assume pi = p i . By setting v = 1/n, the ℓ 2 -norm reduces to Θ(1/ √ n), and the new domain size is Θ(n). 2. Multiplicative upper bounds: Suppose for every i, we are given an estimate pi such that pi ≥ α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 pi . Thus, the sample complexity of the tester in [30] , is:# samples = O new domain size • min ∥p (F ) ∥ 2 , ∥q (F ) ∥ 2 ϵ 2 = O (n + β/v) • v/α ϵ 2 = O n v/α + β/ √ αv ϵ 2 = O n β/α ϵ 2Note 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 β.38th Conference on Neural Information Processing Systems (NeurIPS 2024).We assume we know all of the probability values of the prediction p at all domain elements.If p = q, the standard tester must output accept with high probability. And, if p is ϵ-far from q, the standard tester must output reject with high probability. If p ̸ = q, but is ϵ-close to q, either answer is considered correct.Here, we omit the dependence to other non-varying parameters such as ϵ and n.Without loss of generality assume, f (α-1) = 0.As long as the problem does not have a pre-determined answer that works for all distributions, one must draw at least one sample.It is worth noting that a pair of valid sets may have some overlap. That is, the algorithm can output an answer that is correct for both of the distributions. In other words, technically the algorithm does not have to distinguish any pair of these distributions as long as there is a valid answer which works for both of them.Our parameters in this discussion are not precise. See our proof for a rigorous setting of parameters.A symmetric property of a pair of distribution is a property that does not change if we permute the domain elements of two distribution via the same permutation. The closeness testing of p and q in our context is aFrom CAIDA internet traces 2019, see https://www.caida.org/catalog/datasets/monitors/ passive-equinix-nyc 
 
 Appendix:  


Decision: This paper studies the problem of identity/closeness testing when the tester is augmented with predictions of the unknown distribution involved apriori. Under this setting, the paper presents sample complexity bounds for uniformity, identity, and closeness testing problems and shows them to be smaller than that of their classical counterparts. Moreover, lower bounds are presented to establish the optimality. All the reviewers are positive about the paper's novelty and contributions, which I agree with. Based on the discussion with the reviewers, I think the paper is a nice contribution to the property testing literature and therefore I am recommending acceptance. I strongly urge the authors to incorporate all the reviewer comments along with any clarifications they provided in the rebuttal into the camera-ready version.

Reviews: review: paper_summary: paper_summary: This 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. 

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

For 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$.
strength_weakness: strength_weakness: The 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.
questions: questions: Have 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.

Comment: 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.
limitations: limitations: Yes, they have.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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: limitations: N/A
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: I did not check the proofs in the appendix in detail but the intuition given for the approaches make sense.
questions: questions: - 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.
- 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?

Possible typos:
- Line 266: "...strategy is put the..." should be "...strategy is **to** put the..."?
- First line of the equation between Lines 328-329: Should the last equality be $\leq$?
- Line 375: Should it be $\leq 100 s_f$ instead?
- Line 376: Should it be $\leq 102 n$ instead?
limitations: limitations: Nil
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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: limitations: no limitations are stated.

Comments: comment: Thanks for your explanation. I will raise my rating to 6.
Thank you for your comment. In our future version, we will add the necessary clarification to provide a more comprehensive comparison with [29].
Thank you very much for clearing up my confusion about the sample complexity!

I did read Section 1.3 and I understand that [29] considers a different query model where one can perform stronger queries such as asking for element probabilities $p(i)$. However, from what I understand and what you mentioned, [29] does not get any additional predictor as part of the input. As such, I think it is highly confusing to say that [29] uses "perfect predictors" and I do not think they are directly comparable... In the learning-augmented setup, one can always make a decision by just looking the given predictor without doing any additional work: if the predictor is arbitrarily bad, the answer will just be wrong. Thus, I find it strange to say that an algorithm given perfect predictor even needs any to take any sample.

This is still a good piece of work overall and I stick to my positive review. However, as per my response above, I would recommend the authors to reconsider the rephrasing the comparison of [29] in their revision.
Thank you for the detailed rebuttal. My score remains unchanged.