SketchOps

Core implementation accompanying Mergeable Model-Side Aggregation States for Long-Context Language Models.

Long-context language models become increasingly unreliable at non-additive, set-based aggregation as context length grows. Examples include cardinality estimation, set relationships, and grouped statistics, which provide intermediate evidence for downstream reasoning and decision-making.

SketchOps is a model-side aggregation interface for long-context reasoning that maintains compact sketch states alongside a frozen language model. As the model processes a context, an extractor identifies relevant records, assigns each record a canonical identity, and routes it to the corresponding operand-specific stream. The states are updated online in parallel with the model's forward computation and may be merged across compatible segments of the same stream. A query-specific readout converts one or more states into an estimate and returns it to the model as explicit evidence for downstream reasoning, without an additional generate-execute-return cycle.

Problem setting

Many non-additive aggregation tasks depend on how item identities occur across relevant records, including their uniqueness, repetition, grouping, and relationships across collections. The paper refers to this family as identity-aware aggregation.

The current interface supports distinct count, union size, Jaccard similarity, containment, grouped distinct count, and point frequency. HyperLogLog (HLL) provides the support state for the set-based readouts, the HLL joint maximum-likelihood estimator (JMLE) provides the set-relation readouts, and Count-Min provides a complementary state for point-frequency queries.

HLL and Count-Min are established sketch operators. The contribution of SketchOps is the interface that selects and maintains task-appropriate aggregation states alongside a frozen language model and returns the resulting statistics to the model for downstream reasoning.

Model-side aggregation interface

As the frozen language model processes a context, an extractor identifies the relevant records, assigns each record a canonical identity, and routes it to the operand-specific stream required by the request. Each stream maintains a fixed-budget sketch state. Once the records have been processed, a query-specific readout converts one or more states into an aggregate estimate, which is returned to the model for downstream reasoning.

The interface is designed around three requirements from the paper:

• Model-side operation: states are updated as records are extracted, with low computational overhead.
• Fixed budget: state size does not grow with stream length or the number of distinct identities.
• Aggregation consistency: states constructed from compatible segments of the same stream can be merged into the state for the complete stream.

Model-side aggregation workflow. As the frozen language model processes the context, the extractor assigns canonical identities to relevant records and places them in the streams required by the request. For grouped requests, a group key routes each record to the corresponding group state. Each stream maintains a fixed-budget state, and states built from separate segments can be merged without reading the records again. The query-specific readout obtains the requested aggregate from one or more states and returns it to the model for downstream reasoning. HLL is used for set-based readouts, while Count-Min supports point-frequency queries.

Results

• On all 3,969 aggregate-then-reason tasks from 174 source windows, SketchOps achieved 89.0% on Qwen 3.6 (35B) and 99.2% on Gemma 4 (31B, BF16). Exact aggregation achieved 93.0% and 100.0%, respectively.
• On the 174 matched tasks, SketchOps achieved 92.0% on Qwen and 99.4% on Gemma, compared with 31.0% and 43.1% for the strongest full-context baselines.
• On the fixed 1,200-task cohort preselected for comparison with external code execution, SketchOps achieved 91.1% on Qwen and 99.3% on Gemma while generating 2.2 output tokens per task. Exact code execution achieved 100.0% while generating 1,260.2 and 208.6 output tokens per task, respectively.
• With a fixed 2 KiB state, HLL's median relative error remained between 1.29% and 1.49% as stream length increased from 8K to 200K records. In the one-million-record distinct-count experiment, the mean relative error was 1.6%.
• States built from as many as 256 segments produced no register mismatch or readout difference from a single pass over the same stream.

Main results. Colour indicates the model (blue: Qwen 3.6 (35B); vermillion: Gemma 4 (31B)); solid fill is SketchOps and pale fill is the comparison method. (a) 174 matched tasks against the strongest full-context baseline. (b) All 3,969 tasks; the black tick is exact aggregation and the whisker is the 95% window-cluster confidence interval of the paired gap. (c) Mean output tokens per task against external code execution. (d, e) State-level results without a language model: relative error remains stable in stream length under a fixed 2 KiB budget and follows 1.04/sqrt(m).

Repository scope

This repository contains the core aggregation-state implementation:

• fixed-budget HyperLogLog update, cardinality readout, and registerwise maximum for union;
• HLL joint maximum-likelihood readouts for Jaccard similarity and containment;
• fixed-budget Count-Min state for point-frequency queries; and
• the one-million-record error and segmented-merge checks reported in the paper.

Model weights, datasets, prompts, and benchmark outputs are not included.

Install

python -m pip install -e ".[test]"

Contact

Dachuan Song: dsong25@gmu.edu

Citation

If you use SketchOps in your research, please cite:

@article{song2026mergeable,
  title={Mergeable Model-Side Aggregation States for Long-Context Language Models},
  author={Song, Dachuan and Yin, Junyu and Hu, Zechen and Wang, Xuan},
  journal={arXiv preprint arXiv:2607.26448},
  year={2026}
}

Disclaimer

This software is provided under the MIT License.