PyTorch Sampling from Logits
Sampling from logits picks the next class by drawing from a model's raw scores. Greedy, temperature, and top-k are three controls that shape that pick.
Greedy takes the top score
Logits are raw scores, one per class, before any softmax. Greedy sampling always takes the index of the largest score.
logits = torch.tensor([2.0, 1.0, 0.5, 0.0, -1.0])
print("Logits:", logits.tolist())
print("Greedy index:", logits.argmax().item())Temperature reshapes the odds
Temperature divides every logit before the softmax. A large temperature flattens the odds, and a small one sharpens them.
for temp in (2.0, 0.5):
probs = torch.softmax(logits / temp, dim=0)
print(temp, [round(v, 3) for v in probs.tolist()])Top-k cuts the tail
Top-k keeps only the k largest logits. The rest become minus infinity, so the softmax gives them zero odds.
values, index = torch.topk(logits, 2)
print("Kept values:", values.tolist())
print("Kept index:", index.tolist())
kept = torch.full_like(logits, float("-inf"))
kept[index] = values
print("Masked logits:", kept.tolist())Example
import torch
from torch.distributions import Categorical
logits = torch.tensor([2.0, 1.0, 0.5, 0.0, -1.0])
names = ["alpha", "beta", "gamma", "delta", "omega"]
greedy = logits.argmax().item()
torch.manual_seed(22)
warm = Categorical(logits=logits / 0.8).sample().item()
torch.manual_seed(22)
top_values, top_index = torch.topk(logits, 2)
top_pick_pos = Categorical(logits=top_values).sample().item()
top_k = top_index[top_pick_pos].item()
print("Greedy pick:", greedy, names[greedy])
print("Temperature 0.8 pick:", warm, names[warm])
print("Top-2 pick:", top_k, names[top_k])Key points
- Logits: the raw scores a model gives each class
- Greedy: always pick the class with the largest logit
- Temperature: a divisor on the logits that sharpens or flattens the odds
- Top-k: keep the k largest logits and cut the rest before sampling
- Use Categorical(logits=...) to draw one class from a row of logits
- All three controls reshape the same odds before one class is picked
In this example, one row of logits names five classes, alpha through omega. Greedy picks index 0, the largest logit, with no randomness at all. The temperature draw divides the logits by 0.8 under seed 22 and picks index 2, gamma. The top-2 draw keeps only the two largest logits under the same seed. It picks index 1, beta, one of the two classes the cut left open.
Note
Greedy ignores the seed, because it never draws a random number
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
Temperature 2.0 is applied to logits before sampling. What happens to the odds?