FastPrepBeam Search and Nucleus Decoding

Beam Search and Nucleus Decoding

Microsoft logoMicrosoft● HardFULLTIMEONSITE INTERVIEW
Learn

Problem statement

Implement two next-token decoding strategies for one finite language-model simulation: beam search and nucleus (top-p) sampling. Return the best beam sentence and one nucleus-sampled sentence.

The report names these two strategies but omits its model, interface, scoring, random source, and exact output. The following finite rules are practice assumptions that make the exercise deterministic.

Model and stopping

  • tokens is a distinct vocabulary. Its final entry is <eos>; startToken is the initial context and is not emitted.
  • percentages[i][j] gives the integer percentage for token j after token i. Every row sums to 100.
  • Generate at most maxTokens transitions. Generating <eos> ends a path and contributes its probability, but does not appear in the sentence. An unfinished path at the cutoff is valid. Join emitted tokens with one space; an immediate <eos> produces the empty string.

Beam search

  • Start with one empty path of probability 1. On each transition expand every unfinished path through every positive-probability token; carry finished paths unchanged. The path score is the raw product of its transition probabilities.
  • From the combined pool retain at most beamWidth paths by descending exact probability, then ascending sentence text. If both still tie, finished precedes unfinished, then compare the last token text. Finished paths occupy beam slots.
  • Stop when all retained paths are finished or maxTokens transitions have occurred. The first retained path is the beam result. Do not renormalize over the beam or add a length penalty.

Nucleus sampling

  • At each transition sort positive-probability next tokens by descending percentage, then ascending token text; treat <eos> as the empty text for tie-breaking. Keep the shortest prefix whose total percentage is at least topP.
  • Renormalize within that prefix. draws[step] is a deterministic uniform draw of draws[step] / 10000. Walk the kept tokens in the same order and choose the first whose cumulative renormalized probability is strictly greater than the draw. Continue from the selected token unless it is <eos>.

Return a two-element string array: [bestBeamSentence, nucleusSentence]. Use exact integer comparisons for the stated percentage and draw rules. Different from the earlier greedy-and-beam interview exercise, this task judges top-p sampling and does not request a greedy result or sentence-probability output.

Function

decodeBeamAndNucleus(tokens: String[], percentages: int[][], startToken: int, maxTokens: int, beamWidth: int, topP: int, draws: int[]) → String[]

Examples

Example 1

tokens = ["start","a","b","<eos>"]percentages = [[0,60,40,0],[0,0,50,50],[0,0,0,100],[0,0,0,100]]startToken = 0maxTokens = 3beamWidth = 2topP = 60draws = [0,0,0]return = ["b","a"]

Beam search retains the finished b path at probability 0.4 ahead of either a path at 0.3. The nucleus first contains only a; on its second step the zero draw selects the tied <eos>.

Example 2

tokens = ["seed","apple","berry","<eos>"]percentages = [[0,40,35,25],[0,0,0,100],[0,0,0,100],[0,0,0,100]]startToken = 0maxTokens = 2beamWidth = 2topP = 75draws = [9999,0]return = ["apple","berry"]

Beam search returns apple with probability 0.4. The first nucleus retains apple and berry to reach 75%; the last draw chooses berry.

Example 3

tokens = ["seed","go","<eos>"]percentages = [[0,20,80],[0,0,100],[0,0,100]]startToken = 0maxTokens = 2beamWidth = 2topP = 50draws = [0,0]return = ["",""]

Both strategies can select <eos> immediately. The generated sentence is then empty, so the returned strings are empty.

Constraints

  • 2 <= tokens.length <= 8. Ordinary tokens are distinct lowercase words of length 1..8; the final token is exactly <eos>.
  • percentages is square with one row per token; entries are integers from 0 to 100, and every row sums to 100.
  • 0 <= startToken < tokens.length - 1; 1 <= maxTokens <= 6; 1 <= beamWidth <= 12; 1 <= topP <= 100.
  • draws.length == maxTokens and each draw is an integer in 0..9999.
  • The <eos> row is never queried after termination. Scores can be compared exactly with a common fixed denominator of 100^maxTokens, which fits signed 64-bit arithmetic under these bounds.

More Microsoft problems

See Microsoft hiring insights
public String[] decodeBeamAndNucleus(String[] tokens, int[][] percentages, int startToken, int maxTokens, int beamWidth, int topP, int[] draws) {
    // Return [best beam sentence, nucleus-sampled sentence].
}
tokens["start","a","b","<eos>"]
percentages[[0,60,40,0],[0,0,50,50],[0,0,0,100],[0,0,0,100]]
startToken0
maxTokens3
beamWidth2
topP60
draws[0,0,0]
expected["b", "a"]
Checking account…