Large neural recommendation models can be a challenge to deploy at scale. For recommendation services with a large number of users, the most powerful models may require an impractical amount of space to store the large dense vectors encoding each of the users' tastes. Combining ideas from auto-encoder-based recommender systems, neural discrete representation learning (VQ-VAE), and product quantization (PQ), we propose PQ-VAE, a recommendation model that learns compact, discrete embeddings at only a small cost in accuracy.
dimension, and P is the precision in bits (typically 32). With tens of millions of users, and sometimes multiple vectors per user, that could mean upward of 100G of embeddings.
In this work, we propose PQ-VAE: a neural recommendation model that uses quantization to learn
compact user embeddings at only a small accuracy cost. Our approach builds on three ideas: the
auto-encoder recommenders proposed in [
Experiments in [
We combine this model with discrete representations. VQ-VAE, the model proposed in [
The third innovation on which we built PQ-VAE is product quantization [
PQ-VAE
PQ-VAE, like the state-of-the-art DAE recommender from [
Concretely, the vector of activations corresponding to each example is cut into chunks of around
length 10, and each of these chunks is quantized using a learned codebook. This presents a few
challenges at training time. First, the quantization operation is not diferentiable, making it impossible
to optimize the above model exactly via gradient descent. As in [
In this paradigm, PQ-VAE quantizes queries as
opposed to (more usually in PQ) the candidate
vectors. Therefore, we cannot use PQ-based
nearest neighbor search directly. We found that
is possible to reduce the complexity of the
similarity computation between a query and a
collection of candidate vectors regardless of which
of the two are quantized. However, eficient
distance computations on quantized user
embeddings require the similarity scores between the
query codebook and all candidates to be
precomputed, trading memory footprint for speed.
step simply never happened.1 Second, we would like to learn codes eficiently as part of model training.
To this end, we use exponential moving average k-means to learn codes in a batch-wise fashion, again
following [
In a trained PQ-VAE, the quantized representations in the last hidden layer represent the user, as illustrated in Figure 1. When these user representations are frozen for serving, we can convert them to discrete vectors, allowing us to store them with just 10–100 bytes each, 1-2 orders of magnitudes less that the the embeddings learned by a large DAE.
Finally, the PQ structure of the embeddings alternatively allows for eficient retrieval. At prediction
time, the PQ structure of the embeddings allows us to swap out the model’s final layer for PQ-based
similarity scoring, replacing most of the multiply-adds with a lookup into precomputed scores.2
However, we noticed that without a highly optimized PQ implementation (see e.g. [
USER ITEM INTERACTIONS USER ITEM PREDICTIONS DROPOUT
DENSE +
TANH
DENSE +
TANH
DENSE +
TANH
QUANTIZE CHUNKS
DENSE +
SOFTMAX ranking performance drops by 19%. Preliminary results on the larger internal dataset are consistent with these findings: a 40× user vector compression can be achieved with a relative performance drop of less than 5%.
We show that is possible to learn high-quality, memory-eficient user representations by implementing product quantization inside a neural recommender. In future work we hope to further leverage the discrete nature of the learned user representations to speed up prediction.