Skip to content

Repository files navigation

RecDistillery

RecDistillery is a modular framework for Knowledge Distillation in Recommender Systems, built around a simple idea:

Great recommendation models, like fine spirits, should be distilled -- not diluted.

Large teacher models contain rich collaborative knowledge, complex ranking behaviors, structural information, and latent relational patterns. RecDistillery provides the tools to extract, refine, and transfer this knowledge into lightweight student models through unified and reproducible distillation pipelines.

The current framework supports:

  • teacher import from .pth or .pt embedding checkpoints or .json predictions, exported into .teacher;
  • adapter-backed teacher/student training from RecBole, Elliot, and Lenskit model definitions through a torch-based training pipeline;
  • 6 distillers: DE, RRD, DE+RRD, HTD, FTD, UnKD.

Imported recommendation model definitions currently come from RecBole, Elliot, and Lenskit. RecDistill trains only this subset exposed through PyTorch adapters:

  • RecBole:
    • BPRMF (aliases: BPR, BPRMF)
    • LINE (aliases: LINE)
    • LGCN (aliases: LGCN, LightGCN)
    • NGCF (aliases: NGCF)
    • DGCF (aliases: DGCF)
    • SGL (aliases: SGL)
    • SPECTRALCF (aliases: SpectralCF, SPECTRALCF)
    • NMF (aliases: NMF, NeuMF)
  • Elliot:
    • BPRMF (aliases: BPR, BPRMF)
    • LGCN (aliases: LGCN, LightGCN)
    • NGCF (aliases: NGCF)
    • DGCF (aliases: DGCF)
    • SGL (aliases: SGL)
    • ULTRAGCN (aliases: UltraGCN, ULTRAGCN)
    • NMF (aliases: NMF, NeuMF)
  • Lenskit:
    • BPRMF (aliases: BPRMF)
    • LGCN (aliases: LGCN, LightGCN)

The dataset preprocessing and standardized data management are based on DataRec:

Table of Contents


Architecture

The framework core lives under recdistill/.

recdistill/
  data/                  Interaction datasets and DataRec-compatible loading
  distillers/            DE, RRD, UnKD, HTD, FTD, CompositeDistiller (DE+RRD)
  teachers/              Teacher adapters, import/export, native .teacher format
  samplers/              Negative and distillation samplers
  trainers/              Training loops and optimization helpers
  framework_backbone.py  Framework adapter and model registry
  factories.py           Builders for student models, distillers, and backbone aliases
  config_integration.py  RecDistill config composition and loading
  experiment_runner.py   RecDistillExperimentRunner
  native_runner.py       Native teacher/student model training runner
  checkpointing.py       Teacher/student checkpoint management
  evaluation.py          Shared top-k metrics used by student evaluation and training runners
  model_validation.py    Compatibility and request validation
  supported_models.py    Torch-compatible model registry
  training.py            Shared training utilities
  tracking.py            Experiment logging and metadata
  paths.py               Canonical path resolution helpers
  registry.py            Canonical aliases for models and distillers

Repository Setup

Prerequisites

  • Python 3.12+
  • Conda
  • CUDA 11.8+ (recommended)

Optional:

  • Apple Silicon MPS support

Clone The Repository

git clone https://github.com/sisinflab/RecDistillery.git
cd RecDistillery

Configure The Environment

CUDA:

bash setup/setup_environment.sh cuda

Apple Silicon:

bash setup/setup_environment.sh mps

CPU only:

bash setup/setup_environment.sh

After installation, activate the environment with:

source setup/activate_env.sh

This activates the default distillation environment and sets the project PYTHONPATH.

Download and Prepare Datasets

RecDistillery provides, as an example, ready-to-use data preparation scripts for three datasets:

  • Amazon-CD
  • BookCrossing
  • CiteULike

The preprocessing utilities are based on DataRec, so additional datasets can also be used if they are available in that ecosystem. See DataReHub for more dataset references and formats: https://www.datarechub.com/.

Run:

bash setup/setup_directories.sh
bash setup/setup_datasets.sh

The split files generated are stored under:

data/<dataset>/train.tsv
data/<dataset>/val.tsv
data/<dataset>/test.tsv

DataRec-compatible loading lives in recdistill/data/datarec_loader.py.

Dataset sources:


Workflow

Compatibility Check

Before launching training, you can list the imported model definitions that are torch-compatible, plus the adapter-backed subset already wired to the current RecDistill unified loop:

python scripts/recdistill/welcome.py

Non-Torch teachers can still be used after conversion to .teacher with import_teacher.py.

Adding New Model Adapters

If a model appears in the imported framework definitions but is not listed as adapter-backed, RecDistill cannot train it in the unified PyTorch loop yet. To make it trainable, add a backbone adapter.

  1. Add or extend the framework adapter in recdistill/framework_backbone.py. The adapter must be a torch.nn.Module and expose the same surface used by the distillers:
