Tensorflow: how to look up and average a different amount of embedding vectors per training instance, with multiple training instances per minibatch?

In a recommender system setting: let's say I want to learn to predict future item purchases based on user past purchases using an approach inspired by Youtube's recommender system:

Concretely, let's say I have a trainable content-based network that receives as input an item and, based on its content, returns an embedding for such item. Now, let's say each user has purchased a variable number of items in the past (some users might have purchased 5 items, others maybe 1, others maybe 10, some outliers maybe 100, etc.). I want to generate a user vector, a candidate item vector and then a user-item match score as follows:

1. I map each item purchased by that user to its embedded item vector using the trainable content-based network
2. I calculate the average of all those embedded item vectors (as illustrated in the picture)
3. I apply a couple of ReLu layers on top of this average, thus obtaining a user vector
4. I map a candidate item (to be recommended) to its embedded item vector using the same trainable content-based network of step 1 (the weights of this network are always shared, like a Siamese network so to speak)
5. Finally, I compute the dot product between the user vector and the candidate item vector, apply a cross entropy loss during training, etc.

So my question is about the technical details of how to implement the embedding lookup and average of a variable number of embedded item vectors per user using Tensorflow, considering that during training each mini-batch may contain many training instances, where each training instance possibly consists of a different user with a different amount of purchased items in the past. Although the context is different, my question is very similar to this one, but unfortunately nobody has answered that question up to now.

1 Answer

Use tf.gather().

Single instance case

In the example below, we selected a variable number of embedding vectors from the matrix embedding. The selection indexing vector user can be of variable length. Then we calculate the average embedding.


with tf.Graph().as_default():
embedding = tf.placeholder(shape=[10,3], dtype=tf.float32)
user = tf.placeholder(shape=None, dtype=tf.int32)
selected = tf.gather(embedding, user)
average = tf.reduce_mean(selected, axis=0)

with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
embedding_ = np.random.randn(10,3)
user_ = [1,3,5]
print(sess.run(average, feed_dict={embedding:embedding_, user:user_}))
print(np.mean(embedding_[user_], axis=0))


Multiple instances in a mini-batch

You may manually specify the first vector in embedding to be a zero vector, and patch the above selection vector with 0s. For example

with tf.Graph().as_default():
embedding = tf.placeholder(shape=[10,3], dtype=tf.float32)
user = tf.placeholder(shape=[None, None], dtype=tf.int32)
selected = tf.gather(embedding, user)
non_zero_count =  tf.cast(tf.count_nonzero(user, axis=1), tf.float32)
embedding_sum = tf.reduce_sum(selected, axis=1)
average = embedding_sum / tf.expand_dims(non_zero_count, axis=1)

with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
embedding_ = np.concatenate([np.zeros((1,3)),np.random.randn(9,3)], axis=0)
user_ = [[3,5,7,0], [1,2,0,0]]
print(sess.run(average, feed_dict={embedding:embedding_, user:user_}))
print(np.sum([embedding_[i] for i in user_], axis=1) / np.atleast_2d(np.count_nonzero(user_, axis=1)).T)


You can use tf.gather() like this even if the embedding is a trainable variable instead of a placeholder.

• Oh great, patching with 0s should work. I have a question though, what if some training instances are very large (say, an outlier user who watched a lot of videos or purchased a lot of items), would there be important performance penalties if I force all training instances to be as large as the outlier one? Commented Sep 28, 2018 at 0:29
• In terms of model performance, I don't see that would influence your model's accuracy or something. In terms of computational performance, there should be some negative influence, which might be somehow reduced by sorting you training instances by length before minibatching, which is commonly used in machine translation where sentences have varying lengths. Commented Sep 28, 2018 at 2:33
• How do you train the embedding as youtube paper do? Do you set embedding trainable or average trainable? I guess just put average above embedding, and have both of them trainable? Commented Mar 28, 2019 at 17:29
• mask_zero=True for keras.layers.Embedding does the 0 padding Commented Mar 31, 2019 at 14:47