Beam Search: The Algorithm Behind LLM Text Generation

Posted by

When a Large Language Model (LLM) like GPT generates text, it doesn't simply pick the single most probable next word at every step. Doing so — known as greedy decoding — often leads to repetitive, suboptimal, or even incoherent outputs. Instead, most production-grade language models use a smarter strategy called Beam Search.

In this post, we'll break down:

  • What Beam Search is and why it exists
  • How it compares to Greedy and Exhaustive search
  • The step-by-step algorithm with a visual walkthrough
  • A full implementation in Java and C++
  • Its strengths, weaknesses, and modern variants

The Problem: Decoding from a Language Model

A language model assigns a probability to every possible next token given the preceding context:

P(token | previous tokens)

Generating a complete sequence means making a chain of these decisions. The goal is to find the sequence with the highest joint probability:

P(w1, w2, ..., wN) = P(w1) × P(w2|w1) × P(w3|w1,w2) × ...
Search Space Explosion: If your vocabulary has V words and you generate sequences of length N, there are V^N possible sequences. For V = 50,000 and N = 20, that's astronomically large. This is why we need smart search heuristics.

Approach 1: Greedy Search

Greedy search picks the single most probable token at every step.

At each step t: choose token = argmax P(token | context)

Why Greedy Fails

Consider generating a sentence where:

  • Step 1: "The" → 0.4 (highest)
  • Step 2 given "The": "cat" → 0.3 (highest)
  • Step 3 given "The cat": "sat" → 0.2 (highest)

But an alternative path exists:

  • Step 1: "A" → 0.35
  • Step 2 given "A": "beautiful" → 0.6
  • Step 3 given "A beautiful": "sunset" → 0.9

Joint probability of greedy path: 0.4 × 0.3 × 0.2 = 0.024
Joint probability of better path: 0.35 × 0.6 × 0.9 = 0.189

Greedy search misses the better path because it committed too early.

Approach 2: Exhaustive Search

Exhaustive search explores all possible sequences and picks the best one. It is optimal but computationally infeasible for any realistic vocabulary and sequence length.

Approach 3: Beam Search — The Sweet Spot

Beam Search is a heuristic best-first search that keeps track of the top-K most promising partial sequences (called the "beam") at each step, where K is the beam width.

Key Idea

  • At each step, expand every sequence in the beam by all possible next tokens.
  • Score each expanded sequence.
  • Keep only the top-K scoring sequences for the next step.
  • Repeat until all sequences hit an end token or max length.

This gives us a controllable trade-off:

  • K = 1 → Equivalent to Greedy Search
  • K = ∞ → Equivalent to Exhaustive Search
  • K = 5 to 20 → Practical sweet spot used in real systems
Comparison between search approaches

Step-by-Step Visual Walkthrough

Let's trace Beam Search with beam width K=2 and a tiny vocabulary: {A, B, C, <END>}.



Step 0 — Initialization

Beam: [ ([], score=0.0) ]

We start with one empty sequence.

Step 1 — Expand

Expand the empty sequence with all tokens:

Candidates:
  [A]     → log P(A)     = -0.9
  [B]     → log P(B)     = -0.5   ✓ keep
  [C]     → log P(C)     = -1.2
  [<END>] → log P(<END>) = -2.0

After keeping top-2:

Beam: [ ([B], -0.5), ([A], -0.9) ]

Step 2 — Expand Each Beam Member

Expand [B]:

  [B, A]     → -0.5 + -1.1 = -1.6
  [B, B]     → -0.5 + -0.4 = -0.9   ✓
  [B, C]     → -0.5 + -0.7 = -1.2   ✓
  [B, <END>] → -0.5 + -1.5 = -2.0

Expand [A]:

  [A, A]     → -0.9 + -0.8 = -1.7
  [A, B]     → -0.9 + -0.6 = -1.5
  [A, C]     → -0.9 + -0.3 = -1.2   (tied)
  [A, <END>] → -0.9 + -2.1 = -3.0

All 8 candidates are ranked, top-2 kept:

Beam: [ ([B, B], -0.9), ([B, C], -1.2) ]

