PaperlayerView exact revision on arXiv

Source-linked AI summary

Language Models As Semantic Indexers

Bowen Jin, Hansi Zeng, Guoyin Wang, Xiusi Chen, Tianxin Wei, Ruirui Li, Zhengyang Wang, Zheng Li, Yang Li, Hanqing Lu, Suhang Wang, Jiawei Han, Xianfeng Tang

arXiv:2310.07815v3cs.IRcs.CLcs.LG

TL;DR

Semantic ID learning must capture document meaning and hierarchy despite discrete sequential IDs, limited supervision, and shortcomings in two-stage embedding pipelines. LMIndexer learns IDs directly with a generative language model, using progressive and contrastive training plus self-supervised document reconstruction. It produces consistently strong results across recommendation, product search, and document retrieval on five datasets, though direct optimization can encounter reconstructive and posterior collapse.

  • Problem

    Semantic IDs are discrete and sequentially hierarchical, but semantic ID learning lacks ground-truth document–ID pairs and prior two-stage pipelines can lose information.

    Source: §1 · pp. 1–2, §3.1 · p. 3

  • Method

    LMIndexer jointly learns semantic representations and hierarchical IDs with a generative language model, progressive and contrastive training, and self-supervised document reconstruction.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    LMIndexer significantly and consistently outperforms competitive baselines across recommendation, product search, and document retrieval on five datasets from diverse domains.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    The learned semantic indexer can generate IDs for seen and unseen documents and be fine-tuned for downstream recommendation and retrieval.

    Source: §3.3 · p. 5, §6 · p. 9

    Directly optimizing the reconstruction objective can produce reconstructive collapse and posterior collapse.

    Source: §3.2 · p. 5

Abstract

from arXiv · show

Semantic identifier (ID) is an important concept in information retrieval that aims to preserve the semantics of objects such as documents and items inside their IDs. Previous studies typically adopt a two-stage pipeline to learn semantic IDs by first procuring embeddings using off-the-shelf text encoders and then deriving IDs based on the embeddings. However, each step introduces potential information loss, and there is usually an inherent mismatch between the distribution of embeddings within the latent space produced by text encoders and the anticipated distribution required for semantic indexing. It is non-trivial to design a method that can learn the document's semantic representations and its hierarchical structure simultaneously, given that semantic IDs are discrete and sequentially structured, and the semantic supervision is deficient. In this paper, we introduce LMIndexer, a self-supervised framework to learn semantic IDs with a generative language model. We tackle the challenge of sequential discrete ID by introducing a semantic indexer capable of generating neural sequential discrete representations with progressive training and contrastive learning. In response to the semantic supervision deficiency, we propose to train the model with a self-supervised document reconstruction objective. We show the high quality of the learned IDs and demonstrate their effectiveness on three tasks including recommendation, product search, and document retrieval on five datasets from various domains. Code is available at https://github.com/PeterGriffinJin/LMIndexer.

1. Introduction

LMIndexer addresses the difficulty of learning hierarchical, sequential discrete semantic IDs without semantic supervision by jointly learning document representations and structure. It uses a self-supervised generative framework and reports consistent gains across three downstream tasks and five datasets.

  • Semantic IDs are intended to encode coarse document semantics in early positions and increasingly fine-grained content details in later positions.Source: §1 · p. 1
  • Prior methods typically obtain embeddings with off-the-shelf text encoders and then derive IDs using methods such as rq-VAE or hierarchical clustering.Source: §1 · p. 1
  • The paper formulates semantic ID learning as simultaneous modeling of document semantics and hierarchical structure with a generative language model.Source: §1 · p. 2
  • LMIndexer uses a semantic indexer with progressive training and contrastive learning, plus a reconstructor that rebuilds the original text from sequential discrete IDs.Source: §1 · p. 2
  • LMIndexer outperforms competitive baselines significantly and consistently across recommendation, product search, and document retrieval on five datasets.Source: §1 · p. 2

