Decoding Strategies: Beam Search vs Greedy vs Sampling

Overview

Decoding strategies determine how to select the next token during autoregressive generation. Different strategies offer different trade-offs between quality, diversity, and speed.

• Greedy: Fastest, lowest diversity, sometimes lower quality
• Beam Search: Balanced quality and diversity, slower
• Sampling: Best diversity and quality, slower than greedy
• Nucleus Sampling: Balanced quality and diversity
• Trade-off: Speed vs. generation quality

---

Greedy Decoding

How It Works

At each step, select the token with highest probability:

P(t+1 | context) = softmax(logits)[1:vocab_size]
                 = [0.3, 0.25, 0.15, 0.1, 0.05, ...]

Greedy:
  - Select argmax = token with 0.3 (highest probability)

Result sequence:
  - [The, cat, sat, on, the, mat]
   (always picks most likely next token)

Pseudocode:

def greedy_decode(model, prompt, max_length=100):
    generated = list(prompt)

    for _ in range(max_length):
        # Forward pass
        logits = model(generated)[-1, :]  # Last position

        # Select token with highest probability
        next_token = argmax(logits)
        generated.append(next_token)

    return generated

# Cost: 1 forward pass per token (required)
# Time: Fast (no overhead)
# Diversity: Very low (deterministic)
# Quality: Often acceptable for factual tasks

Advantages

✅ Fast (no overhead)
✅ Deterministic (reproducible)
✅ Works reasonably well for many tasks

Limitations

❌ No diversity (same output every time)
❌ Can get stuck in repetition loops
❌ May select suboptimal token if high-probability path is bad
❌ Doesn't explore alternatives

---

Beam Search

How It Works

Instead of tracking 1 hypothesis, track top-k:

Step 1: Initial (beam size=3)
  - Probabilities: [0.3, 0.25, 0.15, ...]
  - Keep top-3 hypotheses:
    - [token_1, prob=0.3]
    - [token_2, prob=0.25]
    - [token_3, prob=0.15]

Step 2: Expand each hypothesis
  - For [token_1]: compute next probabilities
    - [token_1 → token_A, prob=0.3×0.4=0.12]
    - [token_1 → token_B, prob=0.3×0.3=0.09]
    - [token_1 → token_C, prob=0.3×0.2=0.06]
    - [token_1 → token_D, prob=0.3×0.1=0.03]
│
  - For [token_2]: compute next probabilities
    - [token_2 → token_A, prob=0.25×0.5=0.125]
    - [token_2 → token_B, prob=0.25×0.3=0.075]
    - [token_2 → token_C, prob=0.25×0.15=0.0375]
    - [token_2 → token_D, prob=0.25×0.05=0.0125]
│
  - For [token_3]: ...

Step 3: Select top-3 globally
  - All combinations: 3 × 4 = 12 options
  - Rank by probability:
    - [token_2 → A, 0.125] ← best
    - [token_1 → A, 0.12]
    - [token_2 → B, 0.075]
  - Keep top-3, discard others

Step 4: Repeat until max length or EOS

Implementation

def beam_search(model, prompt, beam_size=3, max_length=100):
    """
    Generate with beam search

    Args:
        beam_size: Number of hypotheses to track (3-5 typical)

    Returns:
        Top-scoring sequences
    """
    vocab_size = model.config.vocab_size
    device = next(model.parameters()).device

    # Initialize beam
    # Each entry: (sequence, score)
    batch_size = 1
    sequences = [(prompt, 0.0)]  # Start with initial prompt

    for step in range(max_length):
        # Candidates for this step
        candidates = []

        for seq, seq_score in sequences:
            # Forward pass
            with torch.no_grad():
                logits = model(seq.unsqueeze(0))[0, -1, :]

            # Get log probabilities
            log_probs = F.log_softmax(logits, dim=-1)

            # Get top beam_size next tokens
            top_log_probs, top_indices = torch.topk(log_probs, beam_size)

            # Expand each hypothesis
            for i in range(beam_size):
                new_seq = torch.cat([seq, top_indices[i].unsqueeze(0)])
                new_score = seq_score + top_log_probs[i].item()
                candidates.append((new_seq, new_score))

        # Keep top beam_size candidates globally
        candidates.sort(key=lambda x: x[1], reverse=True)
        sequences = candidates[:beam_size]

    # Return best sequence
    best_seq, best_score = sequences[0]
    return best_seq

# Cost: beam_size forward passes per token
# Time: ~3-5x slower than greedy (for beam_size=3-5)
# Diversity: Limited (only explores top hypotheses)
# Quality: Better than greedy, can find better paths

Analysis

Beam search quality vs speed:

Beam size  Quality   Speed    Diversity   Use case
──────────────────────────────────────────────────
1          Low       1x       None        Baseline (greedy)
3          Medium    3x       Low         Balanced
5          Slightly  5x       Low         Deliberate tasks
10         Marginal  10x      Low         Research only

Observation:
  - Diminishing returns after beam=3
  - beam=5 rarely better than beam=3
  - Speed cost not worth quality gain
  - Default: beam_size=3 for practical use

---

Sampling

Temperature-Based Sampling

Problem with greedy: Deterministic, can't escape bad paths
Problem with beam search: Still limited to high-prob paths

Solution: Sample from distribution!

Standard greedy:
  - P(next) = [0.3, 0.25, 0.15, 0.1, ...]
  - Sample argmax = token with 0.3
  - Deterministic

