Beam Search: global pruning, EOS and sequence scores
Posted 26 October 2026. Original preparation material, not a claim about any company’s interview requirements.
What you should be able to do
Explain the objective, trace a small example on paper, implement the core without library algorithm shortcuts, and defend boundary behavior. Prerequisites: Python, arrays, basic linear algebra and derivatives.
Intuition
Greedy decoding keeps one partial sequence; beam search keeps a bounded set. Expand all active beams, then prune globally. Keeping k successors per parent without a global top-k accidentally grows the beam. Scores are log probabilities to avoid underflow and turn products into sums.
Formula and derivation
Each extra token adds a nonpositive log probability. Raw scoring therefore often favors shorter sequences. Length normalization changes the objective and its stopping bounds. This implementation uses raw scores, allows completed beams to occupy slots without expansion, and stops when all retained beams finish or max_steps is reached. At the cap, it returns completion flags rather than pretending unfinished sequences are complete.
Paper-and-pencil example
At step one
Implementation contract and common errors
For width k, vocabulary V and T steps, full sorting costs O(T kV log(kV)), excluding model inference. Production systems use top-k kernels and must reorder KV caches by parent beam IDs. EOS consumes one generated token here. Ties use lexical sequence order. Width one matches greedy under this exact scoring and tie policy; larger width does not guarantee better human-perceived quality.
Complete runnable reference
The shared file contains all five algorithms, including multi-head attention. Requires Python 3.11+ and NumPy; validated with Python 3.12.14 and NumPy 2.3.5. Inputs are deterministic; no random seed is needed. No paid API or model download.
Download Python implementation and tests
python -m pip install numpy==2.3.5
python core.py
# PASS: 10 algorithm invariants (synthetic fixtures)Show full source
import numpy as np
def softmax(x):
x = np.asarray(x, dtype=float)
if x.size == 0 or not np.isfinite(x).all():
raise ValueError('nonempty finite logits required')
z = np.exp(x - x.max(axis=-1, keepdims=True))
return z / z.sum(axis=-1, keepdims=True)
def attention(q, k, v, causal=False):
if q.shape[-1] != k.shape[-1] or k.shape[-2] != v.shape[-2]:
raise ValueError('incompatible shapes')
scores = q @ k.swapaxes(-1, -2) / np.sqrt(q.shape[-1])
if causal:
if q.shape[-2] != k.shape[-2]:
raise ValueError('this teaching mask requires equal sequence lengths')
mask = np.tril(np.ones(scores.shape[-2:], dtype=bool))
scores = np.where(mask, scores, -np.inf)
z = np.exp(scores - scores.max(axis=-1, keepdims=True))
weights = z / z.sum(axis=-1, keepdims=True)
return weights @ v
def multi_head(x, wq, wk, wv, wo, heads, causal=True):
b, t, d = x.shape
if heads < 1 or d % heads: raise ValueError('heads must divide model dimension')
def split(w):
return (x @ w).reshape(b, t, heads, d // heads).transpose(0, 2, 1, 3)
y = attention(split(wq), split(wk), split(wv), causal)
return y.transpose(0, 2, 1, 3).reshape(b, t, d) @ wo
def kmeans(x, centers, max_iter=100, tol=1e-8):
x, c = np.asarray(x, float), np.array(centers, float, copy=True)
if x.ndim!=2 or c.ndim!=2 or not len(x) or not len(c) or x.shape[1]!=c.shape[1] or max_iter<1 or not np.isfinite(x).all() or not np.isfinite(c).all():
raise ValueError('invalid data or centers')
for _ in range(max_iter):
labels = ((x[:,None,:]-c[None,:,:])**2).sum(-1).argmin(-1)
new = np.array([x[labels==j].mean(0) if np.any(labels==j) else c[j] for j in range(len(c))])
shift = np.max(np.linalg.norm(new-c,axis=1))
c = new
if shift <= tol: break
labels = ((x[:,None,:]-c[None,:,:])**2).sum(-1).argmin(-1)
return c, labels, ((x-c[labels])**2).sum()
class Adam:
def __init__(self, shape, lr=.001, b1=.9, b2=.999, eps=1e-8):
self.m=np.zeros(shape); self.v=np.zeros(shape); self.t=0
self.lr,self.b1,self.b2,self.eps=lr,b1,b2,eps
def step(self, p, g):
if p.shape!=self.m.shape or g.shape!=p.shape or not np.isfinite(g).all(): raise ValueError('invalid gradient')
self.t+=1
self.m=self.b1*self.m+(1-self.b1)*g
self.v=self.b2*self.v+(1-self.b2)*g*g
mh=self.m/(1-self.b1**self.t); vh=self.v/(1-self.b2**self.t)
return p-self.lr*mh/(np.sqrt(vh)+self.eps)
def beam_search(next_log_probs, width=2, max_steps=4, eos='EOS'):
if width<1 or max_steps<1: raise ValueError('positive limits required')
beams=[((),0.,False)]
for _ in range(max_steps):
candidates=[]
for seq,score,done in beams:
if done: candidates.append((seq,score,True)); continue
for token,lp in next_log_probs(seq).items():
if not np.isfinite(lp) or lp>0: raise ValueError('finite log probabilities required')
candidates.append((seq+(token,),score+lp,token==eos))
if not candidates: raise ValueError('no candidates')
beams=sorted(candidates,key=lambda x:(-x[1],x[0]))[:width]
if all(x[2] for x in beams): break
return beams
if __name__=='__main__':
assert np.allclose(softmax([1000,1000]),[.5,.5])
assert np.allclose(softmax([1,2]),softmax([101,102]))
x=np.arange(24,dtype=float).reshape(1,3,8)/24; eye=np.eye(8)
y=multi_head(x,eye,eye,eye,eye,2); changed=x.copy(); changed[:,2]+=10
assert np.allclose(y[:,:2],multi_head(changed,eye,eye,eye,eye,2)[:,:2])
assert np.allclose(multi_head(x,eye,eye,eye,eye,1),attention(x,x,x,True))
c,l,s=kmeans([[0],[2],[8],[10]],[[0],[10]])
assert np.allclose(c,[[1],[9]]) and s==4
c,l,s=kmeans([[0],[0]],[[0],[8]]); assert np.isfinite(c).all()
opt=Adam((1,)); assert np.allclose(opt.step(np.array([1.]),np.array([2.])),[.999])
assert opt.t==1
def probs(seq):
p={'A':.6,'B':.4} if not seq else ({'EOS':.5,'X':.5} if seq[-1]=='A' else {'EOS':.9,'X':.1})
return {k:np.log(v) for k,v in p.items()}
assert beam_search(probs,max_steps=2)[0][0]==('B','EOS')
assert beam_search(probs,width=1,max_steps=2)[0][0]==('A','EOS')
print('PASS: 10 algorithm invariants (synthetic fixtures)')
Interview follow-up
Why must EOS beams stop expanding?
Hint: Think about probability semantics.
Answer and reasoning
A finished sequence is already a complete hypothesis. Appending further tokens changes its meaning and score. Preserve its score, track completion explicitly, and state whether finished hypotheses share the active beam budget.
Practice without looking
- Spend 3 minutes writing shapes, assumptions and an example.
- Spend 12 minutes writing the core implementation on paper.
- Spend 5 minutes tracing an ordinary input and one failure case.
- Type your handwritten version and run the reference invariants. Explain every mismatch before consulting the solution.
Transfer exercise: change one assumption in the implementation contract above. Write a new test that fails before your modification and passes afterward. A convincing explanation of a limitation is better than silently inventing unsupported behavior.
Sources and next lesson
Primary paper or official documentation
Softmax: stable probabilities and gradients
Attention and Multi-Head Attention: write the shapes first
K-Means: assignment, centroid updates and empty clusters