RAT: Retrieval-Augmented Transformer for Click-Through Rate Prediction

This repository is the official PyTorch implementation of our WWW 2024 short paper.

Catalogue

• 1. Getting Started
• 2. Train/Test
• 3. Trained Models
• 4. Results
• 5. Citation

Getting Started

1. Clone this repository:

git clone https://github.com/YushenLi807/RAT.git
cd RAT

2. Install the dependencies:

• cuda 11.7
• python 3.7.8
• pytorch 1.13.1
• numpy 1.21.6
• h5py 2.10.0

pip install -r requirements.txt

3. Download Datasets: Tmall_002 is the full dataset of Tmall. You can download them from Baiduyun disk.

Dataset	Link
Movielenslatest	Baidu disk
Movielenslatest_10fold_retrieval	Baidu disk
KKBox	Baidu disk
KKBox_10fold_retrieval	Baidu disk
Tmall	Baidu disk
Tmall_002	Baidu disk
Tmall_002_retrieval	Baidu disk

Train/Test

RAT_m0:RAT_JM.

RAT_m1:RAT_CE.

RAT_m2:The default RAT.

RAT_m3:RAT_PA.

To train/test RAT on Movielenslatest_10fold_retrieval:

python run_expid.py --config ./configs/RAT_m2/movielenslatest_x1 --expid RAT_m2_movielenslatest_x1_10fold_retrieval --gpu 0

To train/test RAT on KKBox_10fold_retrieval:

python run_expid.py --config ./configs/RAT_m2/kkbox_x1 --expid  RAT_m2_kkbox_x1_10fold_retrieval --gpu 0

To train/test RAT on Tmall_002_retrieval:

python run_expid.py --config ./configs/RAT_m2/tmall_x1_002 --expid  RAT_m2_tmall_x1_002_retrieval --gpu 0

Trained Models

We provide trained RAT checkpoints. You can download them from Baiduyun disk.

Dataset	Link
Movielenslatest_10fold_retrieval	Baidu disk
KKBox_10fold_retrieval	Baidu disk
Tmall_002_retrieval	Baidu disk

Results

For this repository, the expected performance is:

Model	ML-Tag		KKBox		Tmall
	AUC	Logloss	AUC	Logloss	AUC	Logloss
RATJM	0.9667	0.2003	0.8415	0.4917	0.9581	0.3110
RATCE	0.9736	0.1731	0.8483	0.4831	0.9575	0.3182
RATPA	0.9777	0.1557	0.8484	0.4828	0.9582	0.3177
RAT	0.9809	0.1421	0.8500	0.4812	0.9589	0.3091

Citation

If you find this repository useful, please consider citing our work:

@inproceedings{li2024rat,
  title={RAT: Retrieval-augmented Transformer for Click-through Rate Prediction},
  author={Yushen Li and Jinpeng Wang and Tao Dai and Jieming Zhu and Jun Yuan and Rui Zhang and Shu-Tao Xia},
  booktitle={Companion Proceedings of the ACM Web Conference 2024},
  year={2024}
}

Acknowledgements

Our code is based on the implementation of FuxiCTR and BARS.