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

Question

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.

Answer 1

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.