forward(users, pos_items, neg_items) -> FrameworkBatchOutput
score_items(users, items) -> torch.Tensor
user_embeddings() -> torch.Tensor
item_embeddings() -> torch.Tensor
  1. Wire the model inside the relevant adapter class: RecBoleBackboneAdapter, ElliotBackboneAdapter, or LensKitBackboneAdapter.

  2. Register the model as trainable in recdistill/supported_models.py by adding a TrainableBackbone entry with framework, canonical model name, aliases, adapter name, and implementation path.

  3. Add the model alias in recdistill/registry.py if the model should be resolved from multiple names.

  4. Add model configs for teacher and/or student:

config/teacher/<framework>/<model>.yaml
config/student/<framework>/<model>.yaml
  1. Run the compatibility check again:
python scripts/recdistill/welcome.py --verbose

Teacher Import

RecDistillery treats the .teacher file as the framework-neutral teacher format. Teacher import is handled by a small set of generic adapters that convert external artifacts into either user/item embeddings or precomputed top-k rankings.

The official path is scripts/recdistill/import_teacher.py. The import script currently registers only:

  • CheckpointAdapter: generic torch checkpoints containing a serialized teacher state, embeddings, scores, or top-k tensors.
  • PredictionsJsonAdapter: JSON (or .tsv or .csv) prediction exports with user, item, and optional score/rating/rank fields.
  • RecBolePthAdapter: .pth checkpoints containing user and item embedding tensors.

List the active adapters with:

python scripts/recdistill/import_teacher.py --list-adapters

Import a generic checkpoint:

python scripts/recdistill/import_teacher.py \
  --input path/to/teacher_checkpoint.pt \
  --format checkpoint \
  --framework external \
  --model-name ExternalTeacher \
  --dataset citeulike \
  --embedding-dim 200

Import precomputed recommendation lists:

python scripts/recdistill/import_teacher.py \
  --input path/to/predictions.tsv \
  --dataset amazon_cd \
  --model-name ItemKNN \
  --framework external

Import a .pth checkpoint with user/item embeddings (e.g., from RecBole):

python scripts/recdistill/import_teacher.py \
  --input path/to/model.pth \
  --format recbole_pth \
  --framework recbole \
  --model-name BPRMF \
  --dataset citeulike

External prediction exports (.tsv, .csv, or .json) can contain raw dataset IDs or raw user tokens. When --dataset is provided, import_teacher.py automatically maps these IDs into canonical 0-based dataset indices ($0 \dots N-1, 0 \dots M-1$) across all splits (train, val, test), enabling evaluation with zero dropped interactions.

If an external checkpoint uses framework-internal indices (without raw dataset tokens), supply the framework ID mapping via public_to_local_user_id and public_to_local_item_id metadata, or specify --metadata id_space=dataset_integer if indices match the dataset splits.

Imported teachers are saved as tracked teacher runs:

results/teacher/<timestamp>_<external>_<model>_<dataset>_<experiment_id>/
|-- artifacts/<external>_<model>_<dataset>_<experiment_id>_best.teacher
|-- config/<external>_<model>_<dataset>_<experiment_id>.yaml
|-- logs/import_summary.json
`-- perf/

When using it in a RecDistill config, keep the teacher model field explicit even though it is not trained internally by RecDistillery:

distill_student:
  teacher:
    model: ItemKNN
    path: results/teacher/<run>/artifacts/<teacher>_best.teacher

Teacher Training

Teacher models can also be generated through the framework torch-based training pipeline. Train a native RecDistill teacher and save it as a .teacher artifact with:

python scripts/teacher_training/teacher_training.py \
  --framework recbole \
  --model BPRMF \
  --dataset citeulike

Alternatively, use a complete experiment config file:

python scripts/teacher_training/teacher_training.py \
  --config config/experiments/teacher/<experiment>.yaml

This script trains a teacher model, exports a .teacher file, and makes it available for later distillation or evaluation.

To check that the teacher is readable and compatible before using it in distillation, run this smoke test:

python scripts/recdistill/teacher_smoke.py \
  --teacher-framework <framework> \
  --teacher-model <model> \
  --dataset <dataset> \
  --embedding-dim <embedding_dim> \
  --top-k 20

For imported teachers, pass the artifact explicitly:

python scripts/recdistill/teacher_smoke.py \
  --teacher-path results/teacher/<run>/artifacts/<teacher>_best.teacher \
  --teacher-framework <framework> \
  --teacher-model <model> \
  --dataset <dataset> \
  --embedding-dim <embedding_dim>

Student Training (no distillation)

Train a student model without distillation with:

python scripts/student_training/student_training.py \
  --framework recbole \
  --backbone LGCN \
  --dataset citeulike \
  --distillation none

Or use a config file for the student training run:

python scripts/student_training/student_training.py \
  --config config/experiments/student/<experiment>.yaml

The output is a .student artifact that can be evaluated or later compared with distilled students.

Student Distillation

RecDistill exposes two main student training entry points:

  • scripts/recdistill/train_student.py for direct student training and distillation from a saved teacher.
  • scripts/recdistill/train_student_from_config.py for config-based, reproducible RecDistill experiments.

Example config-driven distillation:

python scripts/recdistill/train_student_from_config.py \
  --config config/experiments/recdistill/de_citeulike_001.yaml

When --config is omitted, pass an explicit teacher artifact path. Do not pass --teacher-model or --teacher-framework; those are reserved for complete configs and native teacher training.

python scripts/recdistill/train_student_from_config.py \
  --dataset citeulike \
  --teacher-path results/teacher/<run>/artifacts/<teacher>_best.teacher \
  --distiller rrd \
  --student-backbone BPRMF \
  --student-framework recbole

Direct distillation from CLI uses the same canonical teacher/student argument names:

python scripts/recdistill/train_student.py \
  --dataset citeulike \
  --teacher-framework recbole \
  --teacher-model BPRMF \
  --teacher-embedding-dim 200 \
  --student-framework recbole \
  --student-backbone LGCN \
  --student-embedding-dim 20 \
  --lambda-de 0.1

Distilled student training produces .distilled_student artifacts, containing model state, config, metadata, and history.

Experiment Launchers

Shell launchers are grouped by:

experiments/recdistill/      RecDistill launchers
experiments/baseline/        Teacher and baseline student launchers

RecDistill examples:

bash experiments/recdistill/train_distiller.sh <distiller> <dataset> [teacher_framework] [teacher_model|ALL] [student_framework] [student_backbone|SAME|ALL] [gpu] [dry_run]
bash experiments/recdistill/train_single_distiller.sh <distiller> <dataset> <teacher_framework> <teacher_model> [student_framework] [student_backbone|SAME] [gpu] [dry_run]

Baseline example:

bash experiments/baseline/teacher_training.sh [teacher_framework] [teacher_model] <dataset> [gpu]

Configs And Experiments

The canonical config package is config/.

config/
  dataset/            Dataset definitions
  teacher/            Teacher model defaults
  student/            Student model defaults
  distillation/       Distiller defaults
  optimization/       Optimization defaults
  runtime/            Runtime defaults
  evaluation/         Evaluation defaults
  composites/         Composable templates
  experiments/        Complete experiment configs

Complete experiment configs are under:

config/experiments/teacher/
config/experiments/student/
config/experiments/recdistill/

Manually planned experiments live there. On-the-fly generated configs are saved as run artifacts under results/<kind>/<run>/config/.


Evaluation

The evaluation in RecDistillery is done as a top-k ranking, typically @20. During student training, the pipeline calculates the top-k recommendations, compares them with the validation/test ground truth, calculates precision, recall, ndcg, and hr and finally saves the best artifact using val.ndcg as selection metric.

The formulas are implemented here: recdistill/evaluation.py. Then the metrics are averaged across evaluable users.

Run scripts/recdistill/evaluate_teacher.py to evaluate .teacher artifacts and scripts/recdistill/evaluate_students.py to evaluate .student or .distilled_student artifacts. Both scripts produce separate JSON/TSV files.

Teacher evaluation example:

python scripts/recdistill/evaluate_teacher.py \
  --teacher-path results/teacher/<run>/artifacts/<teacher>_best.teacher 

Plain student evaluation example:

python scripts/recdistill/evaluate_students.py \
  --student-path results/student/<run>/artifacts/<student>_best.student

Distilled student evaluation example:

python scripts/recdistill/evaluate_students.py \
  --student-path results/recdistill/<run>/artifacts/<recdistill>_best.distilled_student

Results structure

The setup script creates the base runtime directories used by the pipeline, including results/ and data/.

Results use exactly three top-level experiment kinds:

results/teacher/
results/student/
results/recdistill/

Each run is stored as:

results/<kind>/<timestamp>_<framework>_<model>_<dataset>_<experiment_id>/
|-- artifacts/
|-- config/
|-- logs/
`-- perf/

Best artifacts use the same identity as the run:

results/<kind>/<run>/artifacts/<framework>_<model>_<dataset>_<experiment_id>_best.<kind_ext>

For recdistill, the model label is <strategy>_<student_backbone>, for example:

results/recdistill/<timestamp>_elliot_DE_LGCN_citeulike_<experiment_id>/artifacts/elliot_DE_LGCN_citeulike_<experiment_id>_best.distilled_student

Bayesian runs additionally save logs/best_trial.json, logs/optuna_trials.json, logs/optuna_runtime_records.json, and config/<experiment>_best.yaml.


Timing Analysis

Timing analysis is implemented through tracked student runs. The same run that collects the final metrics also records the values needed to measure:

  • total training time,
  • average epoch time,
  • KD overhead,
  • per-epoch training and validation history.

Main Command

python3 scripts/recdistill/train_student_from_config.py \
  --config <config_file> \
  --track

Outputs

results/recdistill/<timestamp>_<student_framework>_<strategy>_<student_model>_<dataset>_<experiment_id>/
|-- artifacts/
|-- config/
|-- logs/run_recap.json
|-- logs/run_recap.tsv
`-- perf/

Releases

Packages

Contributors

Languages