2. Related Work

Prior semantic indexing work combines self-supervised language-model representations with discrete indexing techniques. Related semantic indexers also use discrete latent variables and reconstruction, while document methods commonly separate embedding acquisition from ID construction.

  • Self-supervised language-model training tunes models from unlabeled input corpora, including objectives such as masked language modeling and next sentence prediction.Source: §2 · p. 2
  • Semantic indexers convert inputs into discrete IDs that capture their essence, with auto-encoder designs reconstructing inputs from discrete latent variables.Source: §2 · p. 2
  • Earlier document methods first encode documents with text encoders and then construct IDs from the resulting embeddings using rq-VAE or hierarchical clustering.Source: §2 · pp. 2–3

3. The LMINDEXER Framework

The LMINDEXER framework learns hierarchical sequential semantic IDs directly from documents and trains them through self-supervised reconstruction. Progressive training, codebooks, contrastive learning, and downstream fine-tuning support coarse-to-fine indexing and task use.

3.1. Learning Semantic IDs with Sequential Discrete Auto-Reconstruction

  • The semantic indexer uses an encoder-decoder Transformer to map document text and prior IDs into continuous hidden states, then maps those states to codebook IDs.Source: §3.1 · p. 3
  • LMINDEXER represents documents with sequential discrete IDs whose early positions capture coarse semantics and later positions capture finer details.Different codebooks are maintained at each position to represent different semantic granularities.Source: §3.1 · p. 3
  • A reconstructor receives sequential semantic ID representations and masked document tokens, using the IDs to reconstruct the original document and supply self-supervision.The framework uses a shallow reconstructor so that the semantic indexer must provide high-quality representations.Source: §3.1 · pp. 3–4

3.2. Training Self-Supervised Semantic Indexer

  • Progressive training reduces hints as ID positions increase, encouraging new IDs to encode new information and produce a coarse-to-fine hierarchy.Source: §3.2 · p. 4
  • Contrastive learning separates documents sharing a prefix so later IDs become more diverse and codebook utilization increases.The objective pushes same-prefix documents apart in the current latent space and encourages distinct next-position IDs.Source: §3.2 · p. 4

3.3. Finetuning Semantic Indexer on Downstream Tasks

  • The learned indexer can generate IDs for seen and unseen documents and can be fine-tuned for recommendation and retrieval tasks that output document IDs.Source: §3.3 · p. 5

4. Experiments: Learning Self-Supervised Semantic ID

Experiments evaluate LMINDEXER’s semantic IDs on Amazon product datasets using quantitative, qualitative, and training analyses. The results show semantic-indicative hierarchical IDs and identify warm-up, contrastive learning, and commitment loss as important training components.

Quantitative Results

  • LMINDEXER consistently outperforms baseline methods on AMI between semantic IDs and ground-truth item categories across Amazon datasets.The comparison evaluates whether similar IDs capture similar item semantics.Source: §4.2 · p. 5, Table 1 · p. 6

Qualitative Results

  • The first semantic ID token captures coarse-grained product categories, while the second further divides them into fine-grained categories.On Amazon-Beauty, c1 groups categories such as Bath, Tool, and nail Make-up products, and c2 refines the categories within fixed c1 groups.Source: §4.2 · p. 6

Training Study

  • Without reconstructor warm-up, reconstruction Macro-F1 remains low because the reconstructor cannot provide meaningful training signals to the semantic indexer.Figure 4(a) analyzes this reconstructor-collapse condition.Source: §4.3 · p. 6, Figure 4 · p. 7
  • Without encoder and codebook warm-up, semantic ID perplexity and reconstruction Macro-F1 are low, indicating homogeneous IDs and weak reconstruction signals.These findings correspond to posterior collapse in Figure 4(b,c).Source: §4.3 · p. 6, Figure 4 · p. 7
  • Contrastive loss increases the difference ratio for documents sharing an ID prefix, making their semantic IDs more distinguishable.The loss pushes such documents toward different IDs at the current position, increasing ID diversity and perplexity.Source: §4.3 · p. 7, Figure 4 · p. 7
  • Without commitment loss, perplexity of previously learned IDs decreases while learning a new ID, indicating catastrophic forgetting.The ablation therefore supports commitment loss as necessary to preserve previously learned IDs.Source: §4.4 · p. 7, Table 2 · p. 7

