Beam search decoding in NLP - Model Pipeline Trace

Model Pipeline - Beam search decoding

Beam search decoding is a method used in language models to find the most likely sequence of words by exploring multiple options at each step, keeping only the best few choices to balance quality and speed.

Data Flow - 5 Stages

1Input sequence

1 sequence x 1 token→Start with initial token (e.g., <start>)→1 sequence x 1 token

"<start>" token

↓

2Model prediction

1 sequence x 1 token→Model predicts probabilities for next tokens→1 sequence x vocabulary size probabilities

Probabilities for next words like {"the":0.3, "a":0.2, "cat":0.1, ...}

↓

3Beam expansion

beam width sequences x current length tokens→Expand each sequence by all possible next tokens, score them→beam width * vocabulary size sequences x (current length + 1) tokens

From 3 sequences, each expanded by 5 tokens → 15 sequences

↓

4Beam pruning

beam width * vocabulary size sequences x (current length + 1) tokens→Keep top beam width sequences with highest scores→beam width sequences x (current length + 1) tokens

Keep top 3 sequences out of 15

↓

5Repeat until end token

beam width sequences x tokens→Repeat expansion and pruning until <end> token or max length→beam width sequences x final length tokens

Final 3 sequences like ["<start> the cat <end>", "<start> a dog <end>", ...]

Training Trace - Epoch by Epoch

Loss
2.3 |****
1.8 |***
1.4 |**
1.1 |*
0.9 |

Epoch	Loss ↓	Accuracy ↑	Observation
1	2.3	0.25	High loss and low accuracy as model starts learning
2	1.8	0.40	Loss decreases and accuracy improves
3	1.4	0.55	Model learns better word predictions
4	1.1	0.65	Loss continues to decrease steadily
5	0.9	0.72	Model converges with improved accuracy

Prediction Trace - 5 Layers

Layer 1: Initial token input

Layer 2: Model predicts next token probabilities

Layer 3: Beam expansion with beam width=2

Layer 4: Beam pruning

Layer 5: Repeat until <end> token

Model Quiz - 3 Questions

What does beam width control in beam search decoding?

A Number of sequences kept at each step

B Length of the output sequence

C Size of the vocabulary

D Number of training epochs

Key Insight

Beam search balances exploring multiple possible sequences and focusing on the best ones, improving prediction quality compared to greedy search without excessive computation.