PaperlayerView exact revision on arXiv

Source-linked AI summary

Can Generative Retrievers Learn Semantic IDs Without Forgetting How to Speak?

Junchen Fu, Kleomenis Katevas, Vandana Rajan, Sofía Celi, Hamed Haddadi

arXiv:2609.35430v1cs.IR

TL;DR

Generative retrieval can distort a pretrained model’s natural-language behavior, which matters when the same model must retrieve and respond interactively. SpeakGR combines SID learning with language-preserving regularization, while Adaptive SpeakGR adjusts preservation strength to observed drift. The approaches reduce language drift while retaining effective retrieval, and Adaptive SpeakGR improves retrieval over SpeakGR in most settings.

  • Problem

    SID-based retrieval specialization can substantially distort a pretrained model’s natural-language behavior, which matters when a model must also support interactive responses.

    Source: §1 · p. 2

  • Method

    SpeakGR combines supervised SID learning with on-policy distillation from a frozen original model, and Adaptive SpeakGR adjusts preservation strength according to observed language drift.

    Source: §1 · p. 2, §4.1 · p. 5

  • Results

    SpeakGR substantially reduces language drift while retaining effective retrieval, and Adaptive SpeakGR improves retrieval in five of six settings while keeping drift substantially lower than SFT-only.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    SID-based generative retrieval can be learned while substantially limiting drift from the model’s pretrained natural-language behavior.

    Source: §6 · p. 9

    Achieving state-of-the-art retrieval performance is beyond this paper’s scope; stronger retrieval methods are described as complementary to speak-preserving regularization.

    Source: §5.2.2 · p. 8

Abstract

from arXiv · show

Generative retrieval (GR) enables end-to-end retrieval by generating document semantic identifiers (SIDs). However, retrieval-only fine-tuning can over-specialize pretrained language models to SID prediction, substantially distorting their natural-language distribution and limiting their suitability for interactive systems that must both retrieve documents and generate natural-language responses. We introduce SpeakGR, a dual-objective framework that learns SIDs while preserving language generation. It combines supervised SID learning with speak-preserving regularization: an on-policy distillation objective that aligns the current model with a frozen copy of the original model on student-generated prefixes using forward KL over the original text vocabulary. We further propose Adaptive SpeakGR, which dynamically adjusts the preservation strength based on observed language drift. Compared with SFT-only, SpeakGR reduces WikiText-2 forward KL by 81.3-93.8% on MS MARCO and 81.2-85.2% on Natural Questions (NQ) while retaining effective retrieval across three different LLMs. Adaptive SpeakGR further improves retrieval over SpeakGR in most settings while maintaining substantially lower language drift than SFT-only.

1 INTRODUCTION

Generative retrieval trains language models to generate document identifiers, but SID specialization can distort their pretrained language behavior. SpeakGR combines SID learning with language-preserving regularization, addressing the need for models that can retrieve and continue interactive language tasks.

1 INTRODUCTION

  • Generative retrieval reframes retrieval as autoregressive prediction of document identifiers, but standard objectives do not explicitly preserve pretrained language-modeling behavior.A common design extends a pretrained model’s vocabulary with semantic identifier tokens and trains it to generate them.Source: §1 · p. 1
  • This matters for interactive search because a model may need to answer follow-up questions, explain retrieved information, or clarify a query after retrieval.The paper describes a single model retrieving relevant documents and then using their content in downstream language interactions.Source: §1 · pp. 1–2
  • A model that retains both SID and text generation can provide a unified autoregressive interface and may avoid a separate generation model.The paper identifies retrieval and language generation as uses of the same model, with SID resolution connecting retrieval to subsequent generation.Source: §1 · p. 2
  • Standard SID fine-tuning can distort pretrained language behavior even as it learns effective retrieval, with degradation emerging early in training.Across six backbone–dataset settings, the authors report this tension; Figure 1 shows WikiText-2 PPL rising from 31.69 to 6,956 within 25 updates for Qwen3-0.6B on MS MARCO.Source: Figure 1 · p. 1, §1 · p. 2
  • SpeakGR reduces language drift while retaining effective retrieval, and Adaptive SpeakGR improves retrieval in five of six settings.Adaptive SpeakGR also maintains substantially lower language drift than SFT-only.Source: §1 · pp. 2–3
  • SpeakGR couples supervised SID learning with regularization that preserves pretrained language behavior; Adaptive SpeakGR adjusts preservation strength to observed drift.The paper defines language preservation as retaining natural-language predictive behavior, not every specialized capability such as mathematical reasoning or code generation.Source: §1 · p. 2

