Inside Autoregressive Sampling: Softmax, Logits, Temperature, and Decoding Strategies

Every response generated by a large language model—whether composing an essay, generating a formal proof, or outputting a JSON function call—is the product of a simple, iterative loop: autoregressive token prediction. At time step $t$, given all previous tokens $x_{1:t}$, the neural network predicts the conditional probability distribution over the entire vocabulary:

However, the model itself does not output words or characters. It outputs raw, unnormalized floating-point numbers known as logits. How an inference engine translates those logits into a single discrete token governs whether the output is mathematically rigorous, creatively fluent, or incomprehensibly hallucinated.

1. The Unembedding Projection & Logits

At the final layer of a Transformer, the model state is represented as a dense hidden activation vector $\mathbf{h}_t \in \mathbb{R}^{d_{\text{model}}}$ (where $d_{\text{model}}$ is typically 4,096 in a 7B model or 8,192 in a 70B model).

To convert this geometric vector into token candidates, the model multiplies $\mathbf{h}_t$ by the unembedding weight matrix $W_u \in \mathbb{R}^{|V| \times d_{\text{model}}}$, where $|V|$ is the total vocabulary size (e.g., 128,256 tokens in Llama 3):

The resulting vector $\mathbf{z} = [z_1, z_2, \dots, z_{|V|}]$ contains the logits. A logit $z_i$ represents the unnormalized log-odds that token $w_i$ should be the next token. Because logits range from $-\infty$ to $+\infty$, they cannot be sampled directly.

2. Softmax & Temperature Scaling: The Entropy Thermostat

To map raw logits into a valid probability distribution where $\sum_{i} P(w_i) = 1$ and $0 \le P(w_i) \le 1$, we apply the Softmax function with Temperature scaling ($T > 0$):

The Mathematical Limiting Cases of Temperature ($T$)

• As $T \to 0$ (Greedy Argmax): The gap between the largest logit and all competitors approaches infinity. The probability of the single top token collapses to $1.0$, while all other tokens drop to $0.0$. Output becomes 100% deterministic and reproducible.
• At $T = 1.0$ (Standard Training Distribution): Tokens are sampled exactly according to the model's learned posterior distribution.
• As $T \to \infty$ (Uniform Entropy): Every scaled logit $z_i / T$ approaches $0$. Since $\exp(0) = 1$, the distribution flattens into a uniform random distribution where every token has probability $1 / |V|$, resulting in pure gibberish.

3. Truncation Strategies: Top-K and Top-P (Nucleus Sampling)

Even at moderate temperatures ($T=0.7$), the "long tail" of the vocabulary contains thousands of improbable, irrelevant tokens whose tiny probabilities sum up to a noticeable chance of being selected. If an inference engine accidentally samples a token from this tail, the model can enter a catastrophic failure spiral.

Top-K Truncation

Top-K restricts the candidate pool to strictly the top $K$ tokens with the highest logits, setting all other logits to $-\infty$:

Limitation: When the distribution is highly peaked (e.g. The capital of France is [Paris]), Top-K still forces the model to consider the 49 inferior candidates when $K=50$. When the distribution is flat, $K=50$ might prematurely truncate valid options.

Top-P (Nucleus) Sampling (Holtzman et al., 2019)

Top-P solves this by dynamically adapting the candidate pool size based on certainty. It selects the smallest set of tokens whose cumulative probability equals or exceeds threshold $P$ (typically $0.90$ or $0.95$):

In high-confidence contexts, $V^{(P)}$ might shrink to a single token (e.g. "Paris"). In ambiguous open-ended contexts, $V^{(P)}$ automatically expands to hundreds of plausible candidate tokens.

4. Frequency & Presence Penalties

Autoregressive models frequently fall into repetitive attractor loops (e.g., repeating the same sentence structure or list items). Penalties modify the logits before Softmax based on token history:

Where $c_i$ is the count of times token $i$ has already appeared in the generated sequence, $\alpha_{\text{freq}}$ is the frequency penalty scalar, and $\alpha_{\text{pres}}$ is the presence penalty scalar.

5. Speculative Decoding Acceleration

Because decoding generates one token per forward pass, inference speed is strictly memory-bandwidth bound. Speculative Decoding (Leviathan et al., 2023) breaks this sequential bottleneck without changing the output distribution by a single bit:

# The Speculative Decoding Protocol
1. Draft Phase: A small, fast "draft" model (e.g., Llama-3-8B) generates
   K candidate tokens sequentially: [x1, x2, x3, x4]. (Very low latency).

2. Verification Phase: The large "target" model (e.g., Llama-3-70B) runs
   a SINGLE parallel forward pass over all K tokens simultaneously.

3. Rejection Sampling: The target model verifies each draft token.
   Tokens matching target distribution are accepted; upon first rejection,
   the target model emits its own corrected token and discards the remainder.

In typical code generation tasks where syntax is highly predictable, speculative decoding yields a 2x to 3.5x wall-clock speedup on the exact same server hardware.

6. Complete PyTorch Implementation

The following self-contained PyTorch routine demonstrates the complete sampling pipeline incorporating Temperature, Top-K, and Top-P filtering:

import torch
import torch.nn.functional as F

def sample_next_token(logits: torch.Tensor, 
                      temperature: float = 0.7, 
                      top_k: int = 50, 
                      top_p: float = 0.9) -> int:
    """
    logits: 1D Tensor of raw unnormalized logits [vocab_size]
    """
    # 1. Deterministic Greedy Fallback
    if temperature <= 1e-4:
        return int(torch.argmax(logits).item())

    # 2. Temperature Scaling
    scaled_logits = logits / temperature

    # 3. Top-K Filter
    if top_k > 0:
        indices_to_remove = scaled_logits < torch.topk(scaled_logits, top_k)[0][..., -1, None]
        scaled_logits[indices_to_remove] = -float('Inf')

    # 4. Top-P (Nucleus) Filter
    if 0.0 < top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(scaled_logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

        # Remove tokens with cumulative probability above the threshold
        sorted_indices_to_remove = cumulative_probs > top_p
        # Shift the indices to keep the first token above the threshold
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        scaled_logits[indices_to_remove] = -float('Inf')

    # 5. Softmax & Categorical Multinomial Sampling
    probabilities = F.softmax(scaled_logits, dim=-1)
    sampled_index = torch.multinomial(probabilities, num_samples=1)
    return int(sampled_index.item())

7. Production Parameter Guidelines