constrained decoding

Implemented in mlc-ai/xgrammar

Quote

calculate adaptive token bit-mask per batch

IMPORTANT

operating on string level, not token_id

GrammarMatcher ⇒ FSM in xgrammar

Implemented in (Zheng et al., 2024)

Method 3: Jump-Forward Decoding with compressed FSM

tokenization boundary handling

During decoding, it is preferred to combine multiple characters into a single tokens.

For example, when decoding "Hello" in context of JSON decoding, LLM might output the following token ", He, llo, ",

This may cause some strange behaviour if we combine the last " with , (this regex "[\w\d\s]*" with the last , will lead to endless decoding because this token ", is not valid even if the LM wants to stop.)

Fix:

• implement re-tokenization mechanism during jump-forward phase (append string instead of the tokens, followed with re-tokenization of the entire text) $\to$ add approximately 4% of overhead
• use a comprehensive regex to guide the decoding phase, instead of employing multiple concatenated regex ¹

intuition: Instead of expanding to $n$ state, we can compress certain chunks into one state to reduce the size of said FSM.

A way to adapt character regex to work with tokens in outlines:

import outlines.fsm as fsm
from outlines.fsm.regex import make_deterministic_fsm, create_fsm_index_tokenizer
 
new_fsm, _ = make_deterministic_fsm(fsm)
idx, _ = create_fsm_index_tokenizer(new_fsm, tokenizer)

stateDiagram-v2
    [*] --> InputPrompt: Start

    state "input prompt" as InputPrompt
    state "next-token probability distribution" as GetProb
    state "valid tokens" as ListTokens {
        [*] --> CheckTransitions
        CheckTransitions --> FilterTokens: Get index[0].keys()
        FilterTokens --> [*]
    }
    state "Sample Token" as SampleToken
    state "Update FSM State" as UpdateState

    InputPrompt --> GetProb: "model.generate"
    GetProb --> ListTokens: Get next-token distribution
    ListTokens --> SampleToken: Use filtered token list
    SampleToken --> UpdateState: Selected token X
    UpdateState --> [*]: new_state = index[0]["X"]

idx_with_tokens = {
  state: {tokenizer.tokenizer.decode([key]): value for key, value in transitions.items()}
  for state, transitions in idx.items()
}

example

stateDiagram-v2
    direction LR
    0 --> 2: n
    0 --> 1: t
    1 --> 2: a
    2 --> 4: na
    2 --> 3: a
    3 --> 5: am
    4 --> 6: me
    5 --> 6: me
    2 --> 6: name
    6 --> 7: e
    6 --> 8: c
    7 --> 9: p
    8 --> 9: p
    9 --> 11: Paul
    9 --> 12: Pa
    9 --> 10: Jo
    11 --> 13: aul
    12 --> 14: ul
    10 --> 26: o
    26 --> 27: h
    27 --> 14: n
    13 --> 14: l
    14 --> 16: s
    14 --> 15: s
    15 --> 17: s
    16 --> 17: s
    17 --> 18: a
    17 --> 19: ag
    18 --> 20: ge
    19 --> 20: e
    20 --> 21: 30
    20 --> 22: 20
    21 --> 24: 2
    22 --> 24: 2
    22 --> 23: 3
    24 --> 25: 0
    25 --> [*]

note: each state of FSM represents a forward pass to the LM. In vanilla generation, this is essentially necessary. Thus there is no added overhead of FSM for controlling the generated outputs.

From state 2-6, we observer that there are eight different paths to get the same generations of name. We probably don’t need to do this, given that it will all give us result name

But suffice to say, we can hijack this behaviour to accelerate generations by append either of the following tokens word to currently generated sequence:

• [”name”]
• [”n”, “a”, “m”, “e”]
• [”na”, “m”, “e”]
• [”nam”, “e”]
• [”n”, “am”, “e”]
• [”n”, “ame”]
• [”na”, “me”]
• [”n”, “a”, “me”]

A simplified index can be shown as:

simplified_index = {
    0: {'{"': 2},
    2: {"name": 6},
    6: {'":"': 9},
    9: {'Paul': 14, 'John': 14},
    14: {'","': 17},
    17: {'age': 20},
    20: {'":': 22},
    22: {'20': 24, '30': 24},
    24: {'}': 25},
}

That’s at least a 5x speedup over structured generations, given that out of the 9 tokens, two states are single-state transitions. Therefore we only need to call the model twice!!

difference in sampling distribution

All these paths lead to the same string and the same speedup, however they lead to potentially very different states for the LLM when it reaches state 6. That is, the strings are the same, but each path leads to a different conditional probability distribution in stage 6.

(Willard & Louf, 2023), implemented at dottxt-ai/outlines

assumption: we are building against 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) ²

A pseudocode for sampling procedure is as follow:

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:

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 6 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

finite automaton

We define a finite-state machine, given by $(Q, \Sigma , \delta, q_0, F)$ ³ where character comprising the strings in $\mathcal{V}$ are drawn from $\Sigma$, i.e: $\mathcal{V} \in \mathcal{P}(\Sigma)$

example

target regex: ([0-9]*)?\.?[0-9]*

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

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:

We can then define construction of $\sigma$