Sampling:
  - P(next) = [0.3, 0.25, 0.15, 0.1, ...]
  - Sample from distribution: Can pick any token!
    - 30% chance: pick 0.3 token (most likely)
    - 25% chance: pick 0.25 token
    - 15% chance: pick 0.15 token
    - ...
  - Stochastic: different outputs each time!

Temperature control:

T = 0.5 (sharp):
  - Probabilities sharpen: [0.5, 0.3, 0.15, ...]
  - Most probable token even more likely
  - Conservative, deterministic behavior

T = 1.0 (original):
  - Probabilities unchanged: [0.3, 0.25, 0.15, ...]
  - Standard sampling

T = 2.0 (soft):
  - Probabilities flatten: [0.2, 0.2, 0.15, 0.15, ...]
  - More uniform distribution
  - Diverse, creative outputs

Implementation

def sample_decode(model, prompt, max_length=100, temperature=1.0, top_p=0.9):
    """
    Generate with sampling

    Args:
        temperature: Higher = more random
        top_p: Nucleus sampling (only sample from top p% probability)
    """
    generated = list(prompt)

    for _ in range(max_length):
        # Forward pass
        logits = model(generated)[-1, :]

        # Apply temperature
        logits = logits / temperature

        # Get probabilities
        probs = F.softmax(logits, dim=-1)

        # Nucleus sampling (optional, more stable)
        if top_p < 1.0:
            # Sort by probability
            sorted_probs, sorted_indices = torch.sort(probs, descending=True)
            cumsum_probs = torch.cumsum(sorted_probs, dim=-1)

            # Find threshold: keep top p% probability mass
            threshold_idx = torch.searchsorted(cumsum_probs, top_p)
            threshold_prob = sorted_probs[threshold_idx]

            # Zero out low-probability tokens
            probs[probs < threshold_prob] = 0
            probs = probs / probs.sum()  # Renormalize

        # Sample from distribution
        next_token = torch.multinomial(probs, num_samples=1)
        generated.append(next_token.item())

    return generated

# Cost: 1 forward pass per token (same as greedy)
# Time: Fast (no overhead over greedy)
# Diversity: High (stochastic sampling)
# Quality: Can be better or worse than greedy

---

Nucleus Sampling (Top-p)

How It Works

Pure sampling can generate low-probability nonsense:

Probabilities: [0.4, 0.3, 0.2, 0.05, 0.03, 0.01, 0.01, ...]
                                         ^ These are terrible but still possible!

Nucleus sampling (top_p=0.9):
  - Find cumulative probability threshold
  - Keep tokens until cumsum reaches top_p
    - 0.4 (sum: 0.4)
    - 0.3 (sum: 0.7)
    - 0.2 (sum: 0.9) ← STOP HERE
    - Drop 0.05, 0.03, 0.01, 0.01 (too low)
  - Sample only from kept tokens
  - Result: Probability = [0.4, 0.3, 0.2, 0, 0, 0, ...]

Effect:
  - High probability tokens: always available
  - Medium probability: available if top_p allows
  - Low probability: removed (prevents nonsense)
  - Balances diversity and quality!

Quality

Experiment: LLaMA 7B generation quality

Metric                Greedy  Beam(3)  Sample  Nucleus
─────────────────────────────────────────────────────
ROUGE (factual)       34.2    35.1     33.8    34.8
Diversity (TTR)       1.2     1.5      8.4     5.2
Human rating          3.2/5   3.5/5    3.8/5   3.9/5
Generation time       1x      3x       1x      1x

TTR (Type-Token Ratio): Higher = more diverse vocabulary

Observations:
  - Greedy: Fast, boring, decent quality
  - Beam: Slightly better, slow
  - Sample: Diverse, creative, sometimes weird
  - Nucleus: Best balance of diversity and quality
  - Recommendation: Use nucleus sampling (top_p=0.9)

---

Comparison and Recommendations

Which to Use?

Task                  Recommendation
──────────────────────────────────────
Code generation       Greedy (most consistent)
Factual QA            Beam search (beam_size=3)
Creative writing      Nucleus sampling (top_p=0.9)
Chatbot               Nucleus sampling (top_p=0.9)
Machine translation   Beam search (beam_size=5)
Summarization         Nucleus sampling (top_p=0.9)

Parameter Guidelines

Greedy:
  - No parameters (deterministic)

Beam Search:
  - beam_size: 3-5 (default 3)
  - length_penalty: 1.0 (no penalty)
    - > 1.0 favors longer sequences
    - < 1.0 favors shorter sequences
  - Note: Much slower, marginal quality gain over beam_size=1

Temperature Sampling:
  - T = 0.5-0.7: More deterministic, factual
  - T = 1.0: Neutral
  - T = 1.5-2.0: Creative, diverse
  - Usually use T=0.7-0.9

Nucleus Sampling:
  - top_p = 0.8-0.95: Typical range
  - 0.95: More exploratory
  - 0.85: Balanced
  - 0.75: More conservative
  - Default: 0.9 (good balance)

Combined: Temperature + Nucleus
  - Use together for best results!
  - Nucleus removes low-prob tokens
  - Temperature adjusts prob distribution

---

Key Takeaways

🎯 Greedy: Fast, deterministic, limited quality
📊 Beam Search: Better quality, 3-5x slower
🎲 Sampling: Diverse, creative, same speed as greedy
🎪 Nucleus: Best quality-diversity trade-off
⚙️ Use Nucleus Sampling (top_p=0.9) by default

---

Related Notes

• Speculative Decoding - Optimization technique for sampling
• Medusa (Multi Head Decoding) - Parallel token prediction
• Llm Inference Optimization - Complete inference stack