Lego Sketch: A Scalable Memory-augmented Neural Network for Sketching Data Streams

TL;DR

This work tackles scalable neural sketches for streaming item-frequency estimation by introducing the Lego sketch, a modular memory-augmented neural network that combines a domain-agnostic normalized multi-hash embedding, scalable memory distributed across $K$ bricks, and a memory-scanning and ensemble-decoding pipeline. It provides theoretical guarantees for domain and memory scalability (Theorems 1–2) and a bound on the rule-based estimator error (Theorem 3), and introduces a self-guided weighting loss to balance meta-task objectives. Empirically, Lego surpasses handcrafted and neural baselines across real and synthetic datasets, maintains robustness under distributional shifts, and delivers competitive throughput, with notable gains when extended as a derivative core (e.g., ES as D-Lego). The work advances practical, provably scalable sketching for data streams and opens pathways for integrating Lego as a core in downstream neural-derivative architectures with strong space-accuracy trade-offs.

Abstract

Sketches, probabilistic structures for estimating item frequencies in infinite data streams with limited space, are widely used across various domains. Recent studies have shifted the focus from handcrafted sketches to neural sketches, leveraging memory-augmented neural networks (MANNs) to enhance the streaming compression capabilities and achieve better space-accuracy trade-offs.However, existing neural sketches struggle to scale across different data domains and space budgets due to inflexible MANN configurations. In this paper, we introduce a scalable MANN architecture that brings to life the {\it Lego sketch}, a novel sketch with superior scalability and accuracy. Much like assembling creations with modular Lego bricks, the Lego sketch dynamically coordinates multiple memory bricks to adapt to various space budgets and diverse data domains. Our theoretical analysis guarantees its high scalability and provides the first error bound for neural sketch. Furthermore, extensive experimental evaluations demonstrate that the Lego sketch exhibits superior space-accuracy trade-offs, outperforming existing handcrafted and neural sketches. Our code is available at https://github.com/FFY0/LegoSketch_ICML.

Figures (11)

• Figure 1: Sketch Literatures
• Figure 2: Space-accuracy Trade-off (Aol Dataset)
• Figure 3: Lego Sketch Overview: The Lego sketch enables a scalable and unified framework capable of adapting to different domains and space budgets, mirroring the modular design seen in dedicate creations built from multiple Lego bricks.
• Figure 4: Framework of the Lego Sketch ( When an item $x_i$ is stored, it undergoes $\mathcal{E}$ and $\mathcal{A}$, obtaining the embedding vector $v_i$ and address vector $a_i$, which are subsequently stored into the distributed memory brick $M$ within $\mathcal{M}$. When querying an item $x_i$, its embedding vector $v_i$ and address vector $a_i$ are obtained in the same way. These, combined with stream characteristics $s$ reconstructed from the current memory brick by $\mathcal{S}$, are input into $\mathcal{D}$ for frequency estimation $\hat{f}_i$.)
• Figure 5: Sub-skewness $\alpha'$

Theorems & Definitions (10)

• Theorem 4.1
• Theorem 4.2
• Theorem 4.3
• Definition 4.1
• Theorem 5.1
• proof
• Theorem 6.1
• proof
• Theorem 7.1
• proof