Efficient Clustering with Provable Guardrails for LLM Inference at Scale

AI summary

Overview

Research area: Machine learning — scalable clustering for large language model (LLM) inference, applied to eCommerce recommender systems.

Technical level: Intermediate (readers should be comfortable with clustering concepts, embedding similarity, and Big-O complexity notation).

Scope: This paper proposes Scalable Guardrailed Clustering (SGC), a two-stage clustering algorithm that guarantees a user-specified minimal semantic similarity and exact attribute match between every sample and its cluster representative, and demonstrates it on a deployment covering 38 million customers.

What This Paper Is About

Running an LLM once per user or per item is cheap individually, but over millions of inputs the combined cost and latency become prohibitive. The standard workaround is to cluster the inputs, run the LLM only on cluster representatives, and propagate the outputs to cluster members — but then each member inherits an output that may not fit it, and off-the-shelf clustering methods only optimize average-case quality with no per-sample guarantee. The paper's goal is a clustering method that both scales to tens of millions of inputs and guarantees, for every single sample, a minimum similarity and exact attribute match with its representative.

Key Contributions

1. A scalable clustering algorithm with provable per-sample guardrails. SGC guarantees a user-specified minimal semantic similarity between each sample and its cluster representative, plus exact matching of user-specified attributes, protecting the customer experience (e.g., preventing age-inappropriate product recommendations).
2. A tail-trimming benefit from the greedy selection step. The greedy representative selection produces a heavily skewed cluster size distribution, which allows removing a large percentage of tail clusters at the cost of discarding only a small percentage of samples, further improving data reduction.
3. Benchmarks against common clustering methods on both internal shopping-persona data and the public AG News and Cosmopedia datasets, along with theoretical guarantees and a full time/memory complexity analysis.
4. A real-world deployment that clusters 38 million customers, cutting downstream LLM cost and runtime by 50× and unblocking the launch of a persona-based recommender system that showed gains in revenue and engagement in an A/B test.

Main Findings

• Guardrails are met exactly, not on average. Across Shopping Personas (threshold 0.75), AG News (0.3), and Cosmopedia (0.4), SGC and SGC-R place 0.0% of samples below the threshold, while average-metric baselines range up to 21% below threshold (Spectral Clustering on AG News: 21.2%).
• Large speed and memory advantage over threshold-based baselines. On Shopping Personas, SGC runs in 2.7 seconds versus 768 seconds for SimClus, and the proposed method is described as one to three orders of magnitude faster, with an up-to-1000× runtime and memory advantage in Section 4.1.
• Consistent data reduction at matched thresholds. At the matched thresholds, SGC produced 2545 clusters on Shopping Personas, 1795 on AG News, and 2399 on Cosmopedia, with SGC-R achieving higher average similarity (0.825, 0.490, 0.588 respectively) at the same cluster counts and runtime.
• SimClus produces fewer clusters but is far slower. SimClus yields 1466 clusters on Shopping Personas (0.783 average similarity, 0.0% below threshold) but takes 768 seconds, versus 2545 clusters at 2.7 seconds for SGC; the paper describes SimClus as a special case of SGC with K=1 and no attribute matching or reassignment.
• Community Detection does not actually honor the threshold. 35.7% of its samples fall below α on Shopping Personas, 47.3% on AG News, and 45.9% on Cosmopedia, because its overlap-removal step can reassign a sample to a representative outside its threshold.
• Star Clustering produces more clusters. Using degree-based rather than coverage-based selection, it produces up to 2× more clusters than the proposed method (5173 on Shopping Personas versus 2545 for SGC) and lower average similarity than the reassignment variant.
• Tail trimming gives a large additional gain. Keeping only about 4% of clusters covers 90% of Shopping Personas, offering an additional approximately 25-fold reduction in downstream computation at the cost of removing 10% of samples, which the paper notes can be compensated for by oversampling.
• Data reduction is still strong at a high threshold. Even at a similarity threshold of 0.75, the method achieves 186-fold data reduction.
• Initial cluster count is a scalability knob, not a quality knob. Increasing K reduces Stage 2 time linearly and peak memory quadratically, at the cost of slightly reduced data-reduction efficiency. On 100K shopping personas, varying K from 10 to 250 changed the final cluster count by only about 50% while keeping minimal within-cluster similarity at the user-specified threshold.
• Correctness is independent of Stage 1. Theorem 3.1 holds for any Stage 1 partition, including K=1; the initial clustering affects only the number of final clusters and runtime, never correctness.

