Understanding the softmax output in Youtube's recommender 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?
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.