hexiangnan / attentional_factorization_machine

TensorFlow Implementation of Attentional Factorization Machine
407 stars 156 forks source link
deep-learning factorization-machines neural-factorization-machines recommender-system

attentional_factorization_machine

This is our implementation for the paper:

Jun Xiao, Hao Ye, Xiangnan He, Hanwang Zhang, Fei Wu and Tat-Seng Chua (2017). Attentional Factorization Machines: Learning the Weight of Feature Interactions via Attention Networks IJCAI, Melbourne, Australia, August 19-25, 2017.

We have additionally released our TensorFlow implementation of Factorization Machines under our proposed neural network framework.

Please cite our IJCAI'17 paper if you use our codes. Thanks!

Author: Xiangnan He (xiangnanhe@gmail.com) and Hao Ye (haoyev@gmail.com)

Environments

Quick Example with Optimal parameters

Use the following command to train the model with the optimal parameters:

# step into the code folder
cd code
# train FM model and save as pretrain file
python FM.py --dataset ml-tag --epoch 100 --pretrain -1 --batch_size 4096 --hidden_factor 256 --lr 0.01 --keep 0.7
# train AFM model using the pretrained weights from FM
python AFM.py --dataset ml-tag --epoch 100 --pretrain 1 --batch_size 4096 --hidden_factor [8,256] --keep [1.0,0.5] --lamda_attention 2.0 --lr 0.1

The instruction of commands has been clearly stated in the codes (see the parse_args function).

The current implementation supports regression classification, which optimizes RMSE.

Performance Comparison

Parameters

For the sake of a quick demonstration for the improvement of our AFM model compared to original FM, we set the dimension of the embedding factor to be 16 (instead of 256 in our paper), and epoch as 20.

Train

Step into the code folder and train FM and AFM as follows. This will start to train our AFM model on the dataset frappe based on the pretrained model of FM. The parameters have been initialized optimally according to our experiments. It will loop 20 epochs and print the best epoch depending on the validation result.

# step into the code folder
cd code
# train FM model with optimal parameters
python FM.py --dataset ml-tag --epoch 20 --pretrain -1 --batch_size 4096 --hidden_factor 16 --lr 0.01 --keep 0.7
# train AFM model with optimal parameters
python AFM.py --dataset ml-tag --epoch 20 --pretrain 1 --batch_size 4096 --hidden_factor [16,16] --keep [1.0,0.5] --lamda_attention 100.0 --lr 0.1

After the trainning processes finish, the trained models will be saved into the pretrain folder, which should be like this:

Last Update Date: Aug 2, 2017