Skip to content

[RFC] Add Top-nσ logit truncation example via custom logits processor #2

Description

@Jackie2049

RFC: Add Top-nσ Logit Truncation Example via Custom Logits Processor

Reference: Tang et al., "Top-nσ: Not All Logits Are You Need", ACL 2025. arXiv:2411.07641

Motivation

Top-nσ is a logit-space dynamic truncation method proposed by Tang et al. (ACL 2025). Unlike probability-space filters such as top_p or min_p, it applies before softmax and keeps the candidate set less sensitive to temperature changes.

This issue proposes adding a research-oriented example implementation using vLLM's existing custom logits processor API. It is not a request to add a new built-in sampling parameter.

Scope

This proposal only adds an example file demonstrating how to implement top-nσ as a custom logits processor. It:

  • Does not add a new SamplingParams field
  • Does not change any existing sampling behavior
  • Does not modify the core Sampler class

Algorithm

For each request row:

1. max_logit = logits.max()
2. std_logit = logits.std()
3. threshold = max_logit - n * std_logit
4. logits[logits < threshold] = -inf   (mask unlikely tokens)
5. Let existing vLLM sampler handle temperature, softmax, and sampling

The standard deviation directly measures the "peakiness" of the logit distribution:

  • Sharp distribution (model is certain): small std → narrow threshold → fewer candidates
  • Flat distribution (model is uncertain): large std → wide threshold → more candidates

This is based on the empirical observation (verified in Tang et al., 2025) that logit tails are approximately Gaussian, making the nσ threshold a principled percentile-based cutoff.

Usage (via custom logits processor)

from vllm import SamplingParams

sampling_params = SamplingParams(
    temperature=0.8,
    extra_args={"top_n_sigma": 2.0}
)

The custom logit processor reads top_n_sigma from extra_args and applies truncation before the standard sampling pipeline.

Safety / Edge Cases

Case Behavior
n <= 0 Reject (raise ValueError)
std == 0 (all logits equal) Keep all tokens (no filtering)
NaN/Inf logits Preserve existing behavior (no-op)
After filtering At least the argmax token is always kept (prevent all--inf)

Why This Is Useful

Existing methods don't use the spread of the logit distribution as a signal:

Strategy Signal Limitation
top_k Fixed count Ignores distribution shape entirely
top_p Cumulative probability Over-retains in bimodal distributions
min_p Proportional to max probability Poor for flat distributions

Top-nσ adapts to distribution spread, making it useful for:

  • Models with highly variable logit distributions across prompts
  • Scenarios where top_k/top_p over-filter or under-filter
  • Research comparing logit-space vs probability-space filtering

A key property verified in our experiments: Top-nσ candidate count is invariant to temperature (0.1–10.0), while top_p varies by 19,665× and min_p by 20,000×. See experiment report for details.

Comparison Example

# Logit distribution: max=10.0, std=2.0

top_k=5:           Keeps exactly 5 tokens
top_p=0.9:         Keeps tokens summing to 90% cumulative probability
min_p=0.05:        Keeps tokens with prob >= 5% of max
top_n_sigma=2.0:   Keeps tokens with logit >= 10.0 - 2*2.0 = 6.0

Deliverables (if maintainers are interested)

  1. Example file: examples/features/logits_processor/top_n_sigma.py
  2. Unit tests: Verify masking, edge cases, and compatibility with other samplers
  3. Benchmark sketch: Compare output quality vs top_p/min_p on diverse prompts
  4. Documentation: One-paragraph note in sampling docs

Minimal Prototype (for initial feedback)

Before investing in a full PR, I'll share a standalone prototype script that demonstrates the algorithm. If maintainers find it interesting, I'll follow up with the full example, tests, and docs.


References: Tang et al., "Top-nσ: Not All Logits Are You Need", ACL 2025. arXiv:2411.07641

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions