Alternative to torch.multinomial

parhamg: <p>Hi,</p>
<p>I was in need of performing a weighted random selection in pytorch and at the time I didn’t know about torch.multinomial, so I came up with my own implementation. I’ve since learned about torch.multinomial, but my implementation seems to be significantly faster. I haven’t done detailed performance testing, but some ad-hoc testing shows it to be consistently faster on my machine (cuda). I’d rather use my own implementation since it’s faster, but I’d like to get some feedback to make sure there isn’t an error in my implementation.</p>
<pre><code class="lang-auto">weights = ... # these are calculated somewhere else. Shape = (batchSize, n)

weights = weights.cumsum(-1)
probabilities = torch.rand((batchSize, 1)).expand(-1, probabilities.shape[-1])
delta = probabilities - weights

# I suspect there is a faster implementation for this op
# Assumption: Max value of cumsum of probabilities *should* be 1.0.
delta = torch.where(delta &lt; 0, 2.0, delta)

choices = torch.argmin(delta, -1)
</code></pre>
<p>I will also give an example in plaintext:</p>
<p>Say we have probabilities = [0.3, 0.1, 0.4, 0.1, 0.1]<br>
cumsum(probabilities) = [0.3, 0.4, 0.8, 0.9, 1.0]</p>
<p>Say we choose a random number r =  0.4<br>
delta = r - cumsum(probabilities) = [0.1, 0, -0.4, -0.5, -0.6]</p>
<p>Then we account for the negatives:<br>
delta = delta &lt; 0 ? 2.0 : delta = [0.1, 0, 2.0, 2.0, 2.0]</p>
<p>Finally, we get the argmin of delta, which is where the value is 0 at index 1.</p>

KFrank: <p>Hi Parham!</p>
<aside class="quote no-group" data-username="parhamg" data-post="1" data-topic="190685" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/parhamg/48/64974_2.png" class="avatar"> parhamg:</div>
<blockquote>
<p>Say we have probabilities = [0.3, 0.1, 0.4, 0.1, 0.1]<br>
cumsum(probabilities) = [0.3, 0.4, 0.8, 0.9, 1.0]</p>
<p>Say we choose a random number r =  0.4<br>
delta = r - cumsum(probabilities) = [0.1, 0, -0.4, -0.5, -0.6]</p>
<p>Then we account for the negatives:<br>
delta = delta &lt; 0 ? 2.0 : delta = [0.1, 0, 2.0, 2.0, 2.0]</p>
<p>Finally, we get the argmin of delta, which is where the value is 0 at index 1.</p>
</blockquote>
</aside>
<p>Your algorithm will never generate a <code>choices</code> value equal to the largest<br>
desired value of <code>n - 1</code>.  (In your example case, <code>n - 1 = 4</code>.)</p>
<p>In terms of your example, <code>torch.rand()</code> never generates a result greater<br>
than <em>or equal to</em> <code>1.0</code>.  If <code>rand()</code> generates a value between <code>0.9</code> and <code>1.0</code>,<br>
<code>delta[4]</code> will initially be negative (say, <code>0.95 - 1.0</code>) and will then be replaced<br>
by <code>2.0</code>.  So the <code>argmin()</code> that gives you the <code>choices</code> value will be <code>3</code> (and<br>
never <code>4</code>).</p>
<p>If you wish to follow up, please post a fully-self-contained, runnable script<br>
with your proposed algorithm and a test case or two, together with the output<br>
you get when you run it.</p>
<p>Best.</p>
<p>K. Frank</p>