2 RELATED WORK

Prior work develops generative retrieval and language-preservation methods, while leaving the effect of SID specialization on pretrained language behavior underexplored. This paper studies that effect and formulates SID learning and language preservation as coupled objectives.

2 RELATED WORK

  • Prior generative-retrieval work primarily targets retrieval effectiveness, identifier learning, and generation strategies, while this paper examines effects on pretrained language behavior.The surveyed approaches include alternative identifier designs, retrieval objectives, and decoding strategies.Source: §2 · p. 3
  • Existing forgetting-mitigation approaches use parameter constraints, replay, or distillation; ORBIT instead merges weights after monitoring parameter drift, unlike SpeakGR’s direct distribution regularization.Adaptive SpeakGR adjusts preservation strength online according to observed language drift without weight merging.Source: §2 · p. 3
  • Language preservation during generative-retrieval specialization remains underexplored, and the authors know of no prior work that couples it with SID learning as a training objective.Source: §2 · p. 3

3 PROBLEM FORMULATION

The formulation extends a pretrained language model with SID tokens and trains it to generate document identifiers autoregressively. It defines retrieval supervision and measures language drift by comparing text-token behavior with the frozen pretrained model on natural-language prefixes.

3 PROBLEM FORMULATION

  • The student starts from the pretrained policy with an expanded vocabulary, while the original text-vocabulary policy remains frozen.The expanded vocabulary combines text tokens and SID tokens.Source: §3 · p. 3
  • Each indexed document receives an SID through RQ-VAE or another quantization approach based on its text embedding.Identifier length and level-specific vocabulary sizes depend on the SID construction used in each setting.Source: §3 · p. 3
  • Retrieval generation restricts the student’s next-token distribution to the valid SID vocabulary for each identifier position.The SID distribution is obtained by restricting and renormalizing logits over the level-specific action space.Source: §3 · p. 4
  • Supervision consists of query–identifier pairs, with a separate pair for each query–document association when multiple documents are relevant.The retrieval objective applies level-restricted cross-entropy to supervised SID positions and normalizes by supervised SID tokens; EOS is excluded because identifier length is fixed.Source: §3.1 · p. 4
  • At inference, a trie built from indexed-document identifiers constrains decoding to valid complete identifiers, which map to document buckets ranked internally.The trie is used together with level-specific action spaces.Source: §3.1 · p. 4

3.2 MEASURING LANGUAGE DRIFT

  • Language drift measures how far the specialized model’s text-token distribution departs from the frozen pretrained distribution on natural-language prefixes.Restricting and renormalizing the student distribution over the original text vocabulary isolates changes there from probability assigned directly to SID tokens.Source: §3.2 · p. 4
  • On WikiText-2 test prefixes, a smaller Δlang indicates that the specialized model remains closer to pretrained next-token behavior on natural text.The experiments use fixed prefixes from the WikiText-2 test split as the natural-language prefix distribution.Source: §3.2 · p. 4

4 SPEAKGR: JOINT SID LEARNING AND SPEAK-PRESERVING

SpeakGR trains one student for SID retrieval and natural-language generation, using on-policy distillation to preserve the pretrained text distribution during retrieval specialization. Adaptive SpeakGR changes the preservation weight in response to observed preservation KL while leaving the dual-objective formulation intact.

