DeepFM CTR Prediction Pipeline
A PyTorch-based implementation of the DeepFM model designed for CTR prediction on high-dimensional sparse data. Implements linear Factorization Machines (FM) alongside deep neural networks (DNN) with shared feature embeddings. Achieved published benchmarks (AUC ≈ 0.801, LogLoss ≈ 0.456) on the full Criteo dataset (~45M samples) and evaluated 6 architecture ablation experiments.
Problem Statement
Click-Through Rate (CTR) prediction is central to digital advertising and recommendation systems. Deep recommender systems must efficiently capture both low-order (linear interactions) and high-order (complex non-linear interactions) combinations from extremely sparse, high-dimensional categorical features. Training deep models on the massive Criteo dataset (~45M samples, 39 fields) from scratch presents significant memory overhead, class imbalance, and design choice trade-offs.
Solution
Reproduced the DeepFM architecture in PyTorch, combining a Factorization Machine (FM) component with a Deep Neural Network (DNN) sharing the same embedding layer. Designed an optimized offline preprocessing pipeline using Unix shuffling and sklearn OrdinalEncoding to handle the full dataset. Built a comprehensive ablation testing suite to run 6 distinct architectural configurations (evaluating dropout, pooling, batch norm, skip connections, network depth, and loss formulations).
Architecture
The model fuses a linear FM component (for 1st and 2nd-order feature interactions) and an MLP deep component (for higher-order non-linear combinations). The inputs are mapped through shared embedding lookup tables before being routed to both components in parallel, eliminating the need for manual feature combination engineering.
Key Features
- ▸PyTorch implementation of the dual-component DeepFM recommender architecture
- ▸Shared embedding layer mapping 39 sparse fields into dense vectors
- ▸Memory-efficient data pipeline handling 45 million rows via chunked processing and ordinal encoding
- ▸Automated Weights & Biases (W&B) logging for training, evaluation, and convergence tracking
- ▸Configurable ablation experiment runner to test network variations dynamically
- ▸Class imbalance analysis using specialized BCE with logits loss formulations
Challenges
- ⚡Handling massive dataset memory footprint on limited GPU/RAM environments by designing a generator-backed batch pipeline.
- ⚡Mitigating severe overfitting on sparse embeddings through precise dropout tuning and batch normalization.
- ⚡Balancing the contribution of low-order and high-order components during backpropagation under extreme class imbalance.
Results & Metrics
Successfully reproduced paper results on full Criteo dataset (AUC ≈ 0.801, LogLoss ≈ 0.456)
Optimized ablation config achieved peak AUC of 0.802 and F1-score of 0.56
Processed and trained on 45 million samples across 39 high-dimensional feature fields
Evaluated 6 distinct architectural ablation studies with full W&B logs
Lessons Learned
- 💡Shared embeddings between FM and deep components are highly sensitive to dropout, where lower dropout (0.1) performs better than heavy dropout (0.5).
- 💡Using BCEWithLogitsLoss provides superior numerical stability and helps offset gradient variance under class imbalance.
- 💡Ablation testing reveals that simple flattening of embeddings outperforms average pooling over fields for sparse CTR predictions.
Case Study Overview
Case Study: Reproducing DeepFM for CTR Prediction
Click-Through Rate (CTR) prediction is a foundational problem in recommendation engines and ad-tech. The goal is to estimate the probability that a user will click on a recommended item. Because features consist of high-dimensional categorical fields (like user IDs, ad categories, and location hashes), models must effectively learn combinations of these fields.
This project implements the DeepFM architecture in PyTorch, reproducing the benchmark results on the Criteo dataset.
Architecture & Interaction Mechanics
Unlike traditional architectures that require manual feature combination engineering, DeepFM shares input embeddings between a Factorization Machine (FM) component and a Deep Neural Network (DNN) component.
Output (CTR Probability)
▲
│
┌────────────┴────────────┐
│ Concatenation & Sigmoid │
└────────────┬────────────┘
│
┌──────────────────┴──────────────────┐
│ │
┌──────────┴──────────┐ ┌──────────┴──────────┐
│ FM Component │ │ Deep Component │
│ (1st & 2nd Order) │ │ (High-Order MLP) │
└──────────┬──────────┘ └──────────┬──────────┘
│ │
└──────────────────┬──────────────────┘
│
┌────────────┴────────────┐
│ Shared Embedding Layer │
└────────────┬────────────┘
│
Sparse Input- Shared Embeddings: The sparse input (39 categorical features) is mapped to dense embeddings. Both components share these exact embeddings, enabling simultaneous learning of low- and high-order interactions.
- FM Component: Computes 1st-order linear interactions and 2nd-order dot-product combinations of features.
- Deep Component: Feeds the flattened embeddings through a Multi-Layer Perceptron (MLP) to model non-linear high-order combinations.
Core Pipeline Implementation
1. High-Throughput Data Preprocessing
The Criteo dataset consists of 45 million rows and represents a major memory challenge. To build a robust pipeline:
- Unix Shuffling: Handled shuffling on disk prior to loading via Unix
shufto avoid loading the full text file into RAM. - Ordinal Encoding: Fit
OrdinalEncoderobjects per column to transform sparse categorical labels and continuous integers (bucketed) into tight integer IDs. - Chunked Serialization: Loaded data in chunks, converted them into encoded NumPy
.npyarrays, and serialized encoders to disk withjoblib. - Memory Footprint: Replaced multi-gigabyte pandas DataFrames with memmapped NumPy vectors, enabling fast, zero-copy training batch loads.
2. Ablation Studies & Experimentation
To analyze performance trade-offs, a config-driven ablation runner was developed in config.py to compare six alternative configurations against the baseline model:
BCEWithLogitsLoss3. Training & Evaluation
- Imbalance Handling: The dataset features a ~25% click ratio. Training with
BCEWithLogitsLossprevented loss saturation and gradient clipping. - Early Stopping: Monitored validation AUC with a patience configuration of 3 epochs.
- Analytics: Integrated Weights & Biases (W&B) to log epoch-by-epoch loss, validation AUC, and training speed metrics, ensuring complete reproducibility.
Related Projects
Bitcoin Forecasting Dashboard
An interactive time-series forecasting dashboard for daily BTC/USD prices using Prophet, SARIMA, and Random Forest.
Wikipedia Biographical Clustering
An end-to-end unsupervised clustering pipeline for Wikipedia biographical text using TF-IDF, Word2Vec, and GMM.