Optimized Learned Count-Min Sketch

AI summary

Overview

Research area: Learned data structures — specifically probabilistic frequency-estimation sketches that combine a pre-trained machine learning model with classical Count-Min Sketch (CMS). The paper appears in the "Machine Learning for Systems" workshop track and is categorized under Machine Learning (arXiv:2512.12252v1 [cs.LG], 13 Dec 2025, CC BY 4.0).

Technical level: Advanced. The paper relies on the Karush–Kuhn–Tucker (KKT) conditions, Lagrangian duality, Kullback–Leibler (KL) divergence, Jensen's inequality, and dynamic programming. Readers need comfort with probabilistic data structures and convex optimization.

Scope in one sentence: The paper proposes OptLCMS, a variant of Learned Count-Min Sketch that replaces empirical parameter tuning with an analytical optimization of CMS parameters and score-space partitions, aiming to minimize the probability of "intolerable error" under a fixed memory budget.

Authors and affiliation: Kyosuke Nishishita, Atsuki Sato, and Yusuke Matsui, all at The University of Tokyo.

What This Paper Is About

Count-Min Sketch estimates how often each element appears in a multiset using far less memory than exact counting, but it has a fixed trade-off: less memory means more error. Learned Count-Min Sketch (LCMS) improves this by training a model to predict approximate frequencies and using a threshold to route elements either to an exact dictionary or to CMS, but it must tune its parameters empirically on validation data, which makes construction slow, and it offers no theoretical bound on the probability that an estimate exceeds a user-specified error threshold. OptLCMS attacks both problems: it partitions the model's score space into groups, gives each group its own CMS, and derives the CMS parameters in closed form while choosing the partition boundaries by dynamic programming.

Key Contributions

1. An optimization formulation for partitioned learned frequency estimation. The authors adapt the Partitioned Learned Count-Min Sketch (PL-CMS), originally designed for the heavy hitter problem, to the frequency-estimation setting and formalize parameter selection as a constrained minimization of the upper bound on intolerable error probability, ∑_{g=1}^{G} δ_g q_g, subject to a total memory budget M and the CMS constraint δ_g ≤ 1.

2. A closed-form solution for the CMS parameters. For fixed partition thresholds, the problem is convex in the δ vector, and the KKT conditions yield an analytical expression for each δ_g (Equation 5), with an accompanying constant I and the set D = {g | δ_g < 1}. Integer table dimensions from the ceiling functions are relaxed to continuous variables to enable this analysis.

3. Threshold selection via dynamic programming with approximate feasibility checks. Substituting the analytical (ε, δ) into the objective produces an expression containing the KL divergence between the data distribution and the query distribution within each group. The optimal thresholds maximize this divergence via a DP recurrence, and a lemma proved with Jensen's inequality plus a proposition (δ̂_g < 1 ⟹ δ_g < 1) lets the algorithm check feasibility of a candidate cut without knowing how the remaining score space will be divided.

4. Explicit control of the allowable error threshold. OptLCMS exposes the allowable error ε as a user-set hyperparameter, along with the memory budget and the maximum number of partitions G, whereas the authors state that LCMS's only listed hyperparameter is memory usage.

Main Findings

• Lower intolerable error probability. The paper reports that OptLCMS consistently yields the lowest intolerable error probability across the tested memory budgets, while matching LCMS in average error, under both the uniform and frequency-weighted query patterns. OptLCMS achieves up to about 20× smaller intolerable error probability compared to LCMS.

• Average error is not sacrificed. The authors state that minimizing the upper bound on intolerable error does not degrade average error performance relative to LCMS.

• Dramatically faster construction. Table 1 reports LCMS construction at 10.712 s for frequency-weighted and 10.250 s for uniform queries, versus OptLCMS at 0.003 s (a difference of −10.709 s) for frequency-weighted and 4.873 s (a difference of −5.377 s) for uniform queries. The text describes this as orders-of-magnitude faster construction for frequency-weighted queries and more than twice as fast for uniform queries.

• Baseline comparison. OptLCMS is compared against classical CMS and LCMS under the same memory budget, all sharing the learned model from Hsu et al. with amortized model size 0.0152 MB. For each memory budget M, ε is set to the smallest value feasible under CMS, namely ε = e/M, and OptLCMS uses G = 10.

• Evaluation data. Experiments use the AOL query log: 21M queries, 3.8M unique terms, a Zipfian distribution, with queries from the 5th day used for training/parameter tuning and queries from the 50th day used for evaluation/counting.

• Stated limitations of the evaluation. The authors explicitly note that they did not evaluate on the CAIDA dataset or compare against confidence-aware LCMS; they leave these to future work.

Methodology in Plain English

The starting point is a standard Count-Min Sketch, a grid of counters where each element is hashed into one cell per row and its estimated frequency is the minimum of those cells. The width and depth of the grid depend on two parameters, ε (how large an error you will tolerate, relative to the multiset size N) and δ (the probability of exceeding that error). Classical CMS treats the whole element set with one such grid.

LCMS instead trains a model to predict roughly how frequent each element is, then routes elements by score: very high-scoring elements go into an exact dictionary (a unique bucket) and the rest go into CMS. OptLCMS keeps that routing idea but splits the non-dictionary elements into several score-range groups, each with its own CMS tuned to its own statistics. The key question is how to size each CMS and where to place the cut points.

