structured generation

Algorithm described in paper/blog and implemented in outlines.

For a sequence of tokens $S_t=(s_1,...,s_t)$, predict the next token $s_{t+1}$

$$\begin{array}{rl}\alpha & =\text{LLM}(S_{t,}\theta )\\ s_{t+1}& =\text{Categorical}(\alpha )\end{array}$$

where we can perform multinomial sampling from the probability distribution of output logits.

The idea of guided generation is to restrict support of the possibilities of the vocabulary. Forbidding sequences that do not respect the structure of the possible vocabulary.

naive approach

• apply a boolean mask $m$ that restricts the support of the original distribution
• renormalise output logits
• perform multinomial sampling

$O(N)$ cost of deciding whether each possible token is in restricted vocabulary at each step.

efficient approach (finite state machine)

Use a finite state machine tuple, $(Q,\Sigma ,\delta ,q_0,F)$, to efficiently continue the state machine without reading from the beginning of the growing sample sequence each time.

• $Q$, finite set of states
• $\Sigma$, finite alphabet
• $\delta :Q\times \Sigma \to Q$, transition function
• $q_0$, start state
• $F$, set of accept states

Need to pre-process the vocabulary using the regular expression’s FSM and build an index (hash map for mask $m$ with $O(1)$) to find subsequences of FSM that accept each string:

• consider starting in every viable FSM state
• determine the valid next tokens if transition function produces a viable state.