Methodology in Plain English

The method works in two stages. First, all customer personas are converted to numeric embeddings using a language model — the paper selects all-MiniLM-L6-v2 from the SentenceTransformer model zoo based on dual support for Euclidean distance and cosine similarity, benchmark performance, and latency. Those embeddings are clustered coarsely with Mini-batch K-Means, which is fast and memory-efficient compared with standard K-Means and does not need to be accurate, since a second stage refines it.

Second, within each initial cluster the algorithm computes pairwise cosine similarity and builds a match matrix: two customers count as a "match" only if their similarity exceeds the user-specified threshold and they share the same user-specified attributes. It then greedily picks the customer matching the largest number of still-unmatched customers, makes that customer the representative, assigns everyone they match to that cluster, removes them from the pool, and repeats until everyone is assigned. This stage is the classical Johnson–Chvátal greedy Set Cover heuristic applied inside each initial cluster, which is why the guardrail guarantee follows.

An optional variant, SGC-R, then reassigns each customer to the most similar representative among those whose match relation covers them, which weakly increases within-cluster similarity without breaking the guardrail. The paper proves both properties in Theorem 3.1 and bounds the representative count by (1 + ln|C_k|) · OPT_k via the classical Johnson–Chvátal analysis.

Because the quadratic pairwise computation is confined within each initial cluster rather than computed over the full dataset, total time is O(nd + n²d/K) and total memory is O(nd + n²/K²), versus O(n²d) time and O(n²) memory for Star Clustering, SimClus, and Community Detection. The paper notes Agglomerative Clustering would need roughly 90 TB at n = 5M. All benchmarks ran on a single r7i.12xlarge instance without parallelization.

Why This Matters

• Research impact: The paper treats per-sample quality as a formal guarantee rather than an average-case metric, and shows the guarantee can be preserved while confining quadratic similarity computation to within initial clusters — a structural argument for why threshold-based methods need not be O(n²) at the dataset level.
• Real-world applications:
  • Persona-based product recommendation in eCommerce, where a parent of a toddler must not be grouped with parents of older children and receive age-inappropriate recommendations.
  • Any per-item LLM pipeline (search query generation, product description writing, relevance judging, attribute extraction) where cost scales with the number of inputs.
  • Content selection or marketing generation where the paper cites relevance judging through its fine-tuned LLM-as-a-Judge called Marketing Critic.
  • Quality assurance workflows where verification is expensive, since the skewed cluster distribution lets teams validate only the largest clusters that cover most customers.
• Industry relevance: The motivating workload is concrete: generating recommendations for 38 million customers would have taken 114,000 dollars and 23 days for query generation, and 1,018,500 dollars and 485 days for the Marketing Critic step with the allocated compute. Clustering cut downstream LLM cost and runtime by 50× and unblocked a launch that produced significant revenue and engagement gains in an A/B test. The authors are affiliated with Amazon, and the paper targets the GenAIECommerce'26 workshop co-located with RecSys.

Future Directions

• Handling new customers between re-clusterings. The paper notes that new customers can be assigned to existing representatives whose α-balls cover them, or form new clusters otherwise, but does not report evaluation of this incremental path.
• Quantifying the similarity-to-relevance link. Section 4.3 is described as validating the intuition that higher within-cluster similarity yields more relevant product recommendations; the specific results are not included in the available content.
• Mapping the personalization versus data-reduction tradeoff. Section 4.4 is described as investigating how the required minimal similarity and household attributes affect the number of final clusters; outside the 186-fold figure at threshold 0.75, those results are not reported in the available content.
• Reproducibility constraints. The authors state they cannot release the codebase, the internal persona dataset, or the A/B test business metrics, providing pseudocode, hyperparameters, and public-dataset benchmarks instead.

Target Audience

Practitioners and researchers building LLM-powered systems at scale — especially recommender systems, personalization, and eCommerce engineers — who need to cut inference cost without silently degrading per-sample output quality. It is also relevant to applied ML researchers interested in clustering methods that carry formal per-sample guarantees, and to teams weighing threshold-based clustering (Star Clustering, SimClus, Community Detection) against the scalability wall those methods hit at n ~ 10⁷. Readers without exposure to Set Cover or Big-O analysis will find the theoretical sections demanding, but the problem framing, the two-stage algorithm description, and the benchmark tables are accessible.

Authors’ 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.