AAAI 2025 · Semantic Compression & Clustering

Information-Theoretic Generative Clustering of Documents

Xin Du · Kumiko Tanaka-Ishii

Documents are represented as conditional generative distributions, and KL clustering distortion is optimized directly over the infinite discrete space of possible texts.

From Point Representations to Generative Distributions

Document clustering usually maps each document to one vector and groups the vectors by Euclidean or cosine distance. The convenience comes with a strong prior: each document becomes a single point. Sparse text omits much topic-related knowledge, while one document may invite several queries, paraphrases, or interpretations. Averaging these possibilities into one vector can erase multimodal semantics and the directions that distinguish clusters.

A generative language model offers another representation. Given document $x$, it defines a conditional distribution $p(Y\mid x)$ over all possible texts $Y$. Generated sentences are samples; the representation itself contains the probability assigned to each possible query or description. Each cluster $k$ can therefore have a centre distribution $p(Y\mid k)$, with objective

KL divergence ties semantic difference directly to generation probabilities and preserves multiple latent textual directions. “Generate, embed, then apply $k$-means” does not retain the same structure: it compresses the samples back into one approximately Gaussian point and may compare them in a geometry defined by a different model.

Computation over an Infinite Text Space

The objective is clear but intractable by enumeration. $Y$ is an infinite discrete space, cluster centres are unknown, and every document induces a different distribution. The paper samples $J$ texts from a shared proposal $\phi(Y)$ and estimates the document-to-centre distortion:

The proposal first samples a document uniformly, then generates from its conditional model, covering high-probability regions across the corpus. Cluster assignments and centre updates alternate; the centre probabilities on shared samples are normalized regularized averages over documents in the cluster. The update has the structure of Bregman hard clustering, and each iteration is proven not to increase the objective.

The main difficulty is variance rather than writing a Monte Carlo sum. Language-model importance ratios $p(Y\mid x)/\phi(Y)$ are extremely right-skewed; a few samples dominate the unbiased estimate and make finite-sample clusters unstable. The implementation clips anomalous log probabilities and derives a power-mean proposal to reduce aggregate variance.

Why Regularized Importance Sampling Is Necessary

At $\alpha=1$, the estimator is asymptotically unbiased. Setting $\alpha=0.25$ introduces bias but greatly compresses the dynamic range of weights. Clustering does not require every KL value to be unbiased; it requires relative distortions to be stable enough to recover the grouping.

R2 makes the mechanism visible. At $\alpha=1$, NMI is 25.8; at $\alpha=0.25$, it reaches 77.8. Reducing $\alpha$ further eventually loses too much information. This comparison holds the representation and objective fixed and directly validates the bias–variance mechanism on which the method depends.

Small datasets stabilize once $J\geq50$; Yahoo! Answers stabilizes near $J=384$, still below the 768 dimensions of a common BERT vector. The method continues to outperform the compared SBERT clustering when the requested cluster count $K$ is misspecified between 2 and 20.

Document-Clustering Results

Experiments compare TF–IDF, BERT, SBERT, and deep embedding clustering on R2, R5, AG News, and Yahoo! Answers. Generative clustering obtains the highest ACC, NMI, and ARI on all four. On R2, the scores are 96.1, 77.8, and 84.9; on AG News, 85.9, 64.2, and 67.2; on the larger and finer-grained Yahoo! Answers, 60.7, 43.7, and 35.5.

The gain is not simply a stronger pretrained encoder. Generated query-like texts expand knowledge that is implicit but inferable from the document into a probability structure; the KL objective then compares documents and cluster centres in that same generative space. Representation and clustering geometry remain aligned.

From Flat Clustering to Generative Indices

Generative document retrieval requires hierarchical clusters that become prefix codes. After descending into a subcluster, a global proposal poorly covers local documents. The paper estimates a local proposal $\widetilde\phi$ and resamples existing texts with

without recomputing the complete probability matrix. Every level therefore estimates KL distortion in the text region relevant to its current subcluster.

With the same approximately 30M-parameter T5 retriever, changing only index construction raises MS MARCO Lite Rec@1 from NCI’s 17.91 and BMI’s 23.78 to 32.41—a 36% gain over BMI. NQ320K rises from 55.17 to 56.40. Reducing samples from 4,096 to 1,024 changes MARCO Rec@1 only from 32.41 to 32.16, whereas removing the local proposal reduces it to 30.37.

Scope and Cost

The method uses a language model for clustering through its full conditional distribution rather than through a single generated vector. It applies to topic organization, hierarchical corpus indexing, and generative retrieval, and supplies an information-theoretic template for other multimodal semantic objects.

The primary cost is the document-by-sample probability matrix. State caching and low-precision inference reduce the burden; on R2, the paper computes it in roughly ten minutes on one GPU, with no material degradation under BF16. Continually changing corpora require incremental proposals and probability reuse. The robust choice $\alpha=0.25$ is not universal; adapting regularization through effective sample size remains an important direction.

← Semantic Compression and Clustering