_assumption: we are building against [[thoughts/Autoregressive models|autoregressive transformers models]]_ - Let $\mathcal{F} \subset \mathcal{P}(\mathcal{V})$, where $\mathcal{P}$ is the power set operator, be subset of multi-token string that ends with tokens $\text{EOS} \in \mathcal{V}$. - Text generation tasks is to draw samples from $\mathcal{F}$ Notable sampling methods include greedy decoding (generate tokens recursively with highest probability tokens), beam search (but using heuristic to find the mode of distribution) [^smc] [^smc]: \[@lew2023sequentialmontecarlosteering\] recently proposes a sequential [[thoughts/Monte-Carlo|Monte Carlo steering]]. The idea is to classify causal generations as a _posteriori inference_ problem in a class of discrete probabilistic sequence models. See also [[thoughts/Transformers#Feynman-Kac|Feynman-Kac transformers models]] A pseudocode for sampling procedure is as follow:

"\\begin{algorithm}\n\\caption{LLM token sampling}\n\\begin{algorithmic}\n\\Function{sample}{$L$}\n \\State $s \\gets ()$\n \\For{$i \\gets 1, L$}\n \\State $\\alpha \\gets \\text{LM}(s, \\theta)$\n \\State Sample $s \\sim \\text{Categorical}(\\alpha)$\n \\If{$s = \\text{EOS}$}\n \\State \\textbf{break}\n \\EndIf\n \\State $s \\gets \\text{append}(s, s)$\n \\EndFor\n \\State \\Return $s$\n\\EndFunction\n\\end{algorithmic}\n\\end{algorithm}"

Algorithm 3 LLM token sampling

1:function sample( $L$ )

2: $s \gets ()$

3:for $i \gets 1, L$ do

4: $\alpha \gets \text{LM}(s, \theta)$

5:Sample $s \sim \text{Categorical}(\alpha)$

6:if $s = \text{EOS}$ then

7:break

8:end if

9: $s \gets \text{append}(s, s)$

10:end for

11:

12:return $s$

13:end function

Given that we are dealing with finite discrete distribution, we can then compute an un-normalized conditional distribution by applying a boolean mask $m: \mathcal{P}(\mathcal{V}) \to \{0,1\}^N$, which restricts the support of original distribution: $$ \begin{aligned} \alpha &= \text{LM}(\tilde{S_t}, \theta) \\ \tilde{\alpha} &= m(\tilde{S_t}) \odot \alpha \\ \tilde{s_{t+1}} &\approx \text{Categorial}(\tilde{\alpha}) \end{aligned} $$ > \[!math\] 1. augmentation upon sampling algorithm > >

"\\begin{algorithm}\n\\caption{token sampling with masking}\n\\begin{algorithmic}\n\\Function{sample}{$L$}\n \\State $s \\gets ()$\n \\For{$i \\gets 1, L$}\n \\State $\\alpha \\gets \\text{LM}(s, \\theta)$\n \\State Construct the mask m($s$)\n \\State $\\tilde{\\alpha} \\gets m \\odot \\alpha$\n \\State Sample $\\tilde{s} \\sim \\text{Categorical}(\\tilde{\\alpha})$\n \\If{$\\tilde{s} = \\text{EOS}$}\n \\State \\textbf{break}\n \\EndIf\n \\State $s \\gets \\text{append}(s, \\tilde{s})$\n \\EndFor\n \\State \\Return $s$\n\\EndFunction\n\\end{algorithmic}\n\\end{algorithm}"

> Algorithm 4 token sampling with masking

> 1:function sample( $L$ )

> 2: $s \gets ()$

> 3:for $i \gets 1, L$ do

> 4: $\alpha \gets \text{LM}(s, \theta)$

> 5:Construct the mask m( $s$ )

> 6: $\tilde{\alpha} \gets m \odot \alpha$

> 7:Sample $\tilde{s} \sim \text{Categorical}(\tilde{\alpha})$

> 8:if $\tilde{s} = \text{EOS}$ then

> 9:break

> 10:end if

> 11: $s \gets \text{append}(s, \tilde{s})$

> 12:end for

> 13:

> 14:return $s$

> 15:end function

> \[!tip\] finite automaton > > We define a _finite-state machine_, given by $(Q, \Sigma , \delta, q_0, F)$ [^automaton-definition] where character comprising the strings in $\mathcal{V}$ are drawn from $\Sigma$, i.e: $\mathcal{V} \in \mathcal{P}(\Sigma)$ > > > \[!note\]- example > > > > ![[thoughts/images/vllm/fsm-iterative-generations.webp|FSM illustration]] > > > > target regex: `([0-9]*)?\.?[0-9]*{:rs}` > > > > For simplicity, let the vocabulary $\mathcal{V}$ consists of strings $\{A, ., 42, .2, 1\}$ > > > > - generations start: FSM in state 0, so it masks “A”, since it wouldn’t accepted by the FSM. Then we only sample ”.”, “42”, “.2”, “1” in this case > > - if we sample “.2” then we advance the FSM to state 3. In this case. only “42” and “1” are valid completions, so we mask other values before sampling. > > - If we sample “1” instead, then we advance FSM to state 1, in which case ”.”, “.42”, “.2”, and “1” are valid completions [^automaton-definition]: [[thoughts/DFA|finite state machine]] - $Q$ is a finite set of states - $\Sigma$ is a finite alphabet - $\delta: Q \times \Sigma \to Q$ is the transition function - $q_0 \in Q$ is the start state - $F \subseteq Q$ is the set of all accepted states. > \[!tip\] determinism > > Looping through the vocabulary is still the biggest issue. For that, we preprocess the vocabulary using Regex’s FSM and build a index. > Thus a proceeding for producing matches starting at any point in the FSM is required. We define finding sub-sequences of FSM $M$ that accept string $v$ as follow:

"\\begin{algorithm}\n\\caption{Find sub-sequences of the FSM $M$ that accept the string $v$}\n\\begin{algorithmic}\n\\Function{FindSubSequences}{$M, v$}\n \\State $M = (Q, \\Sigma, \\delta, q_0, F)$\n \\State $\\texttt{res} \\gets ()$\n \\For{$r \\in \\delta^{-1}(\\cdot, v_0)$} \\Comment{$\\text{ Loop through states that read } v_0$}\n \\State $p \\gets (r)$\n \\For{$i \\gets 1, |v| - 1$} \\Comment{$\\text{ Walk the FSM}$}\n \\If{$\\delta(r, v_i) = \\emptyset$} \\Comment{$\\text{ The FSM does not read } v_i$}\n \\State $p \\gets ()$\n \\State \\textbf{break} \\Comment{$\\text{ Stop walking and try the next start state}$}\n \\EndIf\n \\State $r \\gets \\delta(r, v_i)$\n \\State $p \\gets \\text{append}(p, r)$\n \\EndFor\n \\State $\\texttt{res} \\gets \\text{append}(\\texttt{res}, p)$\n \\EndFor\n \\State \\Return $\\texttt{res}$\n\\EndFunction\n\\end{algorithmic}\n\\end{algorithm}"

Algorithm 5 Find sub-sequences of the FSM $M$ that accept the string $v$

1:function FindSubSequences( $M, v$ )

2: $M = (Q, \Sigma, \delta, q_0, F)$

3: $\texttt{res} \gets ()$

4:for $r \in \delta^{-1}(\cdot, v_0)$ do ▷ $\text{ Loop through states that read } v_0$

5: $p \gets (r)$

6:for $i \gets 1, |v| - 1$ do ▷ $\text{ Walk the FSM}$

7:if $\delta(r, v_i) = \emptyset$ then ▷ $\text{ The FSM does not read } v_i$

8: $p \gets ()$

9:break ▷ $\text{ Stop walking and try the next start state}$

10:end if

11: $r \gets \delta(r, v_i)$

12: $p \gets \text{append}(p, r)$

13:end for

14: $\texttt{res} \gets \text{append}(\texttt{res}, p)$

15:end for

16:

17:return $\texttt{res}$

18:end function

We can then define construction of $\sigma$

"\\begin{algorithm}\n\\caption{Construct a map from FSM states to subsets of $\\mathcal{V}$}\n\\begin{algorithmic}\n\\Function{MapStatesToVocab}{$M, \\mathcal{V}$}\n \\State $M = (Q, \\Sigma, \\delta, q_0, F)$\n \\State Initialize the map $\\sigma$ with empty sets for each element in $Q$\n \\For{$v \\in \\mathcal{V}$} \\Comment{$\\text{Loop through the vocabulary}$}\n \\State $Z \\gets \\text{find\\_sub\\_sequences}(M, v)$\n \\For{$z \\in Z$} \\Comment{$\\text{Loop through state sequences accepting } v$}\n \\State $\\sigma(z_0) \\gets \\sigma(z_0) \\cup v$\n \\EndFor\n \\EndFor\n \\State \\Return $\\sigma$\n\\EndFunction\n\\end{algorithmic}\n\\end{algorithm}"

Algorithm 6 Construct a map from FSM states to subsets of $\mathcal{V}$

1:function MapStatesToVocab( $M, \mathcal{V}$ )

2: $M = (Q, \Sigma, \delta, q_0, F)$

3:Initialize the map $\sigma$ with empty sets for each element in $Q$

4:for $v \in \mathcal{V}$ do ▷ $\text{Loop through the vocabulary}$

5: $Z \gets \text{find\_sub\_sequences}(M, v)$

6:for $z \in Z$ do ▷ $\text{Loop through state sequences accepting } v$

7: $\sigma(z_0) \gets \sigma(z_0) \cup v$

8:end for

9:end for

10:

11:return $\sigma$

12:end function