Learning Representations of Categorical Feature Combinations via Self-Attention

1 INTRODUCTION . Self-attention , as a special attention mechanism , has been widely used to model the sequential data and achieved very remarkable results in many applications , e.g. , neural machine translation ( Vaswani et al. , 2017 ) , sentiment analysis ( Lin et al. , 2017b ) , and reading comprehension ( Wang et al. , 2017b ) . Different from recurrent and convolutional neural networks ( Gehring et al. , 2017 ) , self-attention models dependencies without regard to the positions of sequences . In order to capture the sequential behavior , such attention mechanism is often used in conjugate with recurrent neural networks ( Lin et al. , 2017b ) or built upon additional embeddings of the positions ( Vaswani et al. , 2017 ) . On the other hand , when modeling orderless but relational data , the independence on positions makes self-attention a compelling choice . However , as far as we know , there is little work applying self-attention beyond the sequential data . In this work , we extend it to learn representations of multi-field categorical data in prediction tasks . Taking the prediction of recommendation systems in e-commerce as an example , the service provider builds models to predict how much a customer ( with features of gender , age1 , etc . ) does like a product ( with features of category , price , etc ) . The recommendations are then made by ranking the prediction scores . Quite often , input features are not independent , e.g. , female customers are more prone to buying dresses than male customers . Combinations over the two fields , i.e. , gender and category , can thus lead to more accurate predictions . In applications , categorical features are firstly transformed into sparse representations via one-hot encoding . The combinations in linear models are then made by cross product over different fields . Due to the sparsity problem , the combinations rely on much manual work of domain experts ( Lian et al. , 2018 ; Wang et al. , 2017a ; Qu et al. , 2016 ) . Factorization based models instead additionally learn an embedding vector for each feature ( Rendle , 2010 ) . They compute interactions of features by their inner products . Compared with linear models , 1Note that numerical feature can be buketized to category factorization based models can generalize well to unseen or rare features combinations . However , the simple inner product operation limits its approximation ability to some extent . DNN based models , further extend that by learning more complicated interactions . Typically , all embeddings are firstly concatenated into a long vector as a representation of all features and then fed into vanilla multi layer perceptions ( MLPs ) . Recently , several works ( Qu et al. , 2016 ; Guo et al. , 2017 ; Lian et al. , 2018 ) improve the structure by adding product units ( Durbin & Rumelhart , 1989 ) to construct explicit interactions . However , none of these works improve the representation rather than simple concatenation . In parallel , Liu et al . ( 2015 ) consider relating the embeddings of neighboring features with convolutional neural networks . However , it highly depends on the positions of the features . Different orders of features can yield significant different results while it is NP-hard to find the optimal ( Chan et al. , 2018 ) . As introduced earlier , in this work , we instead apply self-attention to transform the feature embeddings due to it is insensitive to positions . The output vector in each position aggregates the input vector with its related feature embeddings . Thus we denote the outputs of self-attention as the representations of feature combinations . The contribution of this work is summarized as follows : ( a ) Under the observation that self-attention models dependency without regard to positions , we extend it to model orderless while relational categorical features in prediction tasks . As far as we know , this is the first work of self-attention applying to non-sequential data . ( b ) By utilizing self-attention , we are able to relate features in different fields and automatically learn representations of sophisticated feature combinations , which are known as the ingredient in many successful linear prediction models . In contrast , most of current DNN-based models simply concatenate all feature embeddings . ( c ) In the original self-attention structure , each vector is aggregated with all vectors weighted by similarity . This reduce effective resolution and Multi-Head Attention is thus proposed . In this work , we further improve that by restricting the similarity weight to top-k non-zero values . This modification is also consistent with the prior knowledge that one feature is almost impossible to be of relevance with all others . The truncation thus additionally works as a regularization ( see Subsection 3.2 ) . ( d ) As a concrete example , we conduct experiments on CTR estimation . Across the two benchmark CTR datasets , i.e. , Cretio and Avazu , our model with top-k restricted self-attention achieves state of the art performance . Compared with vanilla DNN based model , the gain by adding self-attention is significantly larger than that by modifying the network structures , which most current works focus on . Besides , we give some interpretation on the combinations learned via self-attention . 2 EXISTING EMBEDDED BASED MODELS . The embedded based models for prediction over multi-field categorical data can mainly be divided into two categories : Factorization and DNN based models . The first work among factorization based models is Factorization Machine ( FM ) ( Rendle , 2010 ) . In addition to the linear part , it adds an additional interaction part for the prediction , i.e. , f ( x ) = b+ 〈w , x〉+ n∑ i=1 n∑ j=i+1 〈ei , ej〉xixj , ( 1 ) where ei ∈ Rk , k n represents the low dimensional embedding vector for feature i. FM computes all possible second order interactions . For effective feature combinations that are rare or even unseen in the training data , FM can still generalize to the testing set . High order FM extends that by interacting with more features . As pointed out by Xiao et al . ( 2017 ) , some interactions of FM could be noisy as these features may have no intrinsic relations . Xiao et al . ( 2017 ) instead learns an additional weight for each interaction with a neural network for further distinguishing its effectiveness . Juan et al . ( 2016 ) also extend FM to Field aware Factorization Machines ( FFM ) . FFM learns different embeddings for each feature to be interacted with different fields , which achieves better performance at the cost of larger memories . Most of DNN based models , instead , concatenate all looking up embeddings into a long vector , which is then fed into different neural networks . By using vanilla MLP , Zhang et al . ( 2016 ) propose FNN to learn the implicit interactions . The embeddings are initialized with FM , which could also limit their performance . PNN ( Qu et al. , 2016 ) extends that by inserting additional product units ( Durbin & Rumelhart , 1989 ) to the network . Although the product operation can be approximated arbitrarily well by a two layer ’ s neural network with only 4 neurons in the hidden layer , see Theorem ( Lin et al. , 2017a ) , the extra units improve the performance . Recently , Guo et al . ( 2017 ) propose DeepFM by incorporating MLP with FM . To some extent , the structure can also be regarded as adding product units , but to the final layer instead of to the input layer as PNN . xDeepFM ( Lian et al. , 2018 ) further replaces the FM part by constructing higher order interactions . Although achieving better performance than vanilla MLP , all of these works build upon the simple concatenation , which could be insufficient to model the features relations . 3 MODEL STRUCTURE . In this section , we introduce our model for the prediction tasks . Our work builds upon the selfattention structure introduced by ( Vaswani et al. , 2017 ) . Figure 1 gives an overview of the overall architecture . As shown in the left part of Figure 1 , we firstly concatenate the embeddings as a matrix : E = [ . . . , ei , . . . , ej , . . . ] T , ( 2 ) where E ∈ Rm×d , d and m denotes the dimension of the field embedding and the number of fields , respectively . Then the embeddings are transformed via stacked self-attention . After the transformation , the outputs are fed into the final neural network . Note that we use skip connections within self-attention , which guarantees that the model at least will not be worse than the basic DNN model without self-attention ( He et al. , 2016 ) . 3.1 SELF-ATTENTION . Self-attention maps the input with weighted sum to the output by computing its similarity over different positions . It consists of three parts , i.e. , the queries , the keys , and the values . All of them are derived from the same embedding matrix E by linear projection , usually following with the ReLU activation : Q = max ( EWq , 0 ) , K = max ( EWk , 0 ) , and V = max ( EWv , 0 ) , ( 3 ) with Wq ∈ Rd×d , Wk ∈ Rd×d , and Wv ∈ Rd×d being learned parameters . We then compute dot products of the queries with the keys . The resulting matrix is denoted as the attention score map , i.e. , the A in upper Figure 1 . The element aij captures the relation between the i-th position and j-th position . By applying a softmax function to the score map , we obtain the matrix of outputs : Ṽ = softmax ( QKT√ d ) V +E , ( 4 ) where √ d normalizes the weights . Additionally adding E denotes the skip connection . And the softmax operates on each row , ensuring all the computed weights summing up to one . Ṽ is further transformed by a two layers ’ network with ReLU activation : Ẽ = max ( Ṽ W1 , 0 ) W2 + Ṽ , ( 5 ) where the latter part also represents the skip connection . In equation 4 , each row of Ṽ aggregates all values of V with weighted sum . This operation reduces the effective resolution of attention especially for the data with many fields . To alleviate it , the multi-head structure is proposed in the original paper ( Vaswani et al. , 2017 ) . By projecting the embedding E into h different sets of queries , keys , and values , multi-head attention computes the aggregation as Ṽ = [ softmax ( Q1K T 1√ d/h ) V1 , · · · , softmax ( QhK T h√ d/h ) Vh ] +E , ( 6 ) which is then processed by equation 5 . We also stack several layers of self-attention for further performance improvements . 3.2 TOP-k SPARSITY The multi-head attention structure alleviates the resolution problem by operating on several subspaces . In this subsection , we further improve the structure by ensuring sparsity on the weights after the softmax function . Intuitively , each feature is impossible to have relations with all others . Weighting on all the values as equation 4 will include both useful and useless combinations . The useless feature combination may introduce additional noise and adversely harm the performance . Instead , by ensuring sparsity on the weights , we focus on more related features . Besides , the sparsity constraint will filter out some unrelated connections , which to some extent regularize the model . To ensure the sparsity , a feasible solution is to directly set a threshold . The element of the attention score map A is set to −∞ when it is less than the threshold , and is preserved when greater than or equal to the threshold . However , in practice , we find that it is not easy to choose a proper threshold . Instead , we adopt the simple top-k restriction , i.e. , preserving the top-k values in each row of A and set the rest to −∞ . Then equation 4 becomes Ṽ = softmax ( TopK ( QKT , k ) √ d ) V +E , ( 7 ) where the TopK ( · , · ) function is define as TopK ( A , k ) ij = { aij , if aij is the top k elements in the i-th row of A , −∞ , otherwise . ( 8 ) The top-k restriction is also used in ( Shazeer et al. , 2017 ) to threshold the value in the softmax , where it is mainly for reducing computation of each expert network . Here we are to improve the resolution of averaging and further regularize the model . In practice , it is used in conjugate with multi-head structure . As the experiments shows , it can additionally improve the performance .

Summary: The paper proposes to apply self-attention mechanism from (Vaswani et.al.) to the task of click-through rate prediction, which is a task where one has input features which are a concatenation of multiple one-hot vectors (referred to as fields). The paper finds that applying the self-attention mechanism outperforms state of the art approaches for the task on two benchmark datasets. It then proposes a small modification to the self-attention mechanism, retaining only the top-k attentions to sparsify attention, and finds that it leads to marginal improvements.