Skip to content
#data science Preprint

Efficient Clustering with Provable Guardrails for LLM Inference at Scale

Jul 2026 · 0 citations · 32 references
Computer Science Mathematics

TL;DR

This work proposes a two-stage algorithm that generates initial clusters with Mini-batch K-Means, then greedily selects representatives within each initial cluster -- a step equivalent to the Johnson-Chvatal heuristic for Set Cover over alpha-balls in embedding space.

Abstract

Scaling LLM-based applications to millions of users is bottlenecked by the inference cost and latency of modern foundation models. A natural fix is to cluster the inputs and call the LLM only on cluster representatives, letting other members inherit the output -- but this is only safe if each member is measurably close to its representative. Existing clustering methods do not offer such per-sample quality control at scale: none jointly guarantee a minimal within-cluster similarity, exact matching of categorical attributes, and scalability to tens of millions of samples. We propose a two-stage algorithm that generates initial clusters with Mini-batch K-Means, then greedily selects representatives within each initial cluster -- a step equivalent to the Johnson-Chvatal heuristic for Set Cover over alpha-balls in embedding space. The algorithm enforces the similarity and attribute guardrails exactly by construction, and runs in $O(nd + n^2 d/K)$ time and $O(nd + n^2/K^2)$ memory for $n$ samples, feature dimension $d$, and $K$ initial clusters -- linear in $n$ when $K$ grows proportionally with $n$. We provide benchmarks against common clustering methods on internal and public datasets: our method not only delivers per-sample guardrails but also runs 10-1000x faster and scales to data sizes where most standard methods become intractable. Deployed on 38 million customers for a persona-based recommender, the clustering method cut downstream cost and latency by 50-fold while preserving personalization and unblocked the production launch.

View source

Similar papers

Preprint Aug 2026

The Cost of Adaptivity: Matching Lower Bounds Across Learning Problems

A finite-horizon composition law for Gaussian certification from M independent coordinates, a familywise certifier protecting every coordinate and time up to T pays optimal normalized squared half-width of order log(eM) + log log(e^eT), within the sample-mean-centered rectangular class.

Ibne Farabi Shihab, Adria Binte Habib · 0 citations
Preprint Aug 2026

CoRun: Padding is Simple and Efficient for Deterministic LLM Inference

CoRun is presented, a scheduling-based system that achieves deterministic inference without requiring batch invariance, and employs isolated prefill and fixed-shape batched decode to handle the two stages of LLM inference, respectively, leveraging CUDA graphs for efficient execution and simplified implementation.

Shiju Zhao, Jiacheng Yang, Qihang Chen et al. · 0 citations
Preprint Aug 2026

Shape Mutating Expert Compression:LorExperts and BTExperts

LorExperts is introduced, a router-preserving compression method that clusters experts, keeps one full-precision dominant per cluster, and represents the remaining members as low-rank corrections to their local dominant, and BTExperts is introduced, a tree organization of dominants and corrections that enables inference-time amortization of shared computation.

Inesh Chakrabarti, Sourjya Roy, B. Bao et al. · 0 citations
Preprint Aug 2026

Online Differentially Private Consistent Clustering

We study differentially private (DP) $k$-means and $k$-median clustering in the online streaming setting. In this model, points arrive sequentially, and at each time step, we need to output a set of $k$ centers that optimizes the clustering objective for all points seen so far. We give a generic reduction that transforms the (sensitive) input stream into a private stream, which is a semi-coreset of the input stream. This implies that any (non-private) online clustering algorithm, run as a post-processing step, can achieve good utility for the original clustering objective. Our algorithm matches or improves upon the approximation ratio, space usage, and running time of existing algorithms [Epasto et al., 2026, Dupr\'e la Tour et al., 2024]. A key aspect of our reduction is that it inherits desirable properties of the underlying non-private clustering algorithm, such as consistency [Lattanzi and Vassilvitskii, 2017]--a property not satisfied by previous DP algorithms.

Edith Cohen, Vadym Doroshenko, Badih Ghazi et al. · 0 citations
Preprint Aug 2026

DIME: Query-Efficient Framework for Membership Inference on Diffusion Models

Membership inference attacks expose whether individual records were used to train a model, yet existing attacks on diffusion models are largely heuristic and can require substantial query budgets. We introduce DIME (Denoiser Ideal Membership Error), a theoretically grounded and query-efficient framework for membership inference on diffusion models. Our starting point is an exact characterization of the optimal diffusion denoiser for a finite training set, which reveals that membership leakage is governed by the denoiser's implicit reconstruction error. This error decomposes into two complementary signals: a bias term, capturing reconstruction accuracy, and a previously unexplored local crowding term, capturing the geometry of nearby training examples. Both admit efficient estimators using only model queries, yielding a practical attack with as few as two queries. Across CIFAR-10/100, STL10-U, CelebA, and ImageNet, DIME consistently outperforms prior attacks at comparable or substantially lower query cost, improving TPR at 1% FPR by up to $3\times$; remarkably, its two-query variant can outperform existing 30-query baselines. Finally, we suggest, discuss, and evaluate specific defenses to counteract such powerful membership tests.

Tue Do, Daniel Alabi · 0 citations
Preprint Aug 2026

Ranked spreadness and sample-based testing

In this note, we introduce the notion of ranked spreadness, a strengthening of the usual spread condition in which the elements of each member can be ordered so that their one-coordinate marginals decay geometrically with their rank. This additional structure removes the dependence on the maximum set size in random-containment estimates. We prove width-free hitting and weighted-concentration theorems for ranked-spread set systems, together with an elementary kernel-extraction theorem showing that ranked spreadness arises naturally in arbitrary distributions on small sets. Our main application is to the simulation of nonadaptive property testers by sample-based testers. If a one-sided tester has average query complexity $d$ and rejects every far input with probability at least $\delta$, then, for every integer $c>d/\delta$, it admits a one-sided sample-based simulation with expected sample complexity $O_{d,\delta,|\Sigma|}\bigl(n^{1-1/c}\bigr)$. More generally, if positive inputs are rejected with probability at most $\gamma$ and far inputs with probability at least $\delta>\gamma$, the same conclusion holds for every $c>d/(\delta-\gamma)$. In particular, for constant-query nonadaptive testers we obtain an exponent $1-\Theta(1/q)$, matching, up to the dependence on the rejection gap, the exponent conjectured by Fischer, Lachish, and Vasudev.

Gaia Carenini · 0 citations

Related blog posts

Microsoft Research Blog Sep 8, 2026

Called to serve: Tech, research, and positive impact with Chris White

Lab Director Chris White has worked on research challenges with real-world implications—from new approaches to wartime data analysis to tools for combating human trafficking. He talks to program manager Weishung Liu about the influences that led to the work and more. The post Called to serve: Tech, research, and positive impact with Chris White appeared first on Microsoft Research.

We use cookies to run the site and, with your consent, for analytics and to show ads. See our Cookie Policy.