SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Benjamin Coleman*
Google DeepMind  
emanben@google.com
Wang-Cheng Kang*
Google DeepMind  
wckang@google.com
Matthew Fahrbach
Google Research fahrbach@google.com
Ruoxi Wang
Google DeepMind ruoxi@google.com
Lichan Hong
Google DeepMind  
lichan@google.com
Ed H. Chi
Google DeepMind edchi@google.com
Derek Zhiyuan Cheng
Google DeepMind  
zcheng@google.com
Learning high-quality feature embeddings efficiently and effectively is critical for the performance of web-scale machine learning systems. A typical model ingests hundreds of features with vocabularies on the order of millions to billions of tokens. The standard approach is to represent each feature value as a  $d$ -dimensional embedding, introducing hundreds of billions of parameters for extremely high-cardinality features. This bottleneck has led to substantial progress in alternative embedding algorithms. Many of these methods, however, make the assumption that each feature uses an independent embedding table. This work introduces a simple yet highly effective framework, Feature Multiplexing, where one single representation space is used across many different categorical features. Our theoretical and empirical analysis reveals that multiplexed embeddings can be decomposed into components from each constituent feature, allowing models to distinguish between features. We show that multiplexed representations lead to Pareto-optimal parameter-accuracy tradeoffs for three public benchmark datasets. Further, we propose a highly practical approach called Unified Embedding with three major benefits: simplified feature configuration, strong adaptation to dynamic data distributions, and compatibility with modern hardware. Unified embedding gives significant improvements in offline and online metrics compared to highly competitive baselines across five web-scale search, ads, and recommender systems, where it serves billions of users across the world in industry-leading products.
There have been many new and exciting breakthroughs across academia and industry towards large-scale models recently. The advent of larger datasets coupled with increasingly powerful accelerators has enabled machine learning (ML) researchers and practitioners to significantly scale up models. This scale has yielded surprising emergent capabilities, such as human-like conversation and reasoning (Bubeck et al., 2023; Wei et al., 2022a). Interestingly, large-scale model architectures are headed towards two extreme scenarios. In natural language understanding, speech, and computer vision, the transformer (Vaswani et al., 2017) dominates the Pareto frontier. Most transformer parameters are
located in the hidden layers, with very small embedding tables (e.g.,  $\sim 0.63\mathrm{B}$  embedding parameters in the 175B GPT-3 (Brown et al., 2020)). On the other hand, search, ads, and recommendation (SAR) systems require staggering scale for state-of-the-art performance (e.g., 12T parameters in Mudigere et al. (2022)), but here, most of the parameters are in the embedding tables. Typical downstream hidden layers are 3 to 4 orders of magnitude smaller than the embedding tables. Recent work blurs the boundary between these extremes—large transformers started adopting bigger embedding tables for new types of tokens, and new SAR models leverage deeper and ever more-complicated networks. Therefore, embedding learning is a core technique that we only expect to become more critical and relevant to large-scale models in the future.
Motivated by web-scale ML for SAR systems, we study the problem of learning embeddings for categorical (sparse) features. The weekend golf warriors in the ML community may appreciate our analogy between embeddings and the golf club grip (Hogan and Wind, 1985). In golf, the quality of your grip (feature embeddings) dictates the quality of your swing (model performance). Feature embeddings are the only point of contact between the data and the model, so it is critical to learn high-quality representations.
Feature embedding learning has been a highly active research area for nearly three decades (Rumelhart et al., 1986; Bengio et al., 2000; Mikolov et al., 2013; Papyan et al., 2020). In the linear model era, sparse features were represented as one-hot vectors. Linear sketches for dimension reduction became a popular way to avoid learning coefficients for each coordinate (Weinberger et al., 2009). In deep learning models, sparse features are not sketched but rather represented directly as embeddings. Learned feature embeddings carry a significant amount of semantic information, providing avenues to meaningful model interactions in the high-dimensional latent space.
In practice The simplest transformation is to represent each feature value by a row in an  $N\times d$  matrix (i.e., an embedding table). This approach is sufficient for many language tasks, but SAR features often have a massive vocabulary (e.g., tens of billions of product IDs in commerce) and highly dynamic power-law distributions. It is impractical to store and serve the full embedding table.
The "hashing trick" (Moody, 1988; Weinberger et al., 2009) provides a workaround by randomly assigning feature values to one of  $M$  rows using a hash function. Hash embeddings (?) are a variation on the hashing trick where each feature value is assigned to multiple rows in the table. The final embedding is a weighted sum of rows, similar to a learned form of two-choice hashing (Mitzenmacher, 2001). Compositional embeddings (Shi et al., 2020) look up the feature in several independent tables and construct the final embedding from the components by concatenation, addition, or element-wise product. HashedNet embeddings (Chen et al., 2015) independently look up each dimension of the embedding in a flattened parameter space. ROBE embeddings (Desai et al., 2022) look up chunks, instead of single dimensions, to improve cache efficiency. Tsang and Ahle (2022) learn a sparse matrix to describe the row/dimension lookups, and deep hash embeddings (Kang et al., 2021) use a neural network to directly output the embedding.
In theory The analysis of embedding algorithms is heavily influenced by linear models and the work of Weinberger et al. (2009) on the hashing trick. This line of reasoning starts with the assumption that all inputs have a "true" representation in a high-dimensional space that preserves all of the properties needed for learning (e.g., the one-hot encoding). The next step is to perform linear dimension reduction to compress the ground-truth representation into a smaller embedding, while approximately preserving inner products and norms. The final step is to extend the inner product argument to models that only interact with their inputs through inner products and norms (e.g., kernel methods).
In this framework, a key opportunity for improving performance is to reduce the dimension-reduction error. Early methods were based on the Johnson-Lindenstrauss lemma and sparse linear projections (Achlioptas, 2003; Li et al., 2006; Weinberger et al., 2009). Since then, a substantial effort has gone into improving the hash functions (Dahlgaard et al., 2017) and tightening the theoretical analysis of feature hashing (Dasgupta and Gupta, 2003; Freksen et al., 2018). Learned embeddings are now state-of-the-art, but the focus on dimension reduction remains: most recent algorithms are motivated by dimension reduction arguments. Previous attempts to understand hash embeddings do not study the full give-and-take relationship between feature encoding and the downstream learning problem. In this work, we carry the analysis forward to better understand learned embeddings.
<sup>1</sup>For simplicity, we call this problem "feature embedding learning," but we are primarily concerned with representing sparse features such as ZIP codes, word tokens, and product IDs.
Contributions We present an embedding learning framework called Feature Multiplexing, where multiple features share one representation space. Our theoretical analysis and experimental studies on public benchmark datasets show that Feature Multiplexing is theoretically sound and generalizable to many popular embedding methods. We also propose a highly-efficient multiplexed approach called Unified Embedding, which is battle-tested and serves billions of users in industry-leading products via various web-scale search, ads, and recommendation models.
1. Feature multiplexing: We propose the Feature Multiplexing framework, where a single representation space (embedding table) is used to support different sparse features. All of the SOTA benchmark embedding algorithms in our experiments show better results when adopting Feature Multiplexing. Furthermore, the Unified Embedding approach (i.e., multiplexing + hashing trick) outperforms many SOTA embedding methods that are noticeably more complicated to implement and impractical for low-latency serving systems on modern day hardware.
2. Analysis beyond collision counts: We analyze hash embeddings in the standard framework as well as in a supervised learning setting to capture interactions between feature embedding and downstream learning. Our analysis gives insights and testable predictions about unified embeddings that cannot be obtained by dimension reduction arguments alone.
3. Battle-tested approach and industrial experiments at scale: Multi-size Unified Embedding has been launched in over a dozen web-scale SAR systems, significantly improving user engagement metrics and key business metrics. In addition, Unified Embedding presents major practical benefits: simplified feature configuration, strong adaptation to dynamic data distributions, and compatibility with ML accelerators. In Section 5.2, we share our insights, challenges, and practical lessons on adopting Unified Embedding.
SAR applications often present supervised learning problems over categorical (sparse) features that represent properties of users, queries, and items. Each categorical feature consists of one or more feature values  $v$  drawn from the vocabulary of the feature  $V$ . For a concrete example, consider the classic click-through rate (CTR) prediction task for ads systems, where we want to predict the probability of a user clicking on an advertisement. The "ad_id" feature may take on hundreds of billions of possible values—one for each unique ad served by the platform—while "site_id" describes the website where that ad was shown (a.k.a., the publisher).
We embed the values of each feature into a space that is more suitable for modeling, e.g.,  $\mathbb{R}^d$  or the unit  $(d - 1)$ -sphere. Given a vocabulary of  $N$  tokens  $V = \{v_{1}, v_{2}, \ldots, v_{N}\}$ , we learn a transformation that maps  $v \in V$  to an embedding vector  $\mathbf{e} \in \mathbb{R}^d$ . After repeating this process for each categorical feature, we get a set of embeddings that are then concatenated and fed as input to a neural network (Covington et al., 2016; Anil et al., 2022).
Problem statement We are given a dataset  $D = \{(\pmb{x}_1, y_1), (\pmb{x}_2, y_2), \dots, (\pmb{x}_{|D|}, y_{|D|})\}$  of examples with labels. The examples consist of values from  $T$  different categorical features with vocabularies  $\{V_1, V_2, \dots, V_T\}$ , where feature  $t$  has vocabulary  $V_t$ . Unless otherwise stated, we do not consider multivalent features and instead assume that each example  $\pmb{x}$  has one value for each categorical feature, i.e.,  $\pmb{x} = [v_1, v_2, \dots, v_T]$  where  $v_i \in V_i$ . However, this is for notational convenience only—the techniques in this paper extend to missing and multivalent feature values well.
Formally, the embedding table is a matrix  $\mathbf{E} \in \mathbb{R}^{M \times d}$  and the embedding function  $g(\mathbf{x}; \mathbf{E})$  transforms the categorical feature values into embeddings. We let  $h(v): V \to [M]$  be a 2-universal hash function (Vadhan et al., 2012) that assigns a feature value to a row index of the embedding table. We also consider a model function  $f(\mathbf{e}; \boldsymbol{\theta})$  that transforms the concatenated embeddings into a model prediction, which is penalized by a loss  $\ell$ . Using this decomposition, we define the joint feature learning problem as follows:
$$
\underset {\mathbf {E}, \boldsymbol {\theta}} {\arg \min } \mathcal {L} _ {D} (\mathbf {E}, \boldsymbol {\theta}), \quad \text {w h e r e} \quad \mathcal {L} _ {D} (\mathbf {E}, \boldsymbol {\theta}) = \sum_ {(\boldsymbol {x}, y) \in D} \ell (f (g (\boldsymbol {x}; \mathbf {E}); \boldsymbol {\theta}), y). \tag {1}
$$
We use independent hash functions  $h_t(v)$  for each feature  $t \in [T]$  so that tokens appearing in multiple vocabularies have the chance to get mapped to different locations, depending on the feature. We use the notation  $\mathbf{e}_m$  to denote the  $m$ -th row of  $\mathbf{E}$ , which means  $\mathbf{e}_{h(u)}$  is the embedding of  $u$ . Lastly, we let  $\mathbb{1}_{u,v}$  be the indicator variable that denotes a hash collision between  $u$  and  $v$ , i.e.,  $h(u) = h(v)$ .
Figure 1: Embedding methods for two categorical features. We highlight the lookup process for the first value  $v_{1}$  of each feature. Hash tables randomly share representations within each feature, while Unified Embedding shares representations across features. To implement Unified Embedding with different dimensions (multi-size or variable-length), we perform multiple lookups and concatenate the results.
3 Feature Multiplexing: The Inter-Feature Hashing Trick
It is standard practice to use a different hash table for each categorical feature. There are a few approaches that deviate from this common practice—for example, HashedNet and ROBE-Z models use parameter sharing across all features, depending on the implementation. However, conventional wisdom suggests that each categorical vocabulary benefits from having an independent representation, so shared representations are not well-studied.
Parameter reuse We present a Feature Multiplexing framework that only uses one embedding table to represent all of the feature vocabularies. In contrast, typical models in SAR systems have hundreds of embedding tables, each representing one input feature. Multiplexed embeddings allow models to reuse parameters and lookup operations, improving efficiency in terms of both space and time. This is critical for models that must fit into limited GPU/TPU memory or satisfy latency and serving cost constraints. In principle, any feature embedding scheme (e.g., hashing trick, multihash, compositional, or ROBE-Z) can be used as the shared representation for feature multiplexing. In Section 5, we experiment with multiplexed versions of six different embedding algorithms. Here, we focus on a multi-size (multi-probe) version of the multiplexed feature hashing trick (Figure 1), which we call Unified Embedding. Unified Embedding presents unique practical advantages that we describe in Section 5.2.
Tunable embedding width By using a unified table, we make the architectural assumption that all features share the same embedding dimension. However, this is often sub-optimal in practice since independent tuning of individual embedding dimensions can significantly improve model performance. For use-cases where different features require different embedding dimensions, we concatenate the results of a variable number of table lookups to obtain the final embedding. While this limits the output dimension to multiples of the table dimension, this is not a limiting constraint (and is needed by other methods such as ROBE-Z and product quantization).
Shared vocabulary It is common for categorical feature vocabularies to share tokens. For example, two features that represent character bigrams of different text fields (e.g., query and document) might have substantial overlap. Our default approach is to use a different hash function for each feature by using a different hash seed for each feature. However, in cases where many semantically-similar features are present, the performance may slightly improve by using the same hash function.
4 Analysis of Feature Multiplexing
In this section, we present theoretical arguments about feature multiplexing applied to the "multiplexed feature hashing trick" setting. We start with a brief analysis in the dimension-reduction framework, suggesting that multiplexing can optimally load balance feature vocabularies across buckets in the table. This classic framework, however, cannot be used to distinguish between collisions that happen within the vocabulary of a feature (i.e., intra-feature) and those occurring across features (i.e., inter-feature).
To understand these interactions, we also need to consider the learning task and model training. Thus, we analyze the gradient updates applied to each bucket under a simplified but representative model—binary logistic regression with learnable feature embeddings for CTR prediction. Our analysis demonstrates a novel relationship between the embedding parameters and model parameters during gradient descent. Specifically, we find that the effect of inter-feature collisions can be mitigated if the model projects different features using orthogonal weight vectors. In Section 4.2, we empirically verify this by training models on public benchmark on click-through prediction (Criteo) and finding that the weight vectors do, in fact, orthogonalize. We defer the proofs and detailed derivations in this section to Appendix A.
4.1 Parameter Efficiency in The Dimension-Reduction Framework
In the feature hashing problem, we are given two vectors  $\mathbf{x},\mathbf{y}\in \mathbb{R}^N$  that we project using a linear map  $\phi :\mathbb{R}^N\to \mathbb{R}^M$  given by the sparse  $\{\pm 1\}$  matrix in Weinberger et al. (2009). The goal is for  $\phi$  to minimally distort inner product, i.e.,  $\langle \phi (\mathbf{x}),\phi (\mathbf{y})\rangle \approx \langle \mathbf{x},\mathbf{y}\rangle$ . Since the estimation error directly depends on the variance, it is sufficient to compare hashing schemes based on the moments of their inner product estimators. We examine the variance of the hashing trick and multiplexing trick.
To analyze multiple features at once in this framework, we consider the concatenations  $\mathbf{x} = [\mathbf{x}_1,\mathbf{x}_2]$  and  $\mathbf{y} = [\mathbf{y}_1,\mathbf{y}_2]$ . The components  $\mathbf{x}_1\in \{0,1\}^{N_1}$  and  $\mathbf{x}_2\in \{0,1\}^{N_2}$  are the one-hot encodings (or bag-of-words) for the two vocabularies  $V_{1}$  and  $V_{2}$ . The standard hashing trick approximates  $\langle \mathbf{x},\mathbf{y}\rangle$  by independently projecting  $\mathbf{x}_1$  to dimension  $M_{1}$  and  $\mathbf{x}_2$  to dimension  $M_2$ . The multiplexed version uses a single projection, but into dimension  $M_1 + M_2$ . Now, we formalize the hashing trick of Weinberger et al. (2009), but with slightly revised notation. Given a set of tokens from the vocabulary, the embedding is the signed sum of value counts within each hash bucket.
Definition 4.1. Let  $h:V\to \{1,2,\ldots ,M\}$  be a 2-universal hash function and  $\xi :V\rightarrow \{-1, + 1\}$  be a sign hash function. The function  $\phi_{h,\xi}:2^{V}\to \mathbb{R}^{M}$  is defined as  $\phi_{h,\xi}(W) = \sum_{w\in W}\xi (w)\mathbf{u}_{h(w)}$  where  $\mathbf{u}_i$  is the  $i$ -th unit basis vector in  $\mathbb{R}^M$ .
The following result compares multiplexed and standard embeddings in the dimension reduction framework. Note that the randomness is over the choice of hash functions  $h(v)$  and  $\xi(v)$ . This result also directly extends to concatenating  $T$  feature embeddings.
Proposition 4.2. For any  $\mathbf{x}_1, \mathbf{y}_1 \in \{0, 1\}^{N_1}$  and  $\mathbf{x}_2, \mathbf{y}_2 \in \{0, 1\}^{N_2}$ , let  $\mathbf{x} = [\mathbf{x}_1, \mathbf{x}_2]$ ,  $\mathbf{y} = [\mathbf{y}_1, \mathbf{y}_2]$  denote their concatenations. Let  $\mu_U, \mu_H, \sigma_U^2$ , and  $\sigma_H^2$  be the mean and variance of  $\langle \phi_{h,\xi}(\mathbf{x}), \phi_{h,\xi}(\mathbf{y}) \rangle$  for multiplexed and hash encodings, respectively. Then,  $\mu_U = \mu_H = \langle \mathbf{x}, \mathbf{y} \rangle$  and
$$
\begin{array}{l} \sigma_ {U} ^ {2} = \frac {\| \mathbf {x} \| _ {2} ^ {2} \| \mathbf {y} \| _ {2} ^ {2} + \langle \mathbf {x} , \mathbf {y} \rangle^ {2} - 2 \langle \mathbf {x} , \mathbf {y} \rangle}{M _ {1} + M _ {2}}, \\ \sigma_ {H} ^ {2} = \frac {\| \mathbf {x} _ {1} \| _ {2} ^ {2} \| \mathbf {y} _ {1} \| _ {2} ^ {2} + \langle \mathbf {x} _ {1} , \mathbf {y} _ {1} \rangle^ {2} - 2 \langle \mathbf {x} _ {1} , \mathbf {y} _ {1} \rangle}{M _ {1}} + \frac {\| \mathbf {x} _ {2} \| _ {2} ^ {2} \| \mathbf {y} _ {2} \| _ {2} ^ {2} + \langle \mathbf {x} _ {2} , \mathbf {y} _ {2} \rangle^ {2} - 2 \langle \mathbf {x} _ {2} , \mathbf {y} _ {2} \rangle}{M _ {2}}. \\ \end{array}
$$
Observations Proposition 4.2 shows that multiplexed embeddings can do a good job balancing hash collisions across the parameter space. Consider the problem of estimating the inner product  $\langle \mathbf{x},\mathbf{y}\rangle$  Suppose that  $\| \mathbf{x}_1\| _2^2 = \| \mathbf{y}_1\| _2^2 = k_1$  and  $\| \mathbf{x}_2\| _2^2 = \| \mathbf{y}_2\| _2^2 = k_2$  , i.e., features 1 and 2 are multivalent with  $k_{1}$  and  $k_{2}$  values, and assume  $\mathbf{x}$  and  $\mathbf{y}$  are orthogonal, i.e.,  $\langle \mathbf{x},\mathbf{y}\rangle = 0$  . Proposition 4.2 tells us that  $\sigma_U^2 = (k_1 + k_2)^2 /(M_1 + M_2)$  and  $\sigma_H^2 = k_1^2 /M_1 + k_2^2 /M_2$  . The difference  $\sigma_H^2 -\sigma_U^2$  factors as  $(\sqrt{M_1 / M_2} k_1 - \sqrt{M_2 / M_1} k_2)^2 /(M_1 + M_2)$  , which is non-negative, implying  $\sigma_U^2\leq \sigma_H^2$
4.2 Trajectories of Embeddings During SGD for Single-Layer Neural Networks
Intuition suggests that inter-feature collisions are less problematic than intra-feature collisions. For example, if  $x \in V_1$  collides with  $y \in V_2$  but no other values collide across vocabularies, then the embeddings from other tokens in  $V_1$  and  $V_2$  are free to rotate around the shared embedding for  $x, y$  to preserve their optimal similarity relationships.
We analyze a logistic regression model with trainable embeddings for binary classification (i.e., a single-layer neural network with hashed one-hot encodings as input). This corresponds to Eq. (1) where  $f(\mathbf{z};\pmb {\theta})$  is the sigmoid function  $\sigma_{\pmb{\theta}}(\mathbf{z}) = 1 / (1 + \exp (-\langle \mathbf{z},\pmb {\theta}\rangle))$ ,  $\ell$  is the binary cross-entropy loss, and  $y\in \{0,1\}$  (e.g., click or non-click). We concatenate the embeddings of each feature, which means the input to  $f(\mathbf{z};\pmb {\theta})$  is  $\mathbf{z} = g(\mathbf{x};\mathbf{E}) = [\mathbf{e}_{h_1(\mathbf{x}_1)},\mathbf{e}_{h_2(\mathbf{x}_2)},\dots ,\mathbf{e}_{h_T(\mathbf{x}_T)}]$ , where  $\mathbf{e}_{h_t(\mathbf{x}_t)}$  is
Figure 2: Single-layer neural embedding model with per-feature weights  $\theta_{t}$  (left). Mean embedding  $\ell^2$ -norm (middle) and mean angle between all pairs of weight vectors  $\theta_{t_1}, \theta_{t_2}$  (right) as a function of table size for Criteo across all 26 categorical features. Note that the horizontal axes are in log scale.
the embedding for the  $t$ -th feature value in example  $\pmb{x}$ . We write the logistic regression weights as  $\pmb{\theta} = [\pmb{\theta}_1, \pmb{\theta}_2, \dots, \pmb{\theta}_T]$  so that embedding  $\mathbf{e}_{h_t(\pmb{x}_t)}$  for feature  $t$  is projected via  $\pmb{\theta}_t \in \mathbb{R}^M$ . We illustrate this simple but representative model architecture in Figure 2.
The following analysis holds in general for  $T$  features, but we consider the problem for two categorical features with vocabularies  $V_{1}$  and  $V_{2}$  to simplify the presentation. We write the logistic regression objective as follows in Eq. (2), using  $C_{u,v,1}$  and  $C_{u,v,0}$  to denote the number of examples for which  $\pmb{x} = [u,v]$  and  $y = 1$  and  $y = 0$ , respectively. We present a full derivation in Appendix A.2.