Challenge: Inference-Train Mismatch🔗

Agents communicate through structured messages. The assistant generates tool calls, and the tools send responses back. This stream can be represented in three formats: messages, bytes, and tokens.

In a sensible world, these transformations would be bijective and share some notion of concatenation. We do not, of course, live in a sensible world. Decoding and re-encoding tokens need not recover the original token sequence.

Even worse, making changes to the byte string can cause changes to the tokenization that can propagate in both directions.

Messages behave more slightly more sensibly in practice, but are much less formalized. For instance, there used to be a standard practice that every time a user message was added, the rendering for all previous assistant messages would change, e.g. to remove thinking. Message renderers therefore allow arbitrary global Python code.

These issues do not come up in pretraining or agent SFT where you work with a fixed tokenized representation. RL is another story. Because you are working with real tools, you need to alternate between running inference, interacting with the harness, and training on the result. A token mismatch can invalidate the training signal.

This blog explores how we might rule out token mismatches during RL while keeping the process efficient. This is a prototype of what a formally verified component of a large-scale ML system might look like.

Approach: Certified Rust🔗

Rust is a natural choice for writing tokenizers, as in Hugging Face's Tokenizer library. It offers efficiency, memory safety, and, for our purposes, a path to formal verification. We use Aeneas to translate our implementation into Lean so that we can prove properties of its execution.

To make the translation process concrete, consider a simple function used for BPE encoding. This function converts each input byte into its initial token ID before the merge steps begin.

pub fn initial_tokens(bytes: &[u8]) -> Vec<TokenId> {
    let mut tokens = Vec::new();
    for &byte in bytes {
        tokens.push(byte as usize);
    }
    tokens
}

Aeneas translates the loop into the Lean implementation below. The code is a bit more complex because it makes the Rust execution model explicit.

We use TokenId as a readable alias for the extracted integer type:

We trust the extraction from Rust to Lean, then prove properties of the extracted code. The first theorem below shows that initialization succeeds and preserves every input byte's numeric value.

A Spec for BPE🔗

Once we have translated the code, the core challenge is to write a specification for the tokenizer. Tokenizers are surprisingly complex, and certificates for details such as pre-tokenization and special tokens are not especially illuminating. Instead, let's walk through the textbook BPE part. Assessing this specification still requires an understanding of BPE, but it should be easier than reviewing the Rust implementation.

A token can be viewed as a tree whose leaves are bytes and whose internal nodes record merges. For example, the token below represents abc:

Each merge rule names the two adjacent token IDs that it can combine. The extracted Rust type is:

The rule's position in the table gives its rank; its replacement token ID is 256 + rank. The specification starts with token IDs matching the input bytes and requires the result to satisfy the ranked merge relation:

Here initial is the token sequence before any merges: one token ID per input byte. The equality says its numeric values are exactly the input byte values, in the same order. The relation RankedBpe requires each step to choose the applicable pair with the lowest rule rank, using the leftmost position to break ties.

The definitions below are the actual relations used by the proof. PairAt checks adjacent token IDs; Selected matches them to a rule; and Optimal chooses the smallest rule rank and then the leftmost position.

Our specification for RankedBpe permits either a stop when nothing matches or an optimal merge followed by another ranked execution.

The theorem below uses a configuration whose merge rules are valid. It states that the extracted Rust encoder returns the given token sequence successfully if and only if that sequence satisfies BpeSpec.

This is the connection between the extracted Rust code and our BPE spec.

Reversible Message Rendering🔗

Once tokens are generated, they need to be transformed into messages for the harness. This is not technically tokenization, but it tends to live in tokenization libraries, and the interaction between the two matters.

Message rendering often uses a global Jinja template that takes the full message history and produces a single byte string. In principle, adding a new message could change the rendering of any earlier message.

This behavior conflicts with caching in agents, so systems use runtime checks to guard against changes to earlier output. We will build that preservation property directly into our certificate.

