Enhanced Exploration in Neural Feature Selection for Deep Click-Through Rate Prediction Models via Ensemble of Gating Layers

Deep Click-Through Rate (CTR) prediction systems have shown promising performance in many industry-scale CTR prediction tasks, but this usually comes with high computational cost and high memory usage. To make deep learning models more memory efficient, people have proposed many approaches to reduce the size of the networks, such as network pruning (Xia, Zigeng, and Sanguthevar 2019; Gao et al. 2020), training with resource constraint (Srinivas, Subramanya, and Venkatesh Babu 2017), and adjusting feature embedding sizes (Joglekar et al. 2019; Liu et al. 2021). In this work, we focus on improving memory efficiency by performing neural feature selection that chooses a relatively small subset of features with good explanatory power.

We consider end-to-end feature selection approaches, which interpret the feature selection problem as a differentiable neural architecture search (NAS) or network pruning problem, and thus enable people to use many well-studied methods from the two relevant domains. In those approaches, to remove less predictive features, a set of “binary” gates with trainable weights are added to the network to indicate if certain features should remain or drop (Fig. 1). Depending on how discrete relaxation is achieved, there can be different implementations of the binary gates (Sheth and Fusi 2020; Yamada et al. 2020). Similar to differentiable NAS (Han, Ligeng, and Song 2019) and network pruning (Xia, Zigeng, and Sanguthevar 2019; Gao et al. 2020), differentiable neural feature selection methods realize their functionality by iteratively performing feature pruning (updating binary gates) and network fine-tuning (updating network parameters) over a pretrained over-parameterized network. In this paper, we refer to the methods that use such a gating formula as Gating method and the gate output as gating decision.

Since the feature selection problem is essentially a search problem with the search space comprised of all possible gating decisions, to find the most informative subset of features, the gates must fully explore the search space by taking various gating decisions and comparing the outcomes. However, in the Gating approach, the gating exploration is mainly driven by the gradient signal, which doesn’t always provide a strong enough exploration incentive, especially when the binary gates are jointly optimized with the model parameters. Here, we present a special experiment to illustrate how insufficient exploration can result in undesirable outcomes. In this experiment, we use a step function and straight through estimator (Bengio, Léonard, and Courville 2013; Hubara et al. 2016) to implement the binary gates. Details of the implementations and the searching/training process can be found in later sections. But different from normal settings, we deliberately reduce the learning rate of binary gates by a factor of ten. The goal of this setting change is to exaggerate the negative effects of insufficient exploration. Specifically, suppose the gradient signal suggests the $i^{th}$ gate should be changed from “on” to “off” and suppose there is no other exploration incentive except the gradient, since the learning rate of binary gates is abnormally small, it takes more update steps for the $i^{th}$ gate to change its state. However, before the $i^{th}$ gate is turned off, the model parameters have already been fine-tuned according to current gating decision. In this case, changing gate states will instead increase the training loss. We call this phenomenon gating overfitting. The experiment result verifies our hypothesis (Fig. 2). Note that gating overfitting exists even when all the hyper-parameters are properly set.

To enhance gating exploration and prevent gating overfitting, we propose a scalable and effective ensemble learning method, named Ensemble Gating, which creates multiple groups of binary gates and randomly selects one group to perform parameters update at each step. We show that, without adding any computational overhead or introducing any hyper-parameter (except the size of the ensemble), Ensemble Gating can provide meaningful uncertainty-driven exploration. We conduct extensive experiments on two public datasets, namely Avazu and Criteo. The results demonstrate that Ensemble Gating consistently finds better subsets of features on both datasets for three different deep CTR prediction models. Further studies show that Ensemble Gating converges quickly within a few epochs, and the uncertainty-driven exploration is more effective than exploring by injecting random noise into the learning system.

