The following sections provide greater details of the dataset, running training, and the training results.

Scripts and sample code

These are the important scripts in this repository:

• trainer/task.py - Python script for training the Wide & Deep recommender model
• trainer/features.py - Python file describing the request and item level features

Parameters

These are the important parameters in the trainer/task.py script:

--model_dir: Path to model checkpoint directory
--deep_hidden_units: [DEEP_LAYER1 DEEP_LAYER2 ...] hidden units per layer, separated by spaces
--prebatch_size: Number of samples in each pre-batch in tfrecords
--global_batch_size: Training batch size (per all GPUs, must be a multiplicity of prebatch_size)
--eval_batch_size: Evaluation batch size (must be a multiplicity of prebatch_size)
--num_epochs: Number of epochs to train
--linear_learning_rate: Learning rate for the wide part of the model
--linear_l1_regularization: L1 regularization for the wide part of the model
--linear_l2_regularization: L2 regularization for the wide part of the model
--deep_learning_rate: Learning rate for the deep part of the model
--deep_dropout: Dropout probability for deep model
--deep_warmup_epochs: Number of epochs with linear learning rate warmup
--predict: Perform only the prediction on the validation set, do not train
--evaluate: Perform only the evaluation on the validation set, do not train
--gpu: Run computations on GPU
--amp: Enable Automatic Mixed Precision
--xla: Enable XLA
--hvd: Use Horovod for multi-GPU training
--eval_epoch_interval: Perform evaluation every this many epochs

Command-line options

To see the full list of available options and their descriptions, use the -h or --help command-line option:

python -m trainer.task --help

Getting the data

The Outbrain dataset can be downloaded from Kaggle (requires Kaggle account).

Dataset guidelines

The dataset contains a sample of users' page views and clicks, as observed on multiple publisher sites. Viewed pages and clicked recommendations have additional semantic attributes of the documents. The dataset contains sets of content recommendations served to a specific user in a specific context. Each context (i.e. a set of recommended ads) is given a display_id. In each such recommendation set, the user has clicked on exactly one of the ads.

The original data is stored in several separate files:

• page_views.csv - log of users visiting documents (2B rows, ~100GB uncompressed)
• clicks_train.csv - data showing which ad was clicked in each recommendation set (87M rows)
• clicks_test.csv - used only for the submission in the original Kaggle contest
• events.csv - metadata about the context of each recommendation set (23M rows)
• promoted_content.csv - metadata about the ads
• document_meta.csv, document_topics.csv, document_entities.csv, document_categories.csv - metadata about the documents

During the preprocessing stage the data is transformed into 59M rows tabular data of 54 features and eventually saved in pre-batched TFRecord format.

Spark preprocessing

The original dataset is preprocessed using Spark scripts from the preproc directory. The preprocessing consists of the following operations:

• separating out the validation set for cross-validation
• filling missing data with the most frequent value
• generating the user profiles from the page views data
• joining the tables for the ad clicks data
• computing click-through rates (CTR) for ads grouped by different contexts
• computing cosine similarity between the features of the clicked ads and the viewed ads
• math transformations of the numeric features (taking logarithm, scaling, binning)
• storing the resulting set of features in TFRecord format

The preproc1-4.py preprocessing scripts use PySpark. In the Docker image, we have installed Spark 2.3.1 as a standalone cluster of Spark. The preproc1.py script splits the data into a training set and a validation set. The preproc2.py script generates the user profiles from the page views data. The preproc3.py computes the click-through rates (CTR) and cosine similarities between the features. The preproc4.py script performs the math transformations and generates the final TFRecord files. The data in the output files is pre-batched (with the default batch size of 4096) to avoid the overhead of the TFRecord format, which otherwise is not suitable for the tabular data - it stores a separate dictionary with each feature name in plain text for every data entry.

The preprocessing includes some very resource-exhausting operations, like joining 2B+ rows tables. Such operations may not fit into the RAM memory, therefore we decided to use Spark which is a suitable tool for handling tabular operations on large data. Note that the Spark job requires about 1 TB disk space and 500 GB RAM to perform the preprocessing. For more information about Spark, please refer to the Spark documentation.

Training process

The training can be started by running the trainer/task.py script. By default the script is in train mode. Other training related configs are also present in the trainer/task.py and can be seen using the command python -m trainer.task --help. Training happens for --num_epochs epochs with a DNNLinearCombinedClassifier estimator for the model. The model has a wide linear part and a deep feed forward network, and the networks are built according to the default configuration.

Two separate optimizers are used to optimize the wide and the deep part of the network:

• FTLR (Follow the Regularized Leader) optimizer is used to optimize the wide part of the network.
• Adagrad optimizer is used to optimize the deep part of the network.

The training log will contain information about:

• Loss value after every 100 steps.
• Training throughput if --benchmark option is selected.
• Evaluation metrics after every --eval_epoch_interval epochs.

Checkpoints are stored with every evaluation at the --model_dir location.