The authors answer this by writing down the goal mathematically: minimize the sum over groups of δ_g times q_g, where q_g is the probability that a query falls in group g — this sum upper-bounds the chance that any estimate's error exceeds εN. They impose a total memory budget that accounts for every group's CMS table plus a per-element memory cost c for the unique bucket. With the cut points held fixed, the problem becomes convex in δ, and the KKT conditions give a formula for each δ_g directly — no grid search, no validation sweep. The remaining question, where to cut the score space, is handled by substituting the formula back into the objective. What remains is an expression involving the KL divergence between how the data is distributed and how the queries are distributed across groups, so the best cuts are the ones that maximize that divergence. A dynamic programming recurrence finds those cuts, with a special check because each CMS must satisfy δ_g < 1 (otherwise its table effectively does not exist). Since DP cannot see the not-yet-processed part of the score space, the authors define an approximate quantity δ̂_g and prove that δ̂_g < 1 guarantees δ_g < 1 however the remaining space is divided.

In practice this means OptLCMS runs a computation over the validation distribution and produces parameters, whereas LCMS has to repeatedly measure real error for candidate settings.

Why This Matters

Research impact. The paper sits at the intersection of learned data structures and classical sketching theory. It shows that a learned sketch can retain the kind of probabilistic guarantee that classical CMS provides — a bound on the probability that error exceeds a user-specified threshold — rather than trading the guarantee away for better empirical error. It also transfers a partitioning-and-optimization template from the Partitioned Learned Bloom Filter line of work and from PL-CMS to the frequency-estimation problem, giving a reusable formulation: minimize an upper bound on intolerable error, solve analytically via KKT, choose cuts via DP over a KL divergence.

Real-world applications (as motivated by the paper's setting and datasets):

• Search engine query analysis, since the evaluation uses the AOL query log of 21 million search queries.
• Natural language and text processing, where the paper notes Zipf's law governs natural language word frequencies.
• Web traffic and network monitoring, where the paper notes Zipf-like distributions are observed in web traffic patterns.
• Streaming analytics and database systems that need approximate frequency counts within a fixed memory budget, since CMS is described as a widely used memory-efficient sketch.

Industry relevance. The practical selling points are the three the authors emphasize: construction that does not require empirical validation sweeps, explicit control over the allowable error threshold, and the ability to fit a fixed memory budget. Construction time matters whenever a sketch has to be rebuilt — for example across changing query workloads or repeated deployment cycles. On the reported hardware (Intel Core Ultra 7 155H, 16 cores, 3.80 GHz), the reported OptLCMS construction time for frequency-weighted queries was 0.003 s, against 10.712 s for LCMS, and per-element unique-bucket memory cost was set to c = 20 bytes.

Future Directions

1. Evaluation on additional datasets. The authors state that they did not evaluate on the CAIDA dataset used in the original LCMS work; testing on it would show whether the reported gains generalize beyond the AOL query log.

2. Comparison with confidence-aware LCMS. The authors also defer comparison against confidence-aware LCMS, which is another route to giving learned sketches probabilistic guarantees and would be a direct competitor to the analytical approach here.

3. Tightening the gap introduced by relaxation and approximation. The formulation relaxes the ceiling functions on table dimensions to continuous variables, and the DP feasibility check uses the approximate δ̂_g rather than the exact δ_g. Quantifying or reducing the error introduced by these steps is a natural follow-up.

4. Generalizing beyond the heavy hitter / frequency estimation split. PL-CMS was built for the heavy hitter problem while OptLCMS targets frequency estimation; understanding whether the same optimization template transfers to other sketch tasks and other learned probabilistic structures is an open direction raised by the paper's framing.

Target Audience

This paper is best suited to researchers and practitioners already familiar with probabilistic data structures — Count-Min Sketch, Bloom filters, and learned variants such as LCMS, PLBF, and PL-CMS — and comfortable with convex optimization and dynamic programming. It will be most useful to systems and database engineers who deploy sketches under hard memory budgets and care about build time, and to machine-learning-for-systems researchers interested in how classical error guarantees can be preserved when a learned model is inserted into a data structure. Beginners will find the appendices necessary reading: the main text states results but defers the CMS background, the full derivations, and the experimental setup to Appendices A, B, and C.

Authors’ abstract

Count-Min Sketch (CMS) is a memory-efficient data structure for estimating the frequency of elements in a multiset. Learned Count-Min Sketch (LCMS) enhances CMS with a machine learning model to reduce estimation error under the same memory usage, but suffers from slow construction due to empirical parameter tuning and lacks theoretical guarantees on intolerable error probability. We propose Optimized Learned Count-Min Sketch (OptLCMS), which partitions the input domain and assigns each partition to its own CMS instance, with CMS parameters analytically derived for fixed thresholds, and thresholds optimized via dynamic programming with approximate feasibility checks. This reduces the need for empirical validation, enabling faster construction while providing theoretical guarantees under these assumptions. OptLCMS also allows explicit control of the allowable error threshold, improving flexibility in practice. Experiments show that OptLCMS builds faster, achieves lower intolerable error probability, and matches the estimation accuracy of LCMS.