GDMix

What is it

Consider a job recommendation task where two LinkedIn members Alice and Annie have very similar profiles. Both of them have the same user features and respond to the same pair of companies. Their responses are exactly opposite to each other. If we use both member's data to train a machine learning model, the model won't be effective since the training samples contradict each other. A solution is to train a single model for each member based on the member's data. This is an example of personalization.

A possible implementation of personalization is to include all member ID embeddings in a single model. This usually results in a very large model since the number of members can be on the order of hundred of millions. GDMix (Generalized Deep Mixed model) is a solution created at LinkedIn to train these kinds of models efficiently. It breaks down a large model into a global model (a.k.a. “fixed effect”) and a large number of small models (a.k.a. “random effects”), then solves them individually. This divide-and-conquer approach allows for efficient training of large personalization models with commodity hardware. An improvement from its predecessor, Photon ML, GDMix expands and supports deep learning models. Check out our engineering blog for more background information.

Supported models

The current version of GDMix supports logistic regression and DeText models for the fixed effect, then logistic regression for the random effects. In the future, we may support deep models for random effects if the increase complexity can be justified by improvement in relevance metrics.

Logistic regression

As a basic classification model, logistic regression finds wide usage in search and recommender systems due to its model simplicity and training efficiency. Our implementation uses Tensorflow for data reading and gradient computation, and utilizes L-BFGS solver from Scipy. This combination takes advantage of the versatility of Tensorflow and fast convergence of L-BFGS. This mode is functionally equivalent to Photon-ML but with improved efficiency. Our internal tests show about 10% to 40% training speed improvement on various datasets.

DeText models

DeText is a framework for ranking with emphasis on textual features. GDMix supports DeText training natively as a global model. A user can specify a fixed effect model type as DeText then provide the network specifications. GDMix will train and score it automatically and connect the model to the subsequent random effect models. Currently only the pointwise loss function from DeText is allowed to be connected with the logistic regression random effect models.

Other models

GDMix can work with any deep learning fixed effect models. The interface between GDMix and other models is at the file I/O. A user can train a model outside GDMix, then score the training data with the model and save the scores in files, which are the input to the GDMix random effect training. This enables the user to train random effect models based on scores from a custom fixed effect model that is not natively supported by GDMix.

Training efficiency

For logistic regression models, the training efficiency is achieved by parallel training. Since the fixed effect model is usually trained on a large amount of data, synchronous training based on Tensorflow all-reduce operation is utilized. Each worker takes a portion of the training data and compute the local gradient. The gradients are aggregated then fed to the L-BFGS solver. The training dataset for each random effect model is usually small, however the number of models (e.g. individual models for all LinkedIn members) can be on the order of hundred of millions. This requires a partitioning and parallel training strategy, where each worker is responsible for a portion of the population and all the workers train their assigned models independently and simultaneously.

For DeText models, efficiency is achieved by either Tensorflow based parameter server asynchronous distributed training or Horovod based synchronous distributed training.

How to setup

GDMix has mixed implementation of Python and Scala, which used for training model and processing intermediate data on Spark, GDMix requires Python 3.7+ and Apache Spark 2.0+.

As user

User will need to :

Install Spark 2.0+
Compile a gdmix-data fat jar used for intermediate data processing
Install the gdmix-trainer and gdmix-workflow python packages

# The compiled jar will be at directory build/gdmix-data-all_2.11/libs
./gradlew shadowJar
pip install gdmix-trainer gdmix-workflow

See more details in the section Tryout the movieLens example for a hands-on example and the section Run GDMix on Kubernetes for distributed training for distributed training.

As developer

GDMix's Python and Scala module use setuptools and Gradle to manage code, respectively. We recommend to use miniconda to manage Python dev virtual environment. Local build and test is shown below.

# Build scala module
./gradlew clean build

# Create virtural env
conda create --name gdmix python=3.7.4 -y
conda activate gdmix
pip install pytest

# Install from local gdmix-trainer dev code
cd gdmix-trainer && pip install .
# Run pytest
pytest

# Install from local gdmix-workflow dev code
cd gdmix-workflow && pip install .
# Run pytest
pytest

GDMix implementation

Implementation overview

The overall GDMix training flow is shown in Figure 1. The data preparation step is provided by user and the input files should meet the requirements shown in the next section Input data. The fixed effect captures the global trend and the random effects account for the individuality. The complexity of enormous cross-features from a recommender system is overcomed by taking a parallel blockwise coordinate descent approach, in which the fixed effect and random effects can be regarded as “coordinates”, during each optimization step, we optimize one coordinate at a time and keep the rest constant. By iterating over all the coordinates a few times, we arrive at a solution that is close to the solution to the original problem.

<figure> <img src="figures/gdmix-overview.png" alt="" /> <ficaption>Figure 1: GDMix Overview</ficaption> </figure>

Currently, GDMix supports three different operation modes as shown below. | Fixed effect | Random effect | |----------|:-------------:| | logistic regression | logistic regression | | deep NLP models supported by DeText | logistic regression | | arbitrary model designed by a user | logistic regression |

In the last mode, the fixed effect model is trained by the client with his/her own model outside of GDMix, then the scores from that model are treated as input to GDMix random effect training.

<figure> <img src="figures/gdmix-operation-models.png" alt="" /> <ficaption>Figure 2: GDMix Operation Models</ficaption> </figure>

Figure 3 shows jobs of fixed effect and random effect model training for an example with one fixed effect model global and two random effect models per-user and per-movie. The global model is trained on all data and typically requires distributed training; while the random effect models per-user and per-movie are numerous small independent models according to the entity. To accelerate the random effect model training parallelism, a data partition job is needed to group the records by the entity Id (e.g. user id), which can be done efficiently by Spark. In addition, metrics are computed for each model.

<figure> <img src="figures/gdmix-workflow-jobs.png" alt="" /> <ficaption>Figure 3: GDMix Workflow Jobs</ficaption> </figure>

Input data

The input data should be organized in following structure, in which fixed-effect and each random-effect (per-user and per-movie) has a sub-directory containing featureList, metadata, trainingData and validationData directories:

├── fixed-effect
│   ├── featureList
│   │   └── global
│   ├── metadata
│   │   └── tensor_metadata.json
│   ├── trainingData
│   │   └── part-00000.tfrecord
│   └── validationData
│       └── part-00000.tfrecord
├── per-user
│   ├── featureList
│   │   └── per_user
│   ├── metadata
│   │   └── tensor_metadata.json
│   ├── trainingData
│   │   └── part-00000.tfrecord
│   └── validationData
│       └── part-00000.tfrecord
├── per-movie
│   ├── featureList
│   │   └── per_movie
│   ├── metadata
│   │   └── tensor_metadata.json
│   ├── trainingData
│   │   └── part-00000.tfrecord
│   └── validationData
│       └── part-00000.tfrecord

The trainingData and validationData are the preprocessed training and validation data in tfrecord format; featureList directory contains a text file that lists all feature names; metadata directory has a json file that stores the number of training samples and metdata such as name, data type, shape and if is sparse vector, which is used to deserialize the tfrecord data. An example for the global model is shown below:

{
  "features" : [ {
    "name" : "weight",
    "dtype" : "float",
    "shape" : [ ],
    "isSparse" : false

Gdmix

Install / Use

README