4 SPEAKGR: JOINT SID LEARNING AND SPEAK-PRESERVING

  • SpeakGR jointly trains the same student parameters for SID retrieval and pretrained natural-language behavior.Retrieval decoding uses level-specific SID vocabularies and a document trie, while language generation uses the original text vocabulary.Source: §4 · pp. 4–5, Figure 2 · p. 5

4.1 SPEAK-PRESERVING REGULARIZATION

  • On-policy preservation matches the frozen pretrained model’s next-token distributions to the student’s at prefixes sampled from student text-only continuations.Forward KL is computed over the original text vocabulary at sampled response positions; the frozen model supplies distributions rather than target sequences.Source: §4.1 · p. 5
  • Forward KL preserves the pretrained text distribution by penalizing insufficient student probability on tokens favored by the pretrained model.The regularizer constrains natural-language next-token distributions but does not directly optimize a particular downstream task capability.Source: §4.1 · pp. 5–6

4.2 SPEAKGR TRAINING

  • The dual-objective training loss combines supervised SID learning with speak-preserving regularization, separately normalizing the two losses by tokens.The retrieval objective drives SID specialization while the preservation objective regularizes natural-language behavior in the same optimization process.Source: §4.2 · p. 6

4.2.1 ADAPTIVE SPEAKGR

  • Adaptive SpeakGR adjusts only the preservation coefficient, using observed preservation KL to update the weight for the next step.The controller smooths and bounds normalized deviation, applies a deadband, and clips the multiplicative update; the dual-objective formulation otherwise remains unchanged.Source: §4.2.1 · p. 6

5 EXPERIMENTS

The experiments compare retrieval and language preservation across datasets and backbones, then examine alternative preservation strategies and Adaptive SpeakGR’s sensitivity to its reference level. SpeakGR substantially reduces language drift while retaining effective retrieval, though the trade-off varies by setting.

5.2.1 RETRIEVAL AND LANGUAGE PRESERVATION

  • 81.3–93.8% on MS MARCO and 81.2–85.2% on NQ: SpeakGR reduces WikiText-2 forward KL relative to retrieval-only SFT across six settings.Retrieval effects vary by setting; Adaptive SpeakGR reaches R@10 0.6745 versus 0.6696 for SFT-only on MS MARCO with Qwen3-1.7B, while reducing FKL from 11.222 to 1.841.Source: Table 1 · p. 7, §5.2.1 · p. 8

5 EXPERIMENTS

  • Across MS MARCO and NQ, the experiments compare retrieval-only SFT, SpeakGR, and Adaptive SpeakGR on three backbones using fixed held-out retrieval and language evaluation sets.Offline replay and ORBIT-style weight merging are also evaluated under a matched protocol.Source: §5 · pp. 6–7

5.1 DATASETS

  • The retrieval corpora contain 319,872 MS MARCO documents and 109,739 NQ documents, with training staged from document-to-SID indexing through query fine-tuning.Language-preservation prompts include 2,048 WikiText-2 continuation prefixes and 2,048 instruction-style prompts; WikiText-2 test is reserved for evaluation.Source: §5.1 · p. 7

5.2.1 RETRIEVAL AND LANGUAGE PRESERVATION

  • SpeakGR substantially reduces directional drift in middle and later layers while allowing representations to adapt for retrieval; language drift mainly reflects changed probabilities over ordinary text tokens.The analysis attributes less of the degradation to probability mass leaking from text tokens to the newly introduced SID vocabulary.Source: §5.2.1 · p. 8

5.2.2 COMPARISON WITH LANGUAGE-PRESERVATION BASELINES

  • Adaptive SpeakGR achieves R@10/FKL of 0.6498/1.728 on MS MARCO and 0.7285/1.116 on NQ, while replay and ORBIT expose different retrieval–preservation trade-offs.ORBIT’s tight threshold lowers FKL to 0.273 but also lowers R@10 to 0.1980; the paper does not aim to achieve state-of-the-art retrieval, and stronger methods are described as complementary.Source: §5.2.2 · pp. 8–9, Table 2 · p. 8

