Torch-RecHub Documentation

Torch-RecHub provides a collection of PyTorch-based recommendation models organized by their role in the recommendation pipeline. The library covers four primary stages:

1. Ranking Models: Predict click-through rates (CTR) or user preference scores during fine ranking.
2. Matching Models: Retrieve candidates from large item collections during the candidate generation stage.
3. Multi-Task Models: Jointly optimize multiple related tasks to improve model generalization.
4. Generative Recommendation Models: Use generative architectures to produce personalized recommendations.

Torch-RecHub is a lightweight, efficient, and easy-to-use PyTorch-based framework for building industrial-grade recommendation systems. It allows users to implement complex recommendation tasks with minimal code (as little as 10 lines) and provides over 30 out-of-the-box models covering retrieval (matching), ranking, multi-task learning, and generative recommendation.

Key capabilities include:

• Hardware Support: CPU, NVIDIA CUDA, AMD ROCm, and Huawei Ascend NPU.
• Model Library: 30+ algorithms including DSSM, WideDeep, DeepFM, MMoE, etc.
• Deployment: One-click ONNX export for production environments.
• Data Processing: Support for PySpark-based processing and Parquet via PyArrow.
• Experiment Tracking: Integrated support for WandB, SwanLab, and TensorBoardX.
• Modular Design: Easy to extend with new models, datasets, and metrics.

Torch-RecHub provides specialized models categorized by recommendation stage:

Ranking Models (CTR Prediction)

Used for fine-grained scoring of candidates.

• DeepFM: Combines FM and deep networks.
• Wide&Deep: Combines memorization and generalization.
• DCN / DCNv2: Uses explicit feature crossing.
• DIN: Target-aware attention over user history.
• DIEN: Models interest evolution.
• BST: Transformer-based sequence ranking.

Retrieval Models (Matching)

Used for candidate generation/two-tower architectures.

• DSSM: Classic two-tower semantic matching.
• YoutubeDNN: YouTube-style deep retrieval.
• MIND: Multi-interest retrieval using capsules.

Multi-Task Models

Used for joint optimization of multiple objectives.

• MMOE: Multi-gate mixture-of-experts.
• PLE: Progressive layered extraction.

Torch-RecHub supports over 30 mainstream recommendation models categorized into four main types:

1. Ranking Models: Focused on CTR prediction (e.g., DeepFM, Wide&Deep, DCN, DCN-v2, DIN, DIEN, BST, AFM, AutoInt, FiBiNET, DeepFFM, EDCN).
2. Matching Models: Two-tower or retrieval models (e.g., DSSM, YoutubeDNN, YoutubeSBC, MIND, SINE, GRU4Rec, SASRec, NARM, STAMP, ComiRec).
3. Multi-Task Models: Models designed for multiple objectives (e.g., ESMM, MMoE, PLE, AITM, SharedBottom).
4. Generative Recommendation: Cutting-edge generative approaches (e.g., HSTU, HLLM, TIGER).

The repository is organized into core library components, example scripts, and configuration files:

• torch_rechub/: The core library.
  • basic/: Fundamental building blocks like activation.py, features.py, layers.py, loss_func.py, and metric.py.
  • models/: Implementation of recommendation algorithms categorized by task: matching/ (e.g., DSSM, MIND), ranking/ (e.g., WideDeep, DeepFM), and multi_task/ (e.g., MMoE, ESMM).
  • trainers/: Specialized training frameworks for different tasks: ctr_trainer.py, match_trainer.py, and mtl_trainer.py.
  • utils/: Utilities for data.py, match.py, mtl.py, and onnx_export.py.
• examples/: Task-specific scripts for matching/, ranking/, and generative/ (e.g., HSTU, HLLM).
• config/: Configuration files for experiments.
• tutorials/: Jupyter notebooks for guided learning.

PLE is a multi-task learning (MTL) model designed to solve the seesaw phenomenon, where optimizing one task negatively impacts another. Unlike MMOE, which uses only shared experts, PLE utilizes Customized Gate Control (CGC) to assign both task-specific experts and shared experts. This allows for better task separation and adaptive combination of features through gate networks.

Key Components:

• Task-Specific Experts: Dedicated networks for each individual task.
• Shared Experts: Networks shared across all tasks.
• Customized Gate (CGC): Combines task-specific and shared experts for each task.
• Multi-Level Support: Allows stacking multiple CGC levels for progressive feature extraction.
• Task Towers: Independent prediction towers for each task.

DeepFM (Deep Factorization Machine) is a ranking model that combines Factorization Machines (FM) with Deep Neural Networks (DNN). It is designed to capture both low-order (pairwise) and high-order (nonlinear) feature interactions simultaneously without manual feature engineering.

Key Architecture Components:

• FM part: Captures pairwise feature interaction patterns.
• Deep part: Captures higher-order nonlinear interactions via multi-layer fully connected networks.
• Shared embeddings: Both branches share the same bottom embedding layer to reduce parameter count.

Best Use Cases:

• Click-through rate (CTR) prediction.
• Ad recommendation ranking.
• Scenarios requiring both low-order and high-order feature interactions.

DSSM (Deep Structured Semantic Model) is a two-tower architecture used primarily in the retrieval stage of recommendation systems. It consists of a User Tower and an Item Tower that map features into a shared vector space. Matching scores are computed using cosine similarity or dot products.

Key Use Cases:

• Fast filtering of large candidate sets via vector search.
• Search relevance matching.
• Online real-time services (supports offline precomputation of user/item vectors).

MIND (Multi-Interest Network with Dynamic Routing) is a multi-interest retrieval model designed for the retrieval stage of recommendation systems, particularly where users have diverse interests (e.g., e-commerce).

Unlike single-vector models like DSSM, MIND uses a Capsule Network with dynamic routing to extract multiple interest vectors from a user's behavior sequence. This results in a user representation with shape [batch_size, interest_num, embed_dim].

Key Architecture Components:

• Embedding Layer: Encodes user attributes and historical behavior.
• Capsule Network: Extracts multiple interest vectors via dynamic routing.
• Training: Uses list-wise training (Softmax) via mode=2 in MatchTrainer.