The idea of using differentiable binary gates to automatically select important features was discussed in (Sheth and Fusi 2020; Yamada et al. 2020). It has also been widely applied in industry and has been studied in many relevant research directions like gradient-based neural architecture search (Han, Ligeng, and Song 2019; Liu, Simonyan, and Yang 2018) and network pruning (Han et al. 2015; Srinivas, Subramanya, and Venkatesh Babu 2017; Ye et al. 2020), which use learnable binary gates to determine whether to prune certain weight or neuron. Gating approaches in neural feature selection can be viewed as a reduced version of related network pruning approaches in the sense that binary gates are only inserted between certain layers in neural feature selection.

Due to the discrete nature of feature selection/network pruning problems (there are only two possible actions, namely ”prune” or ”keep”), there is no direct way to optimize the binary gates along with the deep neural network in a fully end-to-end manner. Thus, substantial research efforts have been focusing on finding better approximation that relaxes the discrete search space to a continuous space by using some differentiable re-parameterization tricks. Works like (Xia, Zigeng, and Sanguthevar 2019) and (Gao et al. 2020) use straight through estimator (Bengio, Léonard, and Courville 2013; Hubara et al. 2016) to enable gradients to be propagated. Some other works also use Gumbel-Softmax (Jang, Gu, and Poole 2016; Maddison, Mnih, and Teh 2016) to do the discrete relaxation (Xie et al. 2018; Liu, Simonyan, and Yang 2018).

The problem of neural feature selection (NFS) is also related to the neural input search (NIS) problem (Joglekar et al. 2019; Ginart et al. 2019). But the latter aims to reduce memory usage by assigning varying embedding sizes to different features rather than dropping an entire feature. Many works employ a similar differentiable gating formula to tackle the NIS problem (Zhao et al. 2021; Cheng, Shen, and Huang 2020; Liu et al. 2021), but in their frameworks, each binary gate is used to control the dimension of each embedding vector. Note that NIS and NFS are not interchangeable. Firstly, NIS can not be applied to models that require feature embeddings of the same size, such as DeepFM (Guo et al. 2017). Secondly, it can be more effective to perform feature selection before applying NIS or other network pruning techniques, especially for industry-scale CTR prediction models.

Feature selection for CTR prediction systems has been explored previously (Ronen et al. 2013; Koenigstein and Paquet 2013). In this work, we focus on improving the gradient-based feature methods, which are shown to be more computationally efficient with large-scale deep learning models (Sheth and Fusi 2020). Comprehensive surveys on feature selection are provided by Parmezan et al. and Chandrashekar and Sahin.

Figure 1 depicts a typical deep learning CTR prediction model. We assume the model input $X$ consists of $\mathcal{M}$ categorical feature fields. Each raw categorical feature is initially represented by a sparse one-hot vector. Then the feature embedding layer transforms each sparsely encoded feature $x_{i}$ into a $V$-dimensional embedding vector $e_{i}$ as follows:

where $E_{i}\in R^{D\times V}$ denotes the embedding matrix, and $D$ denotes the size of sparse encoding. Note that in this paper, we assume a global embedding size $V$ (default value: 8) is used for all the features. Finally, all the dense feature embeddings are concatenated into an embedding matrix that is used by the remaining parts of the model (i.e. feature interaction layers) to predict the probability that a user likes an item (a.k.a. click-through rate prediction). Let $\theta$ be the parameters of feature interaction layers, and $\Theta=\{E,\theta\}$ be the set of trainable model parameters, the output of the model is given by:

where $\phi$ represents the CTR prediction model, $\hat{y}$ is the model prediction. Note that the architecture of feature interaction layers may vary in different CTR prediction models.

The CTR prediction model $\phi$ is trained to optimize the following objective:

where $\mathcal{D}$ is the size of the training dataset and $\mathcal{L}$ is the model loss function (binary cross-entropy in usual) along with a weight decay regularization term.

In this section, we first formalize the problem of neural feature selection (NFS) and discuss existing solutions that are adapted from gradient-based NAS and network pruning. Then we introduce our Ensemble Gating algorithm that aims to offer effective uncertainty-driven exploration to overcome the gating overfitting problem.