5.3 ABLATION AND ADAPTIVE SENSITIVITY

  • Adaptive SpeakGR’s response to ϵ is non-monotonic: ϵ = 0.25 yields R@10 0.5458 and FKL 2.507, while the best R@10 is 0.6584 at ϵ = 1.5 with FKL 1.566.Moderate values from 0.5 to 2.0 produce a substantially better trade-off region; ϵ is a training reference, not a direct bound on held-out WikiText FKL.Source: §5.3 · p. 9, Figure 3 · p. 9

6 CONCLUSION

The conclusion argues that SID-based generative retrieval can be learned while substantially limiting drift from a model’s pretrained natural-language distribution. SpeakGR jointly trains retrieval and language preservation, while Adaptive SpeakGR adjusts preservation strength based on observed drift.

6 CONCLUSION

  • Retrieval-only fine-tuning can substantially distort pretrained language behavior, but SpeakGR mitigates this while retaining effective retrieval, with trade-offs varying across datasets and backbones.Adaptive SpeakGR adjusts preservation strength in response to observed drift rather than relying only on fixed weighting.Source: §6 · p. 9

AI ASSISTANCE

The authors state that generative AI tools assisted with literature identification, experimental code, prompt-pool curation, figure preparation, and manuscript drafting and polishing.

AI ASSISTANCE

  • The authors report reviewing AI-assisted materials and taking responsibility for the methods, results, and final manuscript.The tools also assisted with identifying literature, experimental implementation code, part of the language-preservation prompt pool, figure preparation, and manuscript text.Source: AI ASSISTANCE · p. 10

A EXPERIMENTAL SETUP

The experiments use a three-stage retrieval schedule across three language-model backbones and two retrieval corpora. Evaluation measures retrieval and language-distribution preservation, while the discussion limits what those measures establish.

A EXPERIMENTAL SETUP

  • The retrieval schedule has three stages, with effective global batches and data order held fixed across GPU counts through gradient accumulation.The stages are document-to-SID indexing, synthetic query-to-SID training, and supervised query-to-SID fine-tuning.Source: Table 4 · p. 13

A.1 MODELS AND TRAINING

  • Experiments fully fine-tune Qwen3-0.6B/1.7B and Gemma-3-1B-IT, using five 256-entry codebooks and five-token SIDs.Within each dataset–backbone comparison, methods share SID assignments and student initialization.Source: §A.1 · p. 13
  • Adaptive SpeakGR uses a controller observing training-stream KL, with settings and stage-specific resets varying by backbone.The main settings use ϵ = 1 nat/token, α = 0.95, c = 0.2, δ = 0.05, η = 0.05, and λmax = 1.Source: §A.1 · p. 13
  • The offline replay baseline matches SpeakGR’s language-data and sequence-length budget, but is a fixed reference rather than a claimed optimum.ORBIT is reimplemented following its original algorithm because no official implementation was publicly available at the time of the experiments.Source: §A.1 · pp. 13–14

A.2 EVALUATION

  • Retrieval is evaluated with Recall@1/10 and MRR@10, while language evaluation uses WikiText-2 forward KL, perplexity, top-1 agreement, and SID probability mass.Comparisons are made within each dataset–backbone setting.Source: §A.2 · p. 14

B DISCUSSION

  • WikiText-2 provides a controlled measure of distributional drift, not a complete measure of language capability.Broader capability preservation is outside the evaluation’s scope.Source: B DISCUSSION · p. 14
  • The experiments do not demonstrate a full end-to-end retrieve–resolve–generate use case; that unified use remains a motivating application.The experiments establish preservation of the pretrained text distribution during retrieval specialization.Source: B DISCUSSION · p. 14
  • Most cross-setting comparisons use one main seed and one SID construction family; specialized capabilities and universal capability-preservation claims remain outside scope.The authors also identify broader multi-seed evaluation and stronger matched controls for the adaptive coefficient as necessary.Source: B DISCUSSION · p. 14

C MECHANISM ANALYSIS

