# Optimal Algorithms for Augmented Testing of Discrete Distributions

arXiv: https://arxiv.org/abs/2412.00974 · HTML: https://mlsanity.com/paper/2412.00974/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 47 of its 79 references that are in the map.
- First posted on arXiv 1 Dec 2024.

## Builds on (most cited first, 30 of 47)

- [Robust Estimators in High Dimensions without the Computational Intractability](https://mlsanity.com/paper/1604.06443/) (2016)
- [Learning-Based Low-Rank Approximations](https://mlsanity.com/paper/1910.13984/) (2019)
- [Optimal Algorithms for Testing Closeness of Discrete Distributions](https://mlsanity.com/paper/1308.3946/) (2013)
- [Optimal Testing for Properties of Distributions](https://mlsanity.com/paper/1507.05952/) (2015)
- [Testing distributional assumptions of learning algorithms](https://mlsanity.com/paper/2204.07196/) (2022)
- [A Moment-Matching Approach to Testable Learning and a New Characterization of Rademacher Complexity](https://mlsanity.com/paper/2211.13312/) (2022)
- [Near-optimal-sample estimators for spherical Gaussian mixtures](https://mlsanity.com/paper/1402.4746/) (2014)
- [The Structure of Optimal Private Tests for Simple Hypotheses](https://mlsanity.com/paper/1811.11148/) (2018)
- [Differentially Private Testing of Identity and Closeness of Discrete Distributions](https://mlsanity.com/paper/1707.05128/) (2017)
- [Testable Learning with Distribution Shift](https://mlsanity.com/paper/2311.15142/) (2023)
- [Differentially Private Chi-Squared Hypothesis Testing: Goodness of Fit and Independence Testing](https://mlsanity.com/paper/1602.03090/) (2016)
- [Faster and Sample Near-Optimal Algorithms for Proper Learning Mixtures of Gaussians](https://mlsanity.com/paper/1312.1054/) (2013)
- [Learning the Positions in CountSketch](https://mlsanity.com/paper/2007.09890/) (2020)
- [Triangle and Four Cycle Counting with Predictions in Graph Streams](https://mlsanity.com/paper/2203.09572/) (2022)
- [Testing Identity of Structured Distributions](https://mlsanity.com/paper/1410.2266/) (2014)
- [Optimal Identity Testing with High Probability](https://mlsanity.com/paper/1708.02728/) (2017)
- [Efficient Testable Learning of Halfspaces with Adversarial Label Noise](https://mlsanity.com/paper/2303.05485/) (2023)
- [An Efficient Tester-Learner for Halfspaces](https://mlsanity.com/paper/2302.14853/) (2023)
- [Locally Private Hypothesis Selection](https://mlsanity.com/paper/2002.09465/) (2020)
- [Tester-Learners for Halfspaces: Universal Algorithms](https://mlsanity.com/paper/2305.11765/) (2023)
- [Learning-based Support Estimation in Sublinear Time](https://mlsanity.com/paper/2106.08396/) (2021)
- [Testing probability distributions underlying aggregated data](https://mlsanity.com/paper/1402.3835/) (2014)
- [Collision-based Testers are Optimal for Uniformity and Closeness](https://mlsanity.com/paper/1611.03579/) (2016)
- [Private Identity Testing for High-Dimensional Distributions](https://mlsanity.com/paper/1905.11947/) (2019)
- [The Price of Tolerance in Distribution Testing](https://mlsanity.com/paper/2106.13414/) (2021)
- [Test without Trust: Optimal Locally Private Distribution Testing](https://mlsanity.com/paper/1808.02174/) (2018)
- [Learning and Testing Junta Distributions with Subcube Conditioning](https://mlsanity.com/paper/2004.12496/) (2020)
- [Learning Intersections of Halfspaces with Distribution Shift: Improved Algorithms and SQ Lower Bounds](https://mlsanity.com/paper/2404.02364/) (2024)
- [Private High-Dimensional Hypothesis Testing](https://mlsanity.com/paper/2203.01537/) (2022)
- [Optimal Algorithms and Lower Bounds for Testing Closeness of Structured Distributions](https://mlsanity.com/paper/1508.05538/) (2015)

## Built on by (most cited first, 2 of 2)

- [Token-Efficient Change Detection in LLM APIs](https://mlsanity.com/paper/2602.11083/) (2026)
- [Optimal Prediction-Augmented Algorithms for Testing Independence of Distributions](https://mlsanity.com/paper/2603.04635/) (2026)

## Abstract (arXiv)

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.