LLM Engineering

A two-step choice. Greedy decoding takes A (0.6) and then the best continuation of A (0.4): a sequence probability of 0.24. B starts less likely (0.4) but continues with 0.9, for 0.36. A beam of width 2 keeps both first steps and finds B–x. Greedy commits too early; beam search looks a little further.

• The most probable text is often bland or repetitive; for open-ended generation, maximising probability is the wrong goal (Holtzman et al. 2020).
• Wider beams can make translation worse beyond small widths (the beam search curse).
• Deterministic decoding still varies across hardware and batch sizes because of floating-point non-determinism.

step1 = {'A': 0.6, 'B': 0.4}
step2 = {'A': {'x': 0.4, 'y': 0.35, 'z': 0.25}, 'B': {'x': 0.9, 'y': 0.05, 'z': 0.05}}
greedy_first = max(step1, key=step1.get)
greedy_best = max(step2[greedy_first].values())
greedy = (greedy_first, step1[greedy_first] * greedy_best)
beam = max((((a, b), step1[a] * pb) for a in step1 for b, pb in step2[a].items()), key=lambda c: c[1])   # width 2 keeps both first steps

Assumes import numpy as np. Configurations named after real models use their published shapes; other numbers are made up or simulated, as the comments say.

def beam_search(model, prompt_ids, beam_width=4, max_len=50):
    beams = [(prompt_ids, 0.0)]  # (sequence, log_prob)
    for _ in range(max_len):
        candidates = []
        for seq, score in beams:
            logits = model(seq)[:, -1, :]
            log_probs = F.log_softmax(logits, dim=-1)
            topk = log_probs.topk(beam_width)
            for i in range(beam_width):
                new_seq = torch.cat([seq, topk.indices[:, i:i+1]], dim=-1)
                new_score = score + topk.values[:, i].item()
                candidates.append((new_seq, new_score))
        # Keep top-K beams
        beams = sorted(candidates, key=lambda x: -x[1])[:beam_width]
    return beams[0][0]  # best sequenceCopy