5. Experiments: Downstream Tasks

LMINDEXER is evaluated on recommendation, product search, and document retrieval across Amazon, Natural Question, and MS MACRO datasets. It consistently outperforms baselines in recommendation and product search, and does so in document retrieval except on MS MACRO-1M.

Sequential Recommendation

  • Recommendation experiments use Amazon-Beauty, Amazon-Sports, and Amazon-Toys, with next-item prediction from users’ interaction histories.The compared methods include sequential recommendation models and generative methods using rq-VAE or hierarchical-clustering semantic IDs.Source: §5.1 · p. 7

Main Results

  • LMINDEXER performs consistently better than all baselines across recommendation and product-search datasets, and beats them in document retrieval except on MS MACRO-1M.The document-retrieval result identifies DPR as best on MS MACRO-1M, leaving room for improved semantic indexers.Source: §5.1 · p. 7, §5.2 · p. 8, §5.3 · p. 9

Semantic ID Length & Codebook Size Study

  • Increasing semantic ID length or codebook size improves recommendation performance on Amazon-Beauty.The study varies ID lengths 1, 2, and 3 and codebook sizes 128, 256, and 512.Source: §5.1 · p. 8, Figure 5 · p. 8

Product Search

  • Product-search experiments reuse the three Amazon domains and compare LMINDEXER with BM25, DPR, rq-VAE, hierarchical clustering, and SEAL.DPR outperforms semantic-indexer baselines in many cases, although LMINDEXER remains the strongest reported method overall.Source: §5.2 · p. 8

Reconstructor Study

  • A more powerful reconstructor decreases semantic-indexer and generated-ID quality on Amazon-Beauty.The study compares reconstructors with two and three layers using AMI, Recall@5, and NDCG@5.Source: §5.2 · p. 8, Table 5 · p. 8

6. Conclusions

LMINDEXER learns discrete semantic embeddings and hierarchical structure simultaneously with a self-supervised generative language model. Its semantic indexer transfers to recommendation and retrieval, where it outperforms competitive baselines across three tasks and five datasets.

Conclusion

  • LMINDEXER combines a generative semantic indexer with progressive training, contrastive learning, and document reconstruction to learn sequential semantic IDs.The framework addresses sequential discrete IDs and deficient semantic supervision within one self-supervised training setup.Source: §6 · p. 9
  • Across recommendation, product search, and document retrieval on five diverse datasets, LMINDEXER significantly and consistently outperforms competitive baselines.The learned semantic indexer can be fine-tuned for downstream recommendation and retrieval tasks.Source: §6 · p. 9

A.1. Datasets

The experiments use Amazon review and product-search data across three domains, plus Natural Question and MS MACRO for document retrieval. Dataset construction applies task-specific interaction, corpus, and label-processing procedures.

Amazon Datasets

  • Recommendation and product search use Amazon-Beauty, Amazon-Sports, and Amazon-Toys, with recommendation histories filtered to users and items having at least five interactions.Recommendation uses each user’s last interaction for testing, the second-last for validation, and earlier interactions for training.Source: §A.1 · p. 11

Dataset Statistics

  • Dataset statistics for recommendation and product search are reported in Table 7, while document-retrieval statistics are reported in Table 8.These tables summarize the datasets used across the downstream experiments.Source: §A.1 · p. 11

