1

I have a huge NNs with tens of million Softmax probabilities at the last layer. For each given x, I need to randomly sample just one of the output nodes according to the probabilities. I use tf.random.multinomial, but its complexity is O(n), where n is the number of nodes at the last layer. What is the most efficient way to do it? Is it possible to do it in O(logn)?

Shan
  • 53
  • 4
  • What is the purpose of the sampling? – gorjan Aug 08 '19 at 11:12
  • I am going to have a generative model that generates words. At the last layer, I have to change the logits to probabilities using softmax and then randomely sample a single word according to the probabilities. Since making Softmax probabilities and doing random selection using those probabilities are in order of O(n), it is time consuming in my case. – Shan Aug 08 '19 at 11:39
  • I assume that you want to compute the loss that way? Then, during inference, you just want to take the word with the highest probability? – gorjan Aug 08 '19 at 11:42
  • No, imagine the last layer of the NN has 1 million nodes that give us the logits for each word in testing time(not training time!). I use Softmax to get probabilities from logits. Then I use those probabilities as a distribution and generate one word. It is possible to generate different words with the same distribution, so its different from Argmax(probabilities). – Shan Aug 08 '19 at 12:37
  • I think the common way of predict words are using word-embedding similar to word2vec. Then each word is not a one hot encoding, but a gathering of multiple output within a range (usually 300 parameters between -1 and 1). I don't think it will make sampling easier tough. But I have not seen any successful project where the prediction uses 1M outputs (a single output each words) as softmax. – KrisR89 Aug 08 '19 at 13:13
  • Thank you. If there were a function similar to Softmax that could give us the CDF instead of PDF, we could use binary search to generate the word in O(logn) because the CDF values are sorted. – Shan Aug 08 '19 at 13:26

0 Answers0