This repository contains the PyTorch implementation for the paper "Uncertainty-aware Generative Recommendation", accepted by KDD 2026.
We identify a pervasive uncertainty-blindness issue in generative recommenders: they treat every prediction as equally reliable and ignore the predictive uncertainty behind each beam. This leads to two failures — during training, all target tokens are optimized uniformly regardless of how confident the model actually is; during inference, items are ranked purely by generation likelihood, leaving no signal that reflects how certain the model truly is about each predicted result.
We address this by turning uncertainty itself into an explicit, learnable signal: UGR augments the vocabulary with confidence tokens and aligns the model via an uncertainty-aware reinforcement learning objective, so that the model can jointly generate a recommendation together with how confident it is in that prediction.
UGR/
├── data/
│ ├── _1/ # Raw data filtering & splitting
│ ├── _2/ # Text embedding generation
│ └── _3/ # RQ-VAE training & SID generation
├── train/
│ ├── SFT/ # Supervised fine-tuning
│ └── RL/ # Uncertainty-aware reinforcement learning
└── eval/ # Evaluation with confidence scoring
Download the raw datasets (e.g., Amazon18) from the Amazon Datasets repository.
bash ./data/_1/amazon18_data_process.shbash ./data/_2/amazon_text2emb.shbash ./data/_3/rqvae.sh
python ./data/_3/generate_indices.pybash ./train/SFT/run_train_SFT.shbash ./train/RL/run_train_RL.shbash ./eval/run_eval.shParts of our implementation build upon, or are inspired by, the following open-source projects. We sincerely thank their authors and the broader community for making their work publicly available:
If you find this repository useful, please consider giving it a star and citing our paper:
@article{fan2026uncertainty,
title={Uncertainty-aware Generative Recommendation},
author={Fan, Chenxiao and Gao, Chongming and Gong, Yaxin and Liu, Haoyan and Feng, Fuli and He, Xiangnan},
journal={arXiv preprint arXiv:2602.11719},
year={2026}
}For questions or issues, please open a GitHub issue or contact the first author.