Understanding the softmax output in Youtube's recommender

Question

This question has been asked before, but never (that I can see) satisfactorily answered.

I'm reading Youtube's paper on their recommender system. The system has two elements, the first of which is a DNN generating 100 "candidate" videos, which are then combined with candidates from other sources and ranked by a second DNN. In the paper they say that they treat the candidate generation problem as extreme multiclass classification with Softmax. That approach is understandable; the indices of the N highest values in the softmax output become the N predicted candidates - easy peasy.

Reading further in the paper however, I began to get confused about exactly what their network is doing. They give a figure of their network's structure:

And it's apparent from that image that there's a whole other step beyond the softmax layer, which is what I'm not understanding. There's also later in the paper the following quote:

The softmax layer outputs a multinomial distribution over the same 1M video classes with a dimension of 256 (which can be thought of as a separate output video embedding).

and also that:

Since calibrated likelihoods from the softmax output layer are not needed at serving time, the scoring problem reduces to a nearest neighbor search in the dot product space for which general purpose libraries can be use.

But I'm confused about what they're actually implementing here. The assertion that the softmax layer outputs...with a dimension of 256 implies to me that they literally have a final layer in the form Dense(256, activation='softmax') instead of Dense(n_classes..., however the presence of "class probabilities" in the figure makes that impossible, since you couldn't convert such an embedding to class probabilities.

So; I'm really confused; does anyone know how to interpret exactly what's being done here?

Answer 1

Part of the confusion comes from different prediction processes used for training and serving.

Softmax is used during training to make predictions for the 1 million video categories. The softmax makes the training of the neural networks gradients easier and works well within deep learning libraries.

Those 1 million video categories also happen to be embedded into 256-dimensional space. During serving, prediction is now a nearest neighbor search. Nearest neighbor search is faster, since it is a dot product operation and can use general purpose libraries.

Answer 2

I am not entirely sure but this is what I understand. The softmax layer is only used at training phase, to learn all the model parameters, especially the final ReLu layer and video embeddings.

At inference time we do not care about the softmax output really. The final ReLu layer has a size of 256 as you can see at Table 1 in the paper. They say that the final ReLu layer can be seen as a user vector. I do not know why, but since the learnt ReLu final layer has the same size as the video embeddings, they are mapped in the same space to compute the approximate nearest neighbours for the user vector. So they try to find the most similar video embeddings to the user vector.