Learning Large-scale Universal User Representation with Sparse Mixture of Experts
Abstract
Learning user sequence behaviour embedding is very sophisticated and challenging due to the complicated feature interactions over time and high dimensions of user features.
Recent emerging foundation models, e .
g ., BERT and its variants, encourage a large body of researchers to investigate in this field.
However, unlike natural language processing (NLP) tasks, the parameters of user behaviour model come mostly from user embedding layer, which makes most existing works fail in training a universal user embedding of large scale.
Furthermore, user representations are learned from multiple downstream tasks, and the past research work do not address the seesaw phenomenon.
In this paper, we propose SUPERMOE, a generic framework to obtain high quality user representation from multiple tasks.
Specifically, the user behaviour sequences are encoded by MoE transformer, and we can thus increase the model capacity to billions of parameters, or even to trillions of parameters.
In order to deal with seesaw phenomenon when learning across multiple tasks, we design a new loss function with task indicators.
We perform extensive offline experiments on public datasets and online experiments on private real-world business scenarios.
Our approach achieves the best performance over stateof-the-art models, and the results demonstrate the effectiveness of our framework.
1. Introduction
Recent works have demonstrated that the pre-trained model plays a critical role on a wide range of applications, e g .
, (Devlin et al., 2018; et al., 2020; Riquelme
First Workshop of Pre-training: Perspectives, Pitfalls, and Paths Forward at ICML 2022 , Baltimore, Maryland, USA, PMLR 162, 2022.
Copyright 2022 by the author(s).
et al., 2021;
Bommasani et al., 2021b;
Geng et al., 2022; Sun et al., 2019; Qiu et al., 2020; Khan et al., 2021; Wu et al., 2020; Xiao et al., 2021; Zeng et al., 2021).
To improve the efficiency and effectiveness of these models, many researchers attempt to exploit transformer in order to capture chronological pattern and dynamics of user intentions (Zeng et al., 2021; Xue et al., 2021).
With the remarkable achievements of pre-trained models, especially BERT-based models (Qiu et al., 2021), the transformer backbone has been utilized to address user data sparsity and cold-start problems in downstream applications (Yuan et al., 2020; Zhang et al., 2020a).
In addition, DNN-based self-supervised learning (SSL) model is designed to improve semantic representations for highly-skewed data distribution, with inadequate explicit user feedback in user behaviour sequence interactions via unlabeled data (Yao et al., 2021; Shin et al., 2021; Zhang et al., 2020b).
However, the existing pre-trained model suffers from many difficulties in achieving good user representations, e .
g ., only a few behaviour channels are used in the model due to the huge sizes of vocabularies and the resulting low training efficiency.
In AETN (Zhang et al., 2020a), only three behaviour channels are utilized, yielding sub-optimal user representations.
Therefore, the motivations of our work are threefold, supported by our practical observations in online production system.
Firstly, most of model parameters come from feature embedding of ID and categorical features, which usually dominate GPU memory usage (Lian et al., 2021).
For example, the number of user IDs are often in the scale of billions, resulting in parameter size of numberIDs × embeddingDIMs .
Secondly, the front embedding layer accounts for the majority of the model’s size, while the rest of model layers are extremely computationally expensive.
Consequently, training feature embedding layer and main neural networks simultaneously and synchronously for model of large scale is challenging, which calls for efficient model training algorithm for sparsity.
Finally, there are multiple training objectives no matter in model pre-training stage or in fine-tuning stage, which often causes pre-trained user embedding models with sub-optimal performance when using simple bottom-shared mechanism for the reason of seesaw phenomenon (Tang et al., 2020) and negative transfer (Ma et al., 2018; Chen et al., 2019).
[Table]
Scalability up to Trillions
[Table]
In this paper, we propose SUPERMOE, a general framework for user sequence behaviour representation and prediction using sparse MoE transformer.
Intuitively, transformer demonstrates the importance of capturing long range dependencies and pairwise or higher order interactions between elements (Bommasani et al., 2021a).
The sparse gating mechanism, such as MoE, has shown its great advantages in multi-objective learning in user recommendation systems.
Therefore, embedding the gating function in transformer would be a good alternative to conventional models in user representation learning.
The comparison of advantages and limitations of the proposed model and the other models is listed in Table 1.
Our contributions can thus be summarized as follows: 1)
We propose a sparse MoE transformer model to deal with huge amount of user behaviour sequence data with high dimensions.
We propose a novel multi-task optimization algorithm in order to address seesaw problem and negative transfer problem across multiple tasks.
We devise a novel method to split feature projection layer in order to address the issue of GPU memory explosion, which successfully integrates hundreds of behaviour channels into model training.
Our method significantly outperforms existing user behaviour representation learning methods.
2. Problem Statement
Generally, we denote a typical one-channel user behaviour sequence as s =
[ s 1 , s 2 , ..., s i , ..., s N ] , where s i indicates the i th user behaviour for this channel, which has length of
N .
A multi-channel user behaviour sequence is denoted as
S = { [ s j 1 , s j 2 , ..., s ji ..., s jN ] } , and [ s j 1 , s j 2 , ..., s ji ..., s jN ] is the j th channel of user behaviour sequence corresponding to
M behaviour channels.
Each instance S in each task contains a userID u ∈ U , and three types of sequence channels, namely, category channel S category , ID channel S ID and dense channel S dense .
Therefore, given a set of N tasks
T = { t 1 , t 2 , ..., t n } with corresponding supervised label
Y = { y 1 , y 2 , ..., y n } , our goal is to learn the base user representations across these tasks in order to apply them to downstream applications.
Following the two-stage training paradigm (Devlin et al., 2018), we pre-train a base model firstly on the huge pre-training dataset and then fine-tune a new model on downstream target dataset with parameters initialized as the pre-trained model.
After the training, our base representation model should be able to produce universal representation H to serve all downstream tasks.
3. Methodology 3.1. User Embedding Pre-training Framework
Pre-training Tasks.
Similar to the pre-training task in (Devlin et al., 2018), a new user representation pre-training task is designed to cater to the attribution of user behaviour data, i.e., masked channel prediction (MCP) task.
Slightly different from masked language modeling (MLM) task in
NLP, not all of the features are masked due to multi-channel problem in user behaviour data which would produce too many feature vocabularies.
Theoretically, in the MCP task, some channel elements in the behaviour sequence are randomly masked with special token [ MASK ] at pre-training stage.
Therefore, an MCP task of one feature channel is elaborated as input =
[ s 1 , s 2 , ..., [ MASK ] i ..., s N ] , with label =
[ MASK ] i .
However, only a few channels are selected to be MCP tasks due to our belief that the more important a user behaviour sequence is, the more likely the sequence is selected as MCP task.
In order to preserve essential information of user behaviours, we choose user ID, location, time interval, payment tool, product, trade amount, super position model (SPM) trace, click, and conversion etc.
Pre-training Objectives.
Formally, we denote s mask as the probability of the estimated activity, and the probability p ( s mask ; Θ) is represented by the product of the conditional distributions over the masked sequence:
[Equation]
Our objective is to maximize p ( s mask ; Θ) , which is equivalent to minimizing the following loss function:
[Equation]
… + user_id evcode pdcode paytool ...
𝐸 !"
" 𝐸 !"
# 𝐸 !"
$ … 𝐸 !"
% 𝐸 !
#" 𝐸 ! ## 𝐸 !
#$ … 𝐸 !
#% 𝐸 !
$" 𝐸 !$# 𝐸 !
$$ … 𝐸 !
$% 𝐸 !&" 𝐸 !&# 𝐸 !&$ … 𝐸 !&% 𝐸 !"
" 𝐸 !"
# … 𝐸 !"
% 𝐸 !
# " 𝐸 !
# # … 𝐸 !
# % 𝐸 !
$ " 𝐸 !$ # 𝐸 !
$ % … … 𝐸 !
𝐸 ’" 𝐸 ’# 𝐸 ’$ 𝐸 ’% ... 𝐸 ()" 𝐸 ()# 𝐸 ()$ 𝐸 ()% 𝐸 &" 𝐸 & 𝐸 * 𝐸 * ...
+ + + +
Query
Key
Value 𝑞 " 𝑘 " 𝑣 " 𝑞 # 𝑘 # 𝑣 # 𝑞 + 𝑘 + 𝑣 + … 𝑞 " 𝑘 " 𝑣 " 𝑞 # 𝑘 # 𝑣 # 𝑞 + 𝑘 + 𝑣 +
Paytool Time Loc
Task ! ……
…… … 𝑥 " 𝑥 # 𝑥
[Figure]
𝑥 +
[Figure]
+
[Figure]
y
[Figure]
FFN 1 FFN
[Figure]
FFN
[Figure]
Query
Key
Value 𝑞 " 𝑘 " 𝑣 " 𝑞 # 𝑘 # 𝑣 # 𝑞 + 𝑘 + 𝑣 + … 𝑞 " 𝑘 " 𝑣 " 𝑞 # 𝑘 # 𝑣 # 𝑞 + 𝑘 + 𝑣 +
Label
[Figure]
Finetune
[Figure]
Loss where S i is the set of positions of masked elements of the i th MCP task, and ˆ s j and s j are the predicted user behaviour and the ground-truth behaviour, respectively.
Notably, user behaviours are of very different statistical characteristics from NLP or CV, e.g., the click and conversion task are sequential tasks.
Hence, we propose a new training objective function:
[Equation]
where δ is the indicator of training samples among k tasks.
Pre-training Model Framework.
The main architecture ingredients of pre-training model are a stack of MoE transformers.
Basically, our MoE transformer’s backbone has a simple structure which consists of a multi-channel feature projection (MFP) layer, a MoE multi-head self-attention (MoE-MSA) layer and two MoE feed-forward network (MoE-FFN) layers.
MFP layer takes the following form:
[Equation]
where [ · ] means the concatenation operator of all vectors.
Each MFP layer in the encoder block is followed by a layer normalization and nonlinear activation layer.
The operator split ( · ) is a model parallel operation, implemented by the whale framework (Xianyan Jia, 2022).
Note that the splitting of MFP layer addresses the issue of GPU memory explosion, which successfully integrate hundreds of behaviour channels into model training.
An MoE-MSA layer takes the output of MFP y m pf as input, formulated as:
[Equation]
( v )) , (5) where q, k, v is the output of an MFP layer, and y msa is output of an MoE-MSA layer, connected by two MoE-FNN layers.
Lastly, the point-wise MoE-FNN(Fedus et al., 2021) can be formulated as:
[Equation]
with FFN
[Equation]
x ) =
w
[Equation]
· Relu ( w
[Equation]
· x ) , G
[Equation]
x )
[Equation]
softmax ( TopK ( h e ( x ) , k )) , where w o and w i are the standard feed-forward networks with the same parameters.
We choose top 1 strategy (Fedus et al., 2021) for TopK ( · ) function.
In summary, y ffn is the output of backbone of an MoE transformer.
Formally, a series of MoE transformer blocks can be described as:
[Equation]
where MoETransformer = MoE FFN ( MoE MSA ( MFP ( · ))) .
The overall pre-training architecture is shown in Figure 1.
3.2. User Embedding Fine-tuning Framework
After pre-training, we adapt the learned user representations to specific downstream tasks, instead of using pre-trained representations directly, which is somehow unrelated to our defacto targets in production environment.
Therefore, we need to develop a new model to fine-tune our user behaviour model across multiple downstream tasks with a unified framework.
Assuming that we have restored and initialized the parameters of the previous pre-trained model, the fine-tuning model shares the parameters of the pre-trained part, and a linear classification layer is placed on the top of the final output without activation function.
Denoting h o as the output of the final MoE-transformer, we have:
[Equation]
and the Tower i is a linear classification layer of the i th fine-tuning task.
Note that the user representation H =
MaxPooling ( h o ) .
The overall architecture of our fine-tuning framework is shown in Figure 1.
3.3. Multi-task Training Optimization
In order to address seesaw and negative transfer problems and to improve learning from multiple tasks, such as regression and classification, we leverage a multi-task optimization strategy, i.e., jointly optimize across multiple tasks, which can be applied in both pre-training stage and fine-tuning stage.
Mathematically, we get k training objectives from equation ( 5 ), and therefore, the total loss can be formulated as:
[Equation]
where Loss (Θ) denotes the total loss and α k is the regularization strength of the k th loss.
Recall that our objective is actually to maximize Area Under Curve (AUC) score, we consider the following bi-level optimization problem:
[Equation]
where AUC val is the AUC score on validation dataset while training.
However, AUC val ( θ λ , λ ) is non-differentiable with the indicator function I ( f ( λ, x + i ) <
f ( λ, x − j )) , and x + i and x − j are the positive and negative samples, respectively.
We therefore employ max { 0 , 1 − ( f ( λ, x + i ) − f ( λ, x − j )) } as a differentiable convex surrogate of the above indicator function.
4. Experimental Methodology
In this section, we demonstrate the online and offline performance of SUPERMOE in generating general embedding for user behaviour sequence.
We evaluate our model in four different real world test datasets, and one for public and three for private datasets respectively.
4.1. Experiment Settings 4.1.1. D ATASET DESCRIPTION
We evaluate the performance of our model on four different downstream applications, i.e., SIUPD, Paytool, MCP, and Fortune.
SIUPD dataset comes from the IJCAI17 contest 1 , which contains 139,6245 users’ shopping logs on
Alipay platform.
Paytool is a user payment preference dataset, which describes the behaviour of using payment tools for online users.
In MCP dataset, we use 103 channels of subscription and redemption behaviour sequences for users.
Fortune dataset includes users ”impression → click” and ”click → purchase” behaviours.
All these four datasets are split into training/test sets with the ratio of 0 .
The statistics of the datasets can be found in Table 3.
4.1.2. B ASELINES
We fine-tune and evaluate our model against four other representative models: MMOE(Ma et al., 2018) , a classical multi-task recommendation model, PLE(Tang et al., 2020) , an extension of MMOE with multiple progressive extraction layers, BERT(Devlin et al., 2018) , a well-famed sequence model widely used in large scale representation learning, especially in NLP and AETN(Zhang et al., 2020a) , a user representation learning model, which combines multi-head attention and Denoising Autoencoder(DAE) model to generate user embeddings.
4.2. Offline Evaluation Results
In order to show the advantages of our model, we conduct the following intrinsic experiments to evaluate offline and online performances.
4.2.1. O FFLINE M ODEL P ERFORMANCE
In this section, we present the results of offline model performance in the downstream tasks.
Table 2 summarizes the overall AUC scores of different models across all datasets.
Taking the evaluation results of SIUPD dataset as an example, it is obvious that our model improves the baseline method MMoE by gains of 2.7 and 1.8, respectively, in two combined tasks, for the reason that our model utilizes more abundant chronological user behaviours to address the behaviour sparsity issue.
Moreover, we outperform the
[Table] [Table]
other two sequential models with gains of 0.63 and 0.19, respectively, benefiting from of our multi-task optimization.
Similar performances can be observed in other three datasets.
It is worth mentioning that our methods all achieve the state-of-the-art performances with significant gains.
4.2.2. O FFLINE E MBEDDING P ERFORMANCE
To evaluate the user embedding quality and efficiency of our model, we conduct six different experiments for comparison, and analyze the effects of different embedding methods, as well as different model capacities.
We select the user’s payment switching task in PAYTOOL dataset to report AUC score, Recall@85 and Recall@50 respectively.
The results are illustrated in Table 4.
Notably, all sequential embedding methods are better than PLE-only model, which demonstrates the advantage of user embedding.
Furthermore, our embedding is more effective than other two sequential models, which takes the same model size of 1 billion.
We also investigate the performance of different model capacities, and it can be seen in Table 4 that MoE with 20 billions parameters performs much better than MoE with 1 billion, which generates gains of 0.67 AUC, 2.01 recall@85, and 1.39 recall@50, respectively.
4.2.3. O NLINE A/B T ESTING
To further investigate the quality and effectiveness of our user embeddings
[Table]
we conduct two A/B testing experiments against online baseline model. ”
Online1” experiment is a payment switching scenario operating on real-world Alipay platform.
In this experiment, our model brings on gains of 13.41% pv, 1.97% in conversion and 21.36% GMV.
In addition, our model achieves gains of 4.95%,9.11% and
[Table] [Table]
OURS
[Table] [Table]
[Table]
Online
PLE
[Table]
BERT
OURS 4.95% 9.11% 25.19% 25.19%, respectively, in ”Online2” experiment, which is a fund subscription and redemption scenario.
These results are summarized in Table 5.
5. Conclusions
In this paper, we investigated the utilization of multi-layer
MoE networks as a practical way to massively increase model capacity and to deal with seesaw phenomenon and negative transfer problem.
To complete this research, we introduce an user behaviour representation pre-training and fine-tuning model using sparse MoE. We have shown that it is possible to learn large scale user embeddings, while capturing ubiquitous high order correlations using sparse
MoE, with our meticulous model architecture.
Moreover, we formulated a bi-level optimization method in order to address multi-task optimization.
Extensive empirical experiments demonstrated the overwhelming superiority of our method on various real-world datasets comparing to other state-of-the-art methods.