We write the chat template directly in Rust so that we can verify its behavior. Here is a fixed template using Qwen-style delimiters.

pub fn header(role: Role) -> Vec<u8> {
    let mut out = Vec::new();
    out = append(out, b"<|im_start|>");
    out = match role {
        Role::System => append(out, b"system"),
        Role::User => append(out, b"user"),
        Role::Assistant => append(out, b"assistant"),
        Role::Tool => append(out, b"tool"),
    };
    append(out, b"\n")
}

pub fn render(messages: &[Message]) -> Vec<u8> {
    let mut out = Text::new();
    let mut i = 0;
    while i < messages.len() {
        out = out.append(&header(messages[i].role));
        out = out.append_escaped(&messages[i].content);
        out = out.append(b"<|im_end|>\n");
        i += 1;
    }
    out.into_bytes()
}

Using the same techniques as above, we prove round-trip consistency. The theorem states that the extracted Rust functions succeed and that the recovered roles and bodies equal the originals, provided the rendered byte length fits in Usize.

We also prove that appending one message only appends its bytes. If history contains the first n−1 messages and message is the nth, rendering all n messages gives exactly the old bytes followed by the new message's bytes:

Application: A Delta Tokenizer🔗

Now let's return to the main application: keeping inference and training tokens consistent during RL. In open-source RL systems, this is handled by a design philosophy known as Tokens In, Tokens Out, as described in No Tokens Left Behind. The key idea is never to retokenize anything generated by the LLM, so that generation and tool responses extend a single append-only token stream.

The design maintains three representations of the history: messages, bytes, and tokens. The LLM generates tokens, which can be converted to bytes and messages. A tool produces a new message, which can be converted to bytes and tokens.

In principle, adding a new message could change the rendered history, so the system needs to ensure that this does not happen. A later code change might also regenerate earlier tokens from messages, so the system needs to protect against that too. This is typically handled through tests or runtime checks.

Lean certificates let us prove that these failures cannot occur within the verified operations and their stated assumptions. For example, we can prove that rendering appended messages is equivalent to appending their local rendering. This property may be stronger than some applications need; one could also imagine a version with state.

Similarly, we can show that a new byte sequence can be encoded into tokens that decode to exactly those bytes.

These properties may seem like common sense, but ML systems are complex. Formalizing them can reduce reliance on runtime checks for the verified operations, helping prevent bugs without repeatedly checking the same properties during execution.

Extension: Inserting a Past Message🔗

The core RL loop can be implemented with an append operation, but related systems can benefit from other kinds of edits. One example is rescoring tokens under alternative contexts, as in self-distillation. For these systems, we would like to insert messages at earlier steps without changing the token IDs that follow them.

The naive approach would be to insert a message and retokenize the whole output. As we have seen, this risks changing the token sequence of the original generation. If we are not careful, inserting a message could also change the rendering of the surrounding messages.

Given our fixed renderer and its locality properties, we can show that inserting a new message, its bytes, and its encoded tokens preserves the surrounding token IDs and recovers the edited conversation when decoded, under the theorem's validity and size assumptions.

Conclusion🔗

We have presented a prototype BPE tokenizer for agent systems, implemented in Rust with Lean-checked properties of the extracted implementation. It demonstrates how Lean can reduce risk in ML systems by proving critical properties under explicit assumptions.

There is still a lot to do. The specifications remain complex and require a substantial amount of Lean code to understand. The implementation is also limited by the subset of Rust that Aeneas supports and does not support threaded tokenization. Finally, we have proved token and message consistency properties relevant to RL; it would be interesting to tackle harder problems, such as cross-tokenizer transfer or guaranteed constrained decoding.

Contents

1. Challenge: Inference-Train Mismatch
2. Approach: Certified Rust
3. A Spec for BPE
4. Reversible Message Rendering
5. Application: A Delta Tokenizer
6. Extension: Inserting a Past Message
7. Conclusion