Semantic Compression and Clustering

Knowledge representation in foundation models is constrained jointly by parameter count, index length, and context windows. We use rate–distortion theory and information bottlenecks to ask which task-relevant information should survive finite capacity, what constitutes a computable semantic distortion, and whether a generative model can directly supply the required probability structure.

Keywords information bottlenecks · rate–distortion · conditional generative distributions · KL divergence · importance sampling

The Information Bottleneck in Generative Retrieval

Generative document retrieval compresses each document $X$ into a short identifier $T$ and trains an autoregressive model to generate $T$ directly from a query $Q$. Search becomes conditional generation, but the short index and finite model parameters must jointly memorize how queries identify documents. Poorly structured indices create errors along shared prefixes, and small models lack the capacity to correct them later in decoding.

Conventional docids are derived from document clusters, random numbers, or text fragments. These methods optimize document similarity, although the task requires preserving only the information needed to distinguish future queries. Index design can instead be stated as task-constrained compression:

The optimal index depends on the joint distribution $p(X,Q)$, especially $p(Q\mid X)$, rather than surface document similarity alone. Similar documents that invite different queries should not share confusable prefixes; differently worded documents serving the same query class may share longer prefixes. The lower-left boundary of the information-bottleneck curve represents minimum query-information loss at a fixed index rate.

Constructing Discrete Indices from Query Distributions

Our ICML 2024 Oral paper introduces Bottleneck-Minimal Indexing (BMI). Because observed queries are sparse, a language model generates synthetic queries to approximate $p(Q\mid X)$; observed queries and document passages provide complementary signals and reduce estimation bias. Their representations are then organized by hierarchical $k$-means into a prefix-structured discrete code.

Experiments estimate where each index lies in the $I(X;T)$–$I(X;Q\mid T)$ plane as well as measuring retrieval accuracy. BMI lies closer to the lower-left boundary: the gain is not explained by longer identifiers or larger clusters but by lower query-relevant information loss at the same rate.

On MARCO Lite with T5-mini, BMI raises Rec@1 from 7.09 to 13.54, an absolute gain of 6.45 points and a relative gain of approximately 91%. On NQ320K, the absolute gain is 7.06 points. Improvements remain with T5-base but shrink to 3.72 and 1.26 points. This capacity dependence is predicted by the bottleneck view: index structure matters most when model parameters cannot memorize an arbitrary code.

Documents as Conditional Generative Distributions

The indexing problem says what finite codes should retain, but it still requires a semantic distortion between documents. Vector clustering assumes that document meaning is approximately unimodal and Euclidean in the chosen embedding. That assumption is restrictive for short texts, implicit knowledge, and multi-topic documents. Using different embedding and generation models also makes the geometry of indexing inconsistent with that of retrieval.

An alternative representation is the conditional distribution $p(Y\mid x)$ over all possible generated texts $Y$. It contains facts, explanations, and contexts that a document can activate, including reliable knowledge not explicitly written in the document. Multiple possible continuations retain multimodal structure. A cluster centre becomes a generative distribution capable of explaining its documents rather than an average vector.

In our AAAI 2025 paper, the distortion from a document to a cluster is $\mathrm{KL}\!\left[p(Y\mid x)\,\|\,p(Y\mid k)\right]$. This yields Bregman hard clustering with alternating assignments and a closed-form centre update; every iteration decreases the objective. The computational obstacle is that $Y$ is the infinite discrete space of all finite texts.

We use regularized importance sampling. All documents share samples from a proposal $\phi$, making their distances comparable. Standard importance weights have extreme variance when proposal and target distributions differ, so an exponent $\alpha$ controls the bias–variance trade-off:

At $\alpha=1$, the estimator is asymptotically unbiased. Smaller $\alpha$ compresses extreme weights and trades limited bias for much lower variance. With $\alpha=1$, NMI on R2 is only 25.8; at $\alpha=0.25$, it reaches 77.8, showing that stable finite-sample comparison matters more than formal unbiasedness.

Generative clustering outperforms the compared vector-clustering baselines on R2, R5, AG News, and Yahoo! Answers. It can also replace the hierarchical clustering in BMI. With the same 30M-parameter retrieval model, Rec@1 rises from 23.78 to 32.41 on MS MARCO Lite and from 55.17 to 56.40 on NQ. Conditional generative distributions therefore change not only offline corpus analysis but also the retrievability of a discrete codebook.

Toward Updatable Knowledge Systems

The two projects form one technical programme. Query-relevant information determines the compression target; conditional generative distributions determine semantic distortion. Together they recast heuristic docid design as an analysable rate–distortion problem.

The remaining limitations also stem from probability estimation. Bias in $p(Q\mid X)$ propagates through a hierarchical index; generative clustering evaluates a document-by-sample probability matrix; a fixed proposal and $\alpha$ need not suit every multimodal distribution. Under data drift, regenerating every query and code can be expensive. Current directions include adaptive proposals, incremental cluster centres, variable-rate indices, and joint allocation of capacity across parametric memory, external indices, and context windows.

Related Papers and Code

Xin Du, Lixin Xiu, and Kumiko Tanaka-Ishii. Bottleneck-Minimal Indexing for Generative Document Retrieval — paper overview. ICML 2024 (Oral). arXiv · Code

Xin Du and Kumiko Tanaka-Ishii. Information-Theoretic Generative Clustering of Documents — paper overview. AAAI 2025. arXiv · Code