RQ-VAE Recommender: semantic ID tokenization plus generative retrieval in one PyTorch repo
[Pytorch] Generative retrieval model using semantic IDs from "Recommender Systems with Generative Retrieval"
At a glance
- What is it?
- This repository implements the two-stage pipeline from Recommender Systems with Generative Retrieval: an RQ-VAE turns catalog items into short tuples of discrete semantic IDs, and a T5 encoder-decoder generates the IDs of likely next items. It is a research harness, not a serving stack, and the MovieLens path stops at tokenization.
- Who is it for?
- Adopt it if you are a recommender-systems researcher who wants the semantic-ID-to-next-item-generation interface in one inspectable PyTorch codebase, and you are working on Amazon Beauty, Sports or Toys. Do not adopt it if you need a serving system, a MovieLens retrieval baseline, or a library of interchangeable quantization layers.
- Can I use it commercially?
- Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 16 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The gap this fills between a quantization layer and a recommender framework
General-purpose vector-quantization libraries give you a reusable neural layer. Broad recommendation frameworks give you many model families, most of them unrelated to this one. The README states the project's position plainly: it focuses on the interface between residual semantic-ID learning and next-item generation. That interface is where the awkward work lives. You need sequential datasets prepared in a consistent order, hierarchical identifiers that stay stable across training runs, a rule for what happens when two catalog items land on the same ID tuple, a sequence model whose output space is restricted to valid IDs, and a ranked evaluation that accounts for both hit rate and NDCG. Each of those is a separate failure point, and each is usually written from scratch by whoever is reproducing the paper.
The intended audience is narrow on purpose. Researchers studying semantic tokenization, vector quantization, cold-item representations and generative retrieval will find the whole path in one place. Someone who wants a production recommender, or who wants to swap in a different retrieval backbone, is outside the target. The repository is MIT licensed and has no homepage, so the README and the docs/ directory are the entry points.
Two stages, three codebooks, and a collision rule
Stage one is the RQ-VAE. It maps an item, represented by an input vector, into a tuple of discrete codes. The quick API example in the README builds an RqVae with input_dim=8, embed_dim=4, hidden_dims=[16, 8], codebook_size=16, n_layers=3 and n_cat_features=0, then calls get_semantic_ids on a batch of five items and asserts the result has shape (5, 3). The n_layers value is what produces the three-part tuple: residual quantization encodes the item, then encodes the residual left over from the first codebook, then the residual from the second. With codebook_size=16 and three layers the ID space is 16 cubed, which is enough for a small catalog and not enough for a real one, but the example is explicitly a shape check against an untrained model.
The README lists K-means initialization for the codebooks and a choice between Euclidean and cosine assignment. Gradient flow through the discrete selection is handled by one of three estimators: Gumbel-Softmax, straight-through, or the rotation trick. QuantizeForwardMode.STE in the example selects straight-through. After tokenization, the second stage is a T5 encoder-decoder that consumes a user's semantic-ID history and generates the IDs of likely next items. Generation is constrained to corpus prefixes, so the decoder cannot emit an ID tuple that no catalog item owns. Collisions, where two items share a tuple, are disambiguated rather than ignored. The repository was used for the public-dataset experiments in a paper titled Stop Treating Collisions Equally, which is consistent with the collision handling being a first-class concern rather than an afterthought.
Installation and the gin files that actually control training
Python 3.10 or newer is required. The README notes that a CUDA or MPS accelerator is recommended for full training runs but is not needed for the unit tests, which run on CPU without downloading datasets or model weights. Installation is a virtual environment plus an editable install:
python -m venv .venv source .venv/bin/activate python -m pip install --upgrade pip python -m pip install -e .
Development and test tools come from the test extra, followed by pytest. A requirements.txt file is kept for environments that use a requirements-based workflow.
Training arguments are managed with gin-config. The train functions in train_rqvae.py and train_decoder.py are configurable, so model dimensions, datasets, optimization settings, evaluation intervals, logging and output paths are recorded in a .gin file rather than passed as flags. The Amazon flow is two commands:
rqvae-train configs/rqvae_amazon.gin rqvae-train-decoder configs/decoder_amazon.gin
The equivalent source-tree invocations are python train_rqvae.py configs/rqvae_amazon.gin and python train_decoder.py configs/decoder_amazon.gin. One detail in the README is worth repeating because it is easy to miss: configs/decoder_amazon.gin currently points to the example Beauty checkpoint in trained_models/rqvae_amazon_beauty/, and train.dataset_split and train.pretrained_rqvae_path must be changed together when using another Amazon split. Changing only one of them is the obvious way to train a decoder against the wrong tokenizer.
Where the pipeline stops: dataset support is not uniform across stages
The README carries a table that separates RQ-VAE tokenization from retrieval-model training. Amazon Reviews (Beauty, Sports, Toys) supports both. MovieLens 1M and MovieLens 32M support tokenization only, with retrieval-model training marked as not currently exposed. That is the single most important constraint for anyone planning an experiment, because it means the checked-in MovieLens configuration trains the semantic-ID tokenizer and nothing else:
rqvae-train configs/rqvae_ml32m.gin
If your research question is about generative retrieval on MovieLens, this repository does not currently give you the second stage for that data. You would be training a tokenizer and then writing the retrieval half yourself, which is most of the work the repository exists to save you.
The second constraint is scale. The README describes the full configurations as research-scale workloads and states that runtime and memory needs vary with the dataset, batch size, text-embedding model and accelerator. It does not publish a runtime or memory figure, and I have not run them, so treat any planning estimate as yours to produce. The docs/reproducibility.md file is described as covering what to record and how reviewers can perform a fast objective check, which is the right place to look before committing to a full run.
What a general vector-quantization library would give you instead
The obvious alternative for stage one is a standalone residual-quantization implementation from a general-purpose library. Those libraries expose the quantizer as a module you drop into your own model, with no opinion about catalogs, semantic IDs, collisions or retrieval. The difference in approach matters more than the difference in code volume. Here, the quantizer is not a component you import in isolation; it is wired to an item representation, a semantic-ID extraction method, a collision disambiguation step, and a downstream T5 decoder whose output is constrained to the ID corpus. A general library leaves all of that to you, which is fine if you already have a retrieval stack and only need the codes, and worse if what you actually want is the paper's end-to-end pipeline.
The trade runs the other way too. Because this repository owns the whole path, replacing the retrieval backbone, the item encoder or the evaluation means working inside its conventions and its gin configuration rather than composing independent pieces. The README is candid that this is a deliberate scope choice: compact implementation of the complete experimental pipeline, not a library of interchangeable parts. If your work is about the quantizer alone, the narrower dependency is the better one.
Maintenance, releases and the MIT licence
The repository is not archived and the most recent push and release both carry a timestamp of 2026-08-31, for version v1.0.2. That follows v1.0.0 and v1.0.1 on 2026-08-24, so the 1.0 line moved quickly across a single week. There is a tests workflow badge and a pytest suite that runs on CPU without downloading datasets or weights, which means a contributor can check a change without a GPU or a network fetch. That is a real maintenance property, not a marketing one: it lowers the cost of verifying a patch.
The licence is MIT, which permits commercial use and modification provided the copyright notice and permission notice are retained. That is the extent of what the repository states, and I am not giving legal advice; if you are incorporating this into a product, read the LICENSE file and the terms of the datasets and the Hugging Face checkpoint separately, because the MIT grant covers the code and not necessarily the weights or the data. Upgrade cost is mostly the cost of re-reading the gin files: because training parameters live in configuration rather than in code defaults, a version bump can change what a config means, and the README's warning about keeping train.dataset_split and train.pretrained_rqvae_path in sync is the kind of coupling that a release note would need to flag. The citation block asks users to cite version 1.0.1 via its Zenodo DOI, while the current release is 1.0.2, so check CITATION.cff for the version you should reference.
What to check before you build on it
Three things are worth confirming in the repository itself rather than from this description. First, open configs/decoder_amazon.gin and confirm the pretrained_rqvae_path and dataset_split pair for the split you intend to train on; the README says the checked-in file points at the Beauty checkpoint, and a mismatch there produces a decoder trained against the wrong ID space without an obvious error. Second, read docs/reproducibility.md, since the README directs reviewers there for what to record and how to run a fast objective check, and that is the closest thing to a verification recipe the project publishes. Third, if your data is MovieLens, decide now whether you are prepared to write the retrieval stage yourself, because the README's support table says the repository will not do it for you.
The published Amazon Beauty RQ-VAE checkpoint on Hugging Face is the fastest way to see whether the semantic IDs look reasonable on your own catalog items before you spend accelerator time on a full tokenizer run. The unit tests, which need neither datasets nor weights, are the cheapest way to confirm the install is sound. Neither of those substitutes for reading the gin files, which is where the pipeline's actual behaviour is defined.
Editorial conclusion
Adopt it if you are a recommender-systems researcher who wants the semantic-ID-to-next-item-generation interface in one inspectable PyTorch codebase, and you are working on Amazon Beauty, Sports or Toys. Do not adopt it if you need a serving system, a MovieLens retrieval baseline, or a library of interchangeable quantization layers. Before committing, verify that configs/decoder_amazon.gin still points at the Beauty checkpoint in trained_models/rqvae_amazon_beauty/, and check whether the MovieLens 1M and 32M retrieval gap described in the README has changed in the current release.
Community notes