Document Retrieval Datasets

  • Document retrieval uses Natural Question and MS MACRO, including an MS MACRO-1M subset and merged TREC-DL 2019–2020 evaluation data.The MS MACRO-1M corpus contains one million documents while retaining original training and validation labels.Source: §A.1 · p. 11

A.3. Implementation Details

LMINDEXER uses T5-base to learn length-three semantic IDs with position-specific codebooks and reconstruction hints. Downstream experiments generate item or document IDs from task inputs under shared generative-model settings.

Semantic ID Learning

  • Self-supervised training uses T5-base, length-three semantic IDs, position-specific codebooks, and reconstruction hints for the first two positions.The final position distinguishes documents sharing the first two prefixes, while codebook size varies with dataset scale.Source: §A.3 · p. 11

Training Procedure

  • Algorithm 1 initializes a semantic encoder and reconstructor, warms up reconstruction and codebooks progressively, then jointly trains the framework before returning semantic IDs.The procedure uses KMeans during codebook warm-up and returns the trained semantic indexer with document IDs.Source: §A.3 · p. 12

Implementation Configuration

  • The implementation reports hyper-parameter configurations for semantic-ID learning and generative recommendation, including dataset-dependent training epochs.The detailed configurations are provided in Tables 9 and 10.Source: §A.3 · p. 12, Table 9 · p. 13, Table 10 · p. 13

Downstream Tasks

  • Recommendation and product-search models generate next-item or relevant-item IDs from textual interaction histories or queries using the same T5-base checkpoint.Generative baselines are trained under matched model initialization and task-specific optimization settings.Source: §A.3 · p. 12

Document Retrieval

  • Document-retrieval training augments Natural Question and MS MACRO with 15 and 20 pseudo queries per document, respectively.Generative retrieval models produce relevant-document IDs from queries and use task-specific training schedules.Source: §A.3 · pp. 12–13

A.4. Definition of AMI

Adjusted Mutual Information (AMI) measures agreement between ground-truth and Semantic ID clusters while correcting for chance.

  • AMI quantifies agreement between ground-truth category clusters and Semantic ID clusters while correcting for chance.Source: §A.4 · p. 13

A.5. Duplication Study of Semantic IDs

LMIndexer’s semantic IDs are designed to reduce duplication and preserve useful distinctions, while downstream quality improves with longer IDs and larger codebooks. The method also shows favorable semantic quality, latency, and zero-shot generalization in the reported studies.

Duplication Study of Semantic IDs

  • Adding a contrastive objective promotes different next-position IDs for documents sharing a prefix, alleviating duplication during self-supervised learning.Source: §A.5 · p. 13
  • Most LMIndexer semantic IDs contain fewer than 5 documents, indicating that the learned IDs are largely distinguishable.Because near-identical documents can still share semantics, the method adds a final ID position to distinguish them.Source: §A.5 · p. 13, Figure 6 · p. 15
  • LMIndexer outperforms baseline methods by a large margin in human evaluation of semantic relatedness for product pairs sharing their first two IDs.The evaluation uses 20 pairs per method and four trained annotators.Source: §A.6 · p. 14, Table 13 · p. 14
  • Recommendation performance increases as semantic ID length grows from 1 to 3 on Amazon-Beauty, Amazon-Sports, and Amazon-Toys.Longer IDs can contain more semantic information.Source: §A.7 · p. 14, Figure 7 · p. 15
  • Downstream performance on Amazon-Beauty increases with codebook size, consistent with larger codebooks containing more information.Source: §A.8 · p. 14, Figure 8 · p. 16
  • LMIndexer has product-search inference latency comparable to rq-VAE and HC indexers and much lower than SEAL, while generalizing better to unseen documents in zero-shot search.The latency comparison uses the full Amazon-Beauty test set; the zero-shot study evaluates Amazon beauty items unseen during semantic-index self-supervised training and downstream fine-tuning.Source: §A.9 · p. 14, Table 14 · p. 14, §A.10 · pp. 14–15

Related papers

Loading 2310.07815v3…