In this section, we present the experiments that evaluate the effectiveness of Ensemble Gating. Firstly, we intend to investigate if our ensemble method can consistently improve Gating. To this end, we compare Ensemble Gating against Gating on two public benchmark datasets (Avazu¹¹1https://www.kaggle.com/c/avazu-ctr-prediction and Criteo²²2https://www.kaggle.com/c/criteo-display-ad-challenge) with three different underlying deep CTR prediction models, namely DeepFM (Guo et al. 2017), AutoInt (Song et al. 2019), and DCN (Wang et al. 2017). Note that depending on the ensemble aggregation method used, there can three different implementations of Ensemble Gating, which are referred to as Ensemble Avg, Ensemble Voting, and Ensemble Min, respectively. Then, to get more insights into Ensemble Gating, we perform a set of complementary experiments, such as varying the numbers of gating layers (ensemble size), and initializing gating parameters in different ways.

Datasets and Tasks To comply with the policies of most companies, we refrain from presenting the empirical results obtained on any industrial large-scale dataset. Instead, we chose to use two largest open CTR prediction datasets, namely Avazu and Criteo. In Avazu, there are 22 feature fields and the task is to select 6 features that have the highest explanatory power. In Criteo, there are totally 39 feature fields and the task is to select a subset of 19 features.

We randomly split the dataset into 80% and 20% for training and testing. The testing dataset is further randomly split in half for validation and testing in our complementary experiments where network parameter update and gating parameter update are performed on different subsets of data. The evaluation metric we adopt is the AUC (Area Under the ROC Curve) achieved on the testing set by the model that is retrained with the selected features.

Training Setup Both Ensemble Gating and the baselines follow the same training process. Firstly, we pretrain a CTR prediction model with all input features for 10 epochs. We apply the Adam optimizer (Kingma and Ba 2015) with a learning rate of 0.001 for network parameters $\Theta$. The mini-batch size is set to 2048 and the weight of weight decay regularizer is set to $1e^{-6}$. The weights of pretrained model are saved for later steps.

After obtaining the pretrained model, the gating layer(s) will be inserted between the embedding layers and the interaction layers. Then we iteratively updated the network parameters $\Theta$ and the gating parameters $\alpha$ for 8 epochs. We use SGD with a learning rate of 0.001 to optimize the gating parameters.

Following the Lottery Ticket Hypothesis (Frankle and Carbin 2018), our last step is to retrain the model from scratch with selected features on the training set (for 10 epochs) and test it on the testing set. The final gating decision of Ensemble Gating is calculated according to the particular aggregation method. In our preliminary experiments, we observed that in both Gating and Ensemble Gating, the number of open gates at the end of training is occasionally one or two smaller than the target value $\mathcal{N}$. The most likely cause of this phenomenon is the choice of the hyperparameter $\beta_{s}$ (weight of sparse regularizer). Considering that finding the “optimal” value of $\beta_{s}$ can be very expensive in practice, we adopt a simple workaround, which adds unselected features with the largest gate weights (average gate weights in ensemble method) to the set of selected features. This process resembles adjusting the gating threshold that is set to zero in the binarized function (Eq. 6). Our preliminary results suggest that this post-processing doesn’t have any observable negative impact on the final outcome.

In this work, we presented a novel ensemble method that offers a more efficient exploration strategy to find the optimal subset of features. We show that our proposed method, Ensemble Gating, consistently improves Gating method on two public datasets with three different underlying CTR prediction models. The proposed method is computationally efficient and can be easily applied to a variety of CTR prediction models.

Future work may include studying whether our ensemble method can benefit related research problems, such as differentiable network pruning and neural input search. Insufficient exploration can be detrimental in any search problem, and gating overfitting may exist as long as similar gating formulations are employed. Moreover, we note that we have used our ensemble method to improve Gating method implemented with step function and STE. Future work can experiment with other binarize functions.