top_p returns wrong values and re-orders the data

from routing_transformer.autoregressive_wrapper import top_p, top_k
import torch
import torch.nn.functional as F

test_tensor = torch.tensor([[0.1, 0.2, 0.15, 0.4, 0.3, 0.001, 0.01]])
threshold=0.3

print(test_tensor)
print(F.softmax(test_tensor, dim=-1))

print("Top K")
print(top_k(test_tensor, thres=threshold))
print(F.softmax(top_k(test_tensor, thres=threshold), dim=-1))

print("Top P")
print(top_p(test_tensor, thres=threshold))
print(F.softmax(top_p(test_tensor, thres=threshold), dim=-1))

Output:

tensor([[0.1000, 0.2000, 0.1500, 0.4000, 0.3000, 0.0010, 0.0100]])
tensor([[0.1325, 0.1464, 0.1393, 0.1789, 0.1618, 0.1200, 0.1211]])
Top K
tensor([[  -inf, 0.2000, 0.1500, 0.4000, 0.3000,   -inf,   -inf]])
tensor([[0.0000, 0.2338, 0.2224, 0.2855, 0.2584, 0.0000, 0.0000]])
Top P
tensor([[  -inf,   -inf, 0.3000,   -inf, 0.4000,   -inf,   -inf]])
tensor([[0.0000, 0.0000, 0.4750, 0.0000, 0.5250, 0.0000, 0.0000]])

Thanks for writing this library! I think there is a bug in top_p, with two symptoms:

1. The wrong results are filtered. It defines the threshold in the opposite way as top_k. So setting thres=0.9 results in everything being returned until a cumulative probability of 0.1 is reached.
2. The results themselves are gathered twice (instead of gathered and scattered) and as a result the returned tensor’s values are distributed in a nearly random order.

I’ll send over a PR in a few, hope it’s helpful!

[lucidrains]

@tomweingarten Hi Tom, just wanted to let you know that I’ve made an important update in the latest version per the author’s suggestion

[tomweingarten]

Very well so far! Only just beginning to run tests to compare it to XLNet but it seems very promising so far.