EARN

This is the pytorch implementation of our paper:

EARN: Efficient Inference Acceleration for LLM-based Generative Recommendation by Register Tokens (KDD 2025)

Below, we take LC-Rec as an example. For HSTU, please refer to HSTU/README.md.

Environment

• Python 3.10.15
• pytorch 2.1.2
• transformers 4.41.2

Train

First, replace the parameters in LC-Rec/script/finetune_register.sh with your own parameters, such as OUTPUT_DIR, LOG_DIR, etc.

OUTPUT_DIR=YOUR_OUTPUT_DIR
LOG_DIR=YOUR_LOG_DIR

Then, run the following command to train EARN.

cd LC-Rec
bash script/finetune_register.sh

Inference

First, replace the parameters in LC-Rec/script/test_register.sh with your own parameters, such as MODEL_PATH, LOG_DIR, etc.

MODEL_PATH=YOUR_MODEL_PATH
LOG_DIR=YOUR_LOG_DIR

Then, run the following command to test the performance.

cd LC-Rec
bash script/test_register.sh