This process continues until <END> tokens are generated or maximum length is reached. The sequence with the highest cumulative log-probability is the final output.

Why Log Probabilities?

Multiplying many small probabilities leads to numerical underflow. By working in log space, we convert multiplications into additions:

log(P1 × P2 × P3) = log(P1) + log(P2) + log(P3)

Higher (less negative) log-probability = better sequence.

Java Implementation

Here's a clean, self-contained Beam Search implementation in Java.

import java.util.*;

public class BeamSearch {

    static double[][] logProbTable = {
        /* step 0 */ {-0.9, -0.5, -1.2, -2.0},
        /* step 1 */ {-1.1, -0.4, -0.7, -1.5},
        /* step 2 */ {-0.8, -0.6, -0.3, -1.0},
        /* step 3 */ {-1.3, -0.9, -0.5, -0.4},
    };

    static String[] vocab = {"A", "B", "C", "<END>"};
    static int END_TOKEN = 3;

    static class Hypothesis implements Comparable<Hypothesis> {
        List<Integer> tokens;
        double score; 
        boolean finished;

        Hypothesis(List<Integer> tokens, double score, boolean finished) {
            this.tokens = new ArrayList<>(tokens);
            this.score = score;
            this.finished = finished;
        }

        @Override
        public int compareTo(Hypothesis other) {
            return Double.compare(other.score, this.score);
        }

        public String toStringSequence() {
            StringBuilder sb = new StringBuilder();
            for (int t : tokens) {
                sb.append(vocab[t]).append(" ");
            }
            return sb.toString().trim();
        }
    }

    static double[] getLogProbs(List<Integer> context) {
        int step = Math.min(context.size(), logProbTable.length - 1);
        return logProbTable[step];
    }

    static List<Hypothesis> beamSearch(int beamWidth, int maxLength) {
        List<Hypothesis> beam = new ArrayList<>();
        beam.add(new Hypothesis(new ArrayList<>(), 0.0, false));

        List<Hypothesis> completedSequences = new ArrayList<>();

        for (int step = 0; step < maxLength; step++) {
            List<Hypothesis> candidates = new ArrayList<>();

            for (Hypothesis hyp : beam) {
                if (hyp.finished) {
                    candidates.add(hyp);
                    continue;
                }

                double[] logProbs = getLogProbs(hyp.tokens);

                for (int tokenId = 0; tokenId < vocab.length; tokenId++) {
                    List<Integer> newTokens = new ArrayList<>(hyp.tokens);
                    newTokens.add(tokenId);
                    double newScore = hyp.score + logProbs[tokenId];
                    boolean isFinished = (tokenId == END_TOKEN);

                    candidates.add(new Hypothesis(newTokens, newScore, isFinished));
                }
            }

            Collections.sort(candidates);

            beam = new ArrayList<>();
            for (int i = 0; i < Math.min(beamWidth, candidates.size()); i++) {
                Hypothesis h = candidates.get(i);
                if (h.finished) {
                    completedSequences.add(h);
                } else {
                    beam.add(h);
                }
            }

            if (beam.isEmpty()) break;
        }

        completedSequences.addAll(beam);
        Collections.sort(completedSequences);
        return completedSequences;
    }

    public static void main(String[] args) {
        List<Hypothesis> results = beamSearch(2, 4);
        for (int i = 0; i < Math.min(3, results.size()); i++) {
            Hypothesis h = results.get(i);
            System.out.printf("Rank %d: %-25s score=%.4f%n", i + 1, h.toStringSequence(), h.score);
        }
    }
}

C++ Implementation

#include <iostream>
#include <vector>
#include <algorithm>
#include <string>

using namespace std;

const vector<string> vocab = {"A", "B", "C", "<END>"};
const int END_TOKEN = 3;
const int VOCAB_SIZE = 4;

const double logProbTable[4][4] = {
    {-0.9, -0.5, -1.2, -2.0},
    {-1.1, -0.4, -0.7, -1.5},
    {-0.8, -0.6, -0.3, -1.0},
    {-1.3, -0.9, -0.5, -0.4},
};

struct Hypothesis {
    vector<int> tokens;
    double score;
    bool finished;

