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$.
Strengths
Identity/closeness testings are fundamental problems in many different areas. In the standard setting, the sample complexity scales polynomially with respect to the domain size, making the tests prohibitively expensive for distributions with large supports. The authors demonstrated that the sample complexity may be significantly improved if the tester is augmented with a prediction of the unknown distribution. The surprising part is that the prediction need not to be an accurate one. While the tester draws significantly fewer samples if the prediction is accurate, even if the prediction is not, the tester is guaranteed not to be misguided. In particular, it can identify inaccuracies in the prediction, and simply fall back to the standard testing approach in that case. Lastly, the bounds are optimal up to constant factors and technically solid.
Weaknesses
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
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.