Paper: Partitioned-Learned Count-Min Sketch (ICLR-2024)

Reviewer comment (Evaluation):
The robustness (i.e., what happens if the prediction is completely wrong) is not discussed/evaluated. Indeed, one major motivation for algorithms with predictions is to utilize ML predictions while still preserving the worst-case guarantee.

Author response:
Given fixed thresholds, even when the predictions are completely inaccurate, PL-CMS optimizes the space allocations of each CMS accordingly, which means that PLCMS is naturally robust. Specifically, we use the predictions to estimate the quantities $E_i$ and $F_i$ as well as the number of unique buckets from the training dataset. If, for example, we incorrectly place a large number of infrequent items in the first region, $F_1$ will be large and we will assign $\text{CMS}_1$ a lot of space to lower its FPR. 

As an example, consider the case when our ML model is just a random oracle that uniformly assigns elements of the stream into different regions of PL-CMS, without any correlation to their actual frequencies. Further, assume that the number of unique buckets is small (which is typically true in practice). In this setting, one can show that the expected error of any item under PL-CMS is roughly equal to its expected error under the standard Count-Min sketch -- that is, when the learned model degenerates to random guesses, PL-CMS still matches the performance of the unlearned algorithm, indicating a high level of robustness. Since a random oracle will allot every region a roughly equal fraction of infrequent items $F_i$ and a roughly equal total frequency of items $E_i$ and, as a result, each CMS will be assigned the roughly same amount of buckets $m_i$, we will have $\frac{E_i}{m_i}\approx \frac{n/c}{m/c} = \frac{n}{m}$, where $c$ is the number of regions. This is the same ratio (and in turn yields the same false positive rate) as if we assigned all items to a single CMS with $m$ buckets. Note that a similar argument would hold if the model assigned items to different regions with different probabilities. In this case, letting $p_i$ be the probability of an item being assigned to region $i$, we would have $m_i \approx m \cdot p_i$.

To verify the above argument, we have included additional experiments in section A.3.3 of the Appendix, where we show that given a learned model that outputs random predictions, PL-CMS performs similarly to standard CMS and LCMS on the Zipfian datasets except for 3-CMS on the Macbeth dataset with $n/k=16$ which we are currently investigating.
We would also like to point out that while it is important to consider the robustness of algorithms with predictions, none of the prior works in the area (LCMS of Hsu et al. and Partionioned Learned Bloom Filters of Vaidya et al.) provide such an analysis.

Reviewer comment (Evaluation):
It would also be better to have a measure of the prediction error, as well as relate the performance of your algorithm to that error measure to obtain a smooth tradeoff between robustness and consistency.

Author response:
We agree that establishing a measure of prediction error is a crucial next step to optimizing the threshold values of PL-CMS and studying its consistency.

Reviewer comment (Theory):
It seems your E_i and F_i are dependent on the thresholds t_i’s. However, it seems you need to estimate E_i and F_i, and then use them to find the t_i’s. This does not make sense to me.

Author response:
The threshold values do not depend on $E_i$ and $F_i$. We first perform a grid search to find the threshold values and then estimate $E_i$ and $F_i$ given the chosen thresholds.

Reviewer comment (Theory):
Page 6, third paragraph. Can you give more intuition on why the theoretical upper bound does not work well? Is there any rationale for introducing a p in the exponent? This way of adding a new parameter seems quite random to me.

Author response:
The theoretical bound on the false positive rate is loose due to applying Markov's inequality in the analysis. The term inside the parentheses found in Equation 6 is the failure probability for any repetition which is set to $1/e$ by picking $m_i$ appropriately. We then raise it to the number of repetitions $t_i$. Thus, adding a constant $p$ is equivalent to estimating the per repetition failure probability as $1/e^p$ instead of $1/e$ which makes it closer to reality.

Reviewer comment (Evaluation):
In page 6 you mentioned that using larger number of score thresholds typically does not improve the performance — why is this? Also, this somewhat contradicts the claim in Section 5, where you mention that “determining the optimal number of partitions is a crucial next step that can significantly enhance its performance”.

Author response:
This is a great question. We agree that optimizing the thresholds themselves rather than the number of regions is more important to improving our algorithm's performance. Finding a way to optimize the number of regions could potentially allow us to better understand why using a small number of regions often works well in practice.

Reviewer comment (Theory):
Page 7, why the space of the learned model is of the lower order so that it can be ignored?

Author response:
Thank you for pointing this out. We follow the approach of Hsu et al. in which we focus on the space complexity of the data structure instead of its learned model since the space occupied by the oracle can be amortized over time. We will include a clarification in the paper.

Reviewer comment (Theory):
In this paper, the definition of the heavy items we are interested in is the i such that \( f_i \ge n/k \). In a number of the works, the heavy hitter also be defined as \( f_i \ge \sum_j f_j / k \), can the analysis in this work be extended to this model?

Author response:
We define $n$ as the sum of all frequencies in the data stream so the definitions are equivalent. The number of distinct items in the stream does not directly come into our bounds.

Reviewer comment (Evaluation):
In the experiments, the authors study the performance of the algorithm using both the ideal prediction and the noisy prediction. The result shows that there are still some gaps in performance between the two cases. I think it would be an interesting part if the author could give a (brief) analysis of the precision of the current prediction.

Author response:
We believe the reviewer asks about the effect of estimating the $E_i$, and $F_i$ values on the algorithm's performance. We agree that it is an interesting direction to pursue. While we currently do not have a formal analysis of the estimates' precision, since we assume that the training dataset comes from the same distribution as the testing dataset, we expect the approximated $E_i$ and $F_i$ to be close to the true values (which we observe in practice).

Reviewer comment (Experiments):
Comparison with other 2 methods of Zhang et. al. (2020): The paper hinted that their approach involves highly accurate learned model and hence omitted for comparison. Are there ways to compare them on equal grounds?

Author response:
We were not able to compare the two algorithms on equal grounds on the Beijing PM2.5 dataset. We trained a learned model according to the approach of Zhang et al. with the same architecture and specified $5\%$ training dataset split size as well as the reported parameters but our resulting model was not accurate and had an MSE of $31552.73$. Since the authors' algorithm relies on accurate predictions and they did not describe how to optimize its parameters, we were not able to proceed further. We also tried to run Zhang's algorithm using our model's architecture and the same train-test ratio of 8:2. For $n/k=420$ with space allocation of $1954$ units, Zhang's algorithm (with parameters reported in their paper) yields an FPR of $0.9$, whereas PLCMS achieves a much lower FPR of $0.367$ for the same space budget.

Reviewer comment (Experiments):
What is the model size for LCMS and PL-CMS for each datasets? Can we train them to achieve similar accuracy levels as in Zhang et. al. (2020)? What are the bottlenecks?

Author response:
We use the same learned model for LCMS and PL-CMS. The size of the model for the Zipfian datasets is equal to the size of the dictionary that stores the rank of each word in the training dataset. Following the approach of Hsu et al., we focus on the space complexity of the data structure instead of its learned model which can be amortized over time (i.e., if identifying frequent items over many days or other time periods, the model remains fixed, while a new data structure is used in each round). 

Since the approach of Zhang et al. requires adjusting the frequencies of items during training, it cannot be directly applied to the Zipfian datasets. Furthermore, Zhang's algorithm requires retraining the model for every different $n/k$ value as well as fine-tuning many of its parameters.

Reviewer comment (Reproducibility):
Is the code available for replication of plots/results?

Author response:
While the code is not currently published, we are working to make it available.