    Hypothesis(vector<int> t, double s, bool f) : tokens(t), score(s), finished(f) {}

    bool operator<(const Hypothesis& other) const {
        return score < other.score;
    }
};

const double* getLogProbs(const vector<int>& context) {
    int step = min((int)context.size(), 3);
    return logProbTable[step];
}

vector<Hypothesis> beamSearch(int beamWidth, int maxLength) {
    vector<Hypothesis> beam;
    beam.emplace_back(vector<int>{}, 0.0, false);
    vector<Hypothesis> completed;

    for (int step = 0; step < maxLength; step++) {
        vector<Hypothesis> candidates;

        for (const Hypothesis& hyp : beam) {
            if (hyp.finished) {
                candidates.push_back(hyp);
                continue;
            }

            const double* logProbs = getLogProbs(hyp.tokens);

            for (int tokenId = 0; tokenId < VOCAB_SIZE; tokenId++) {
                vector<int> newTokens = hyp.tokens;
                newTokens.push_back(tokenId);
                double newScore = hyp.score + logProbs[tokenId];
                bool isFinished = (tokenId == END_TOKEN);
                candidates.emplace_back(newTokens, newScore, isFinished);
            }
        }

        sort(candidates.begin(), candidates.end(), [](const Hypothesis& a, const Hypothesis& b) {
            return a.score > b.score;
        });

        beam.clear();
        for (int i = 0; i < min(beamWidth, (int)candidates.size()); i++) {
            if (candidates[i].finished) {
                completed.push_back(candidates[i]);
            } else {
                beam.push_back(candidates[i]);
            }
        }

        if (beam.empty()) break;
    }

    for (auto& h : beam) completed.push_back(h);
    sort(completed.begin(), completed.end(), [](const Hypothesis& a, const Hypothesis& b) {
        return a.score > b.score;
    });

    return completed;
}

Complexity Analysis

Approach Time Complexity Space Complexity Quality
Greedy Search O(N × V) O(N) Poor
Beam Search (width K) O(N × K × V) O(K × N) Good
Exhaustive Search O(V^N) O(V^N) Optimal

Where: N = sequence length, V = vocabulary size, and K = beam width. Beam Search is linear in both N and K, making it extremely practical.

Length Normalization

A critical issue with raw log-probability scoring: longer sequences always have lower (more negative) cumulative scores because we're adding more negative numbers. This biases beam search toward shorter sequences.

The fix is length normalization:

normalized_score = cumulative_log_prob / (sequence_length ^ α)

Where α (typically 0.6–0.8) controls how much to penalize/reward length.

// Length-normalized score calculation
double normalizedScore(double cumulativeLogProb, int length, double alpha) {
    double lengthPenalty = Math.pow(length, alpha);
    return cumulativeLogProb / lengthPenalty;
}

Beam Search in Real LLMs

Here's how Beam Search fits into the actual text generation pipeline:



Limitations of Beam Search

  • Exposure Bias: Mismatch between training (ground-truth context) and inference (generated context).
  • Generic / Boring Outputs: Prefers safer, common phrases over creative paths.
  • Repetition: Can looped word patterns ("the the the...") without n-gram repetition penalties.
  • Not Globally Optimal: Being a heuristic with finite K, it can still miss the true optimal sequence.

Modern Alternatives and Extensions

Method Description Best For
Greedy Pick argmax at each step Speed-critical, draft generation
Beam Search Keep top-K beams Machine translation, summarization
Top-K Sampling Sample from top-K tokens Creative text generation
Top-P (Nucleus) Sampling Sample from top-P probability mass Chatbots, story generation
Temperature Scaling Sharpen/flatten probability distribution Controlling creativity
Diverse Beam Search Penalize similar beams Getting varied outputs
Constrained Beam Search Force certain tokens to appear Template-guided generation

Conclusion

Beam Search is one of the most elegant solutions to the sequence decoding problem — a principled middle ground between the recklessness of Greedy Search and the infeasibility of Exhaustive Search. By maintaining a small set of the most promising partial sequences at every step, it finds high-quality outputs efficiently.

Happy coding — and may your beams always converge on the right answer! 🚀

Comments

Post a Comment

Please post your valuable suggestions