Can Generative Retrievers Learn Semantic IDs Without Forgetting How to Speak?
Junchen Fu, Kleomenis Katevas, Vandana Rajan, Sofía Celi, Hamed Haddadi
cs.IR
2026-09-28
SpeakGR keeps generative retrievers speaking: on-policy forward KL cuts WikiText-2 drift 81.3–93.8% vs SFT-only, with retrieval intact across three LLMs.
Generative retrieval turns document lookup into autoregressive generation of a short semantic ID. The model grows its vocabulary with SID tokens and emits an identifier instead of scoring an inverted index. Interactive search wants that same network to retrieve, then answer, follow up, or clarify.
Standard GR training never constrains the text-token policy. On Qwen3-0.6B with MS MARCO, WikiText-2 perplexity jumps from 31.69 to 6956 in 25 document-to-SID updates, while Recall@10 is still climbing. Across six backbone-dataset pairs, retrieval-only SFT leaves WikiText forward KL between 5.545 and 13.005. Prior GR papers mostly report recall. Language collapse is the missing axis. Until speaking survives SID training, a single model that both retrieves and replies cannot even be tested.
SpeakGR trains one student with two losses. Retrieval is level-restricted cross-entropy over SID tokens: each position is renormalized inside that level's codebook, identifier length is fixed, and EOS is teacher-forced but left out of the retrieval loss. Preservation follows on-policy distillation. The student masks SID logits, samples a short text continuation, and a frozen copy of the original model scores the same prefixes. Forward KL over the original text vocabulary pulls the student toward tokens the base model still assigns mass to. Prefixes come from the student, so the regularizer tracks language states that SID training actually visits.
The static variant adds the two token-normalized losses with equal weight. Adaptive SpeakGR only changes the coefficient. It tracks an EMA of the preservation KL and multiplies λ up or down, with a deadband and clips. The default target is 1 nat per token. Training still uses three stages: 2000 document-to-SID steps, 8000 pseudo-query steps, 2000 supervised query steps. SIDs are 5-token RQ-VAE codes from Qwen3-Embedding, 256 entries per level. The preservation pool has 4096 prompts, half WikiText-2 prefixes and half instruction-style. The WikiText-2 test split is eval-only.
Backbones are Qwen3-0.6B, Qwen3-1.7B, and Gemma-3-1B-IT. Corpora are processed MS MARCO (319872 docs) and Natural Questions (109739). Language metrics average 16384 WikiText-2 positions. Decoding is trie-constrained.
Versus SFT-only, SpeakGR cuts WikiText forward KL by 81.3% to 93.8% on MS MARCO and 81.2% to 85.2% on NQ. Perplexity falls from five or seven digits back to two. Top-1 agreement with the base model returns to roughly 46% to 58% from near zero. SID probability mass, which hit 26.37% under SFT on Gemma plus MS MARCO, drops to around 1% or less. Retrieval mostly holds. The Qwen models lose more recall on NQ under fixed SpeakGR; Gemma slightly gains on both datasets. Adaptive SpeakGR beats fixed SpeakGR on retrieval in five of six settings and stays far below SFT on drift.
On Qwen3-0.6B with MS MARCO:
| Method | R@10 | FKL | PPL | Top-1 |
| Base | n/a | 0 | 31.69 | 100% |
| SFT-only | 0.6386 | 8.103 | 53967 | 0.13% |
| SpeakGR | 0.6349 | 1.519 | 88.48 | 46.15% |
| Adaptive | 0.6498 | 1.728 | 105.86 | 43.24% |
Offline replay in the same setting still has FKL 8.672 and R@10 0.5903. ORBIT-style merge depends sharply on its threshold: τ=0.10 keeps R@10 at 0.6584 with FKL 7.391; τ=0.007 gets FKL 0.273 and R@10 0.1980. SpeakGR lands where retrieval still works and the text distribution does not explode. SID leakage explains only 0.58% of SFT's extra full-vocab NLL; the damage sits inside the renormalized text distribution. Fixed λ in {0.1, 0.5, 1.0, 1.5} all cut FKL from 8.103 to about 1.43 to 1.53. An overly strict adaptive target (ε=0.25) hurts both axes (R@10 0.5458, FKL 2.507). The best moderate point is R@10 0.6584 at ε=1.5 with FKL 1.566.
If one LLM is supposed to emit SIDs and then talk, this paper measures a failure that retrieval papers usually skip: SID indexing can wreck next-token behavior in 25 steps. SpeakGR is a practical add-on for DSI-style pipelines, retrieval CE plus forward KL on student-sampled prefixes. Adaptive weighting avoids committing to one λ for the whole run. A well-chosen fixed weight is already strong; the paper does not claim adaptation wins that comparison.
This is not a retrieval SOTA paper. Lower WikiText KL means the next-token distribution is closer to the pretrained model. Math, code, and multi-turn clarification are untested.
Language eval is almost entirely WikiText-2 distributional drift. There is no retrieve-then-generate end-to-end task, and no math or coding probe. Backbones top out at 1.7B. Corpora are MS MARCO and NQ subsets. Most comparisons use one seed and one RQ-VAE family. Replay and ORBIT are single operating points. Adaptive SpeakGR still lacks a static run matched on average λ. Figure 1's PPL of 6956 is at step 25 of Stage 1; the main table's finished SFT PPL is 53967 on the same pair, so early collapse and final collapse are different snapshots.
Colliding SID buckets are reranked with a frozen Qwen3-Embedding-0.6B encoder. The preserved speaker still hands off to a separate embedder inside the bucket.