From 9f4bc9aed4d97d1567454809c0dbe97af79caef5 Mon Sep 17 00:00:00 2001 From: winklemad Date: Sun, 12 Jul 2026 12:33:59 +0530 Subject: [PATCH] Fix top-p sampling excluding the threshold-crossing token `top_p_logits` masked tokens with `cumsum <= top_p` on the inclusive cumulative probability, which drops the token that crosses the threshold. The kept nucleus therefore has cumulative mass strictly *below* top_p -- e.g. with probabilities [0.4, 0.3, 0.2, 0.1] and top_p=0.8 it keeps only [0, 1] (mass 0.70) instead of [0, 1, 2] (mass 0.90), biasing every fractional-top_p generation toward the mode. Keep a token while the cumulative mass *before* it is below top_p, so the crossing token is retained and the kept mass reaches top_p, matching the standard nucleus (Holtzman et al., 2019) definition. Added a regression test asserting the kept set reaches >= top_p. --- esm/utils/sampling.py | 10 ++++++++-- esm/utils/sampling_test.py | 16 +++++++++++++++- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/esm/utils/sampling.py b/esm/utils/sampling.py index 40c61b67..03e856f4 100644 --- a/esm/utils/sampling.py +++ b/esm/utils/sampling.py @@ -299,8 +299,14 @@ def top_p_logits(logits: torch.Tensor, top_p: float | torch.Tensor) -> torch.Ten # Sort logits in descending order and extract the mask for the top_p sorted_logits, sorted_indices = torch.sort(logits, dim=-1, descending=True) - cumsum_logits = sorted_logits.softmax(-1).cumsum(-1) - top_p_mask = cumsum_logits <= top_p[:, None] + sorted_probs = sorted_logits.softmax(-1) + cumsum_logits = sorted_probs.cumsum(-1) + # Keep the smallest set of tokens whose cumulative probability reaches top_p, + # including the token that crosses the threshold (the cumulative mass *before* + # it is still below top_p). Masking on the inclusive cumulative sum + # (`cumsum <= top_p`) drops that crossing token, leaving the kept mass below + # top_p. + top_p_mask = (cumsum_logits - sorted_probs) < top_p[:, None] # Make sure at least one token is sampled top_p_mask[:, 0] = True diff --git a/esm/utils/sampling_test.py b/esm/utils/sampling_test.py index ee2dc610..4755ce3f 100644 --- a/esm/utils/sampling_test.py +++ b/esm/utils/sampling_test.py @@ -1,7 +1,7 @@ import pytest import torch -from esm.utils.sampling import sample_logits +from esm.utils.sampling import sample_logits, top_p_logits def test_sample_logits(): @@ -35,4 +35,18 @@ def test_sample_logits(): ) +def test_top_p_logits_reaches_threshold(): + # top-p keeps the smallest set of tokens whose cumulative probability reaches + # top_p, including the token that crosses the threshold; the kept mass must be + # >= top_p, not < top_p. + probs = torch.tensor([0.4, 0.3, 0.2, 0.1]) + logits = torch.log(probs).unsqueeze(0) + for top_p, expected_kept in [(0.8, [0, 1, 2]), (0.7, [0, 1]), (0.35, [0])]: + out = top_p_logits(logits.clone(), top_p=top_p) + kept = (out.squeeze(0) > torch.finfo(out.dtype).min).nonzero().squeeze(-1) + assert kept.tolist() == expected_kept + assert probs[kept].sum().item() >= top_p - 1e-6 + + test_sample_logits() +test_top_p_logits_reaches_threshold()