On Qwen3-0.6B with MS MARCO, SpeakGR and Adaptive SpeakGR reduce language-distribution drift and increase shared predictive mass relative to SFT. The section also examines representation changes and describes the prompts used for preservation training.

C MECHANISM ANALYSIS

  • SpeakGR and Adaptive SpeakGR reduce WikiText-2 FKL from SFT’s 8.10 to 1.52 and 1.73, while shared mass rises from 2.85% to 47.85% and 44.52%.SID leakage accounts for only 0.58% of SFT’s excess full-vocabulary NLL; the main degradation is within the normalized text distribution.Source: Figure 4 · p. 15, C MECHANISM ANALYSIS · p. 15
  • At a fixed prefix, the preservation gradient corrects student text-logit probability differences from the frozen base model; reducing KL raises a lower bound on shared probability mass.Transfer of this effect to held-out text is measured empirically.Source: C MECHANISM ANALYSIS · p. 15
  • Final-block directional drift is 1.61 for SFT, 0.23 for SpeakGR, and 0.24 for Adaptive SpeakGR; representation geometry varies with position selection.The evidence supports later-layer directional alignment, not preservation of all representation geometry.Source: Figure 5 · p. 15, C MECHANISM ANALYSIS · p. 16

D LANGUAGE-PRESERVATION PROMPT POOL

B.2. MS MARCO training-query QA (217)

  • The MS MARCO training-query QA subset includes 111 direct-QA prompts and 106 document-conditioned QA prompts.Source: §B.2 · p. 16

B.3. Programmatically instantiated task templates (623)

  • The pool has an instruction subtotal of 2,048 and a complete pool of 4,096.Source: §B.3 · p. 16
  • Students consume selected prompts but not their stored answers: they generate continuations that the frozen base model scores at the same prefixes.Prompt construction includes varying inputs such as numbers for arithmetic and entities for logic.Source: §B.3 · pp. 16–17

E EVALUATION METRICS

Evaluation combines retrieval ranking measures with language-distribution fidelity and SID probability mass, using WikiText-2 positions for the language metrics.

E EVALUATION METRICS

  • Retrieval uses Recall@1, Recall@10, and MRR@10; each query is scored against one designated relevant document.Recall measures whether that document appears in the top 1 or top 10, while MRR@10 averages its reciprocal rank within the top 10.Source: E EVALUATION METRICS · p. 17
  • Language metrics average over the same WikiText-2 positions: forward KL and perplexity are lower-is-better, while top-1 agreement is higher-is-better.Forward KL compares the base distribution with the text-normalized student distribution.Source: E EVALUATION METRICS · p. 17
  • SID mass is the average probability percentage assigned to SID tokens before text-vocabulary masking; lower values indicate less identifier probability in language contexts.Source: E EVALUATION METRICS · p. 17

F FIXED-WEIGHT SENSITIVITY

Fixed SpeakGR preservation coefficients substantially reduce WikiText FKL while retaining effective retrieval, but the coefficient affects the retrieval–preservation balance. Adaptive weighting is presented as an alternative control strategy, not a universally superior choice.

F FIXED-WEIGHT SENSITIVITY

  • Different fixed coefficients yield different retrieval–preservation operating points, and appropriately chosen fixed weights can achieve some of the strongest results.Source: F FIXED-WEIGHT SENSITIVITY · p. 18
  • Adaptive SpeakGR adjusts λ online to observed preservation KL, offering an alternative when the suitable preservation strength is unclear or may vary during training.The authors do not claim it is intrinsically superior to a well-chosen fixed weight.Source: F FIXED-WEIGHT SENSITIVITY · p. 18
  • WikiText FKL falls from 8.103 for retrieval-only SFT to 1.435–1.532 across non-zero SpeakGR coefficients λ ∈ {0.1, 0.5, 1.0, 1.5}, while retrieval remains effective.The sensitivity study indicates that the preservation effect does not depend on one carefully tuned fixed coefficient.Source: F FIXED-WEIGHT SENSITIVITY · pp. 17–18, Table 7 · p. 18
Loading 2609.35430v1…