Token Healing: Why a Trailing Space Breaks Your LLM Output

Take a code-completion prompt that ends with const userNa. The model should obviously continue with me = .... Instead it emits _ = or restarts the identifier or produces a stray newline. Nothing is wrong with the model, the sampler, or the temperature. The problem is that userNa got tokenized as if the string genuinely ended there, and that token sequence is one the model saw approximately never during training.

This is the boundary artifact that token healing fixes. It is one of the highest-leverage 30-line fixes in a completion stack, and almost nobody implements it until they hit it.

TL;DR

• BPE tokenizes greedily, so the last token of your prompt is chosen without knowing what follows. "userNa" splits differently than the "userName" the model actually trained on.
• The model is then conditioned on a token sequence that is off-distribution, and the highest-probability continuation (the merged token) is literally unreachable from that state.
• Token healing = drop the trailing token(s), then constrain the first generated token to those whose byte string starts with the dropped text. Costs no extra forward passes.
• Backtracking must be recursive: if the sampled token is only a prefix of the dropped bytes, keep constraining on the remainder.
• Trailing whitespace is the worst case, because in byte-level BPE a space almost always merges rightward into the next word. This is why the Anthropic API rejects an assistant prefill ending in whitespace.

Why does the last prompt token go off-distribution?

Because tokenization is greedy and left-to-right, and your prompt boundary is not a real boundary — it is an artifact of where you happened to stop typing.

Run this:

import tiktoken
enc = tiktoken.get_encoding("o200k_base")

enc.encode("https://")     # ':' and '/' merge into a single '://' piece
enc.encode("https:")       # ends with a lonely ':' piece
enc.encode("the ")         # trailing ' ' becomes its own token
enc.encode(" the")         # ' the' is ONE token

During training the model saw " the" as a single unit millions of times. It essentially never saw a bare " " token followed by a bare "the" token, because the tokenizer would never produce that split for natural text. When your prompt ends with a trailing space, you have handed the model a state it has no good statistics for, and you have simultaneously made the correct answer — the merged " the" token — impossible to emit, since it would duplicate the space.

The model does the only thing it can: it puts mass on the weird continuations that do follow a standalone space in the training data. Encoding errors, broken formatting, truncated markup. Output quality falls off a cliff for reasons that look mystical from the API layer.

The same thing happens on every partial-word boundary in fill-in-the-middle code completion, on prompts ending in Answer:, on URLs cut at https:, on any template where you carefully appended a separator character.

Note this is strictly a boundary problem. The interior of the prompt is fine, because the tokenizer saw the whole string. Only the final token is chosen under incomplete information.

What is token healing, exactly?

Token healing removes the last token of the prompt, remembers its byte string, and forces generation to begin with a token that has those bytes as a prefix.

You are letting the model re-derive the boundary itself, with the full context available, instead of freezing a guess made by a tokenizer that could not see the future.

Formally: a language model is a distribution over token sequences, but what you actually want is a distribution over strings. Many tokenizations map to the same string, and the naive approach conditions on exactly one arbitrary tokenization of the boundary. Healing marginalizes over the tokenizations of the last piece. Full marginalization over every boundary in the prompt is intractable; the last one is where all the damage is, and it costs nothing.

How do you implement token healing?

Two pieces: a prefix lookup over the vocabulary, and a mask applied for the first few decode steps.

The vocabulary lookup is the part people over-engineer with tries. You do not need one. Sort the vocabulary by raw bytes once, and every prefix query becomes a contiguous range you find with two binary searches.

import bisect
import tiktoken

enc = tiktoken.get_encoding("o200k_base")

# id -> exact bytes. Special tokens are not decodable; skip them.
_pairs = []
for i in range(enc.n_vocab):
    try:
        _pairs.append((enc.decode_single_token_bytes(i), i))
    except KeyError:
        continue
_pairs.sort()
KEYS = [k for k, _ in _pairs]
IDS  = [i for _, i in _pairs]
EXACT = {k: i for k, i in _pairs}


def _upper(prefix: bytes):
    """Smallest byte string greater than every string starting with prefix."""
    p = bytearray(prefix)
    while p and p[-1] == 0xFF:
        p.pop()
    if not p:
        return None                      # prefix is all 0xFF: no upper bound
    p[-1] += 1
    return bytes(p)


def extensions_of(prefix: bytes) -> list[int]:
    """Token ids whose bytes start with prefix (includes prefix itself)."""
    lo = bisect.bisect_left(KEYS, prefix)
    ub = _upper(prefix)
    hi = len(KEYS) if ub is None else bisect.bisect_left(KEYS, ub)
    return IDS[lo:hi]


def proper_prefixes_of(pending: bytes) -> list[int]:
    """Token ids that are a strict prefix of pending -- at most len(pending)."""
    return [EXACT[pending[:k]] for k in range(1, len(pending)) if pending[:k] in EXACT]

extensions_of is the set that finishes the job in one step. proper_prefixes_of is the set that makes multi-step healing possible, and it is the part most implementations get wrong.

Now the decode loop:

class Healer:
    def __init__(self, pending: bytes):
        self.pending = pending

    def allowed(self):
        if not self.pending:
            return None                                    # unconstrained
        return set(extensions_of(self.pending)) | set(proper_prefixes_of(self.pending))

    def advance(self, token_bytes: bytes):
        if not self.pending:
            return
        if token_bytes.startswith(self.pending):
            self.pending = b""                             # boundary resolved
        elif self.pending.startswith(token_bytes):
            self.pending = self.pending[len(token_bytes):] # partial, keep going
        else:
            raise AssertionError("mask leaked a non-matching token")


def heal(prompt: str):
    ids = enc.encode(prompt)
    if not ids:
        return ids, Healer(b"")
    return ids[:-1], Healer(enc.decode_single_token_bytes(ids[-1]))


prompt_ids, healer = heal("const userNa")
# generation loop
for _ in range(max_new_tokens):
    logits = model_step(prompt_ids)
    mask = healer.allowed()
    if mask is not None:
        logits = restrict(logits, mask)        # set everything else to -inf
    tok = sample(logits)
    healer.advance(enc.decode_single_token_bytes(tok))
    prompt_ids.append(tok)

Prefill is now one token shorter. There are no extra forward passes. The only overhead is a masked softmax on a handful of steps.

Why does token healing need multi-token backtracking?

Because the dropped bytes may not be completable by any single token that also extends past them.

Drop "://". Suppose the vocabulary has ":" and "//" but the highest-probability path is to emit ":" first. If your mask only allowed extensions of "://", you would have banned ":" and forced a low-probability path — reintroducing the exact distortion you were trying to remove. So you allow proper prefixes too, then set pending = b"//" and constrain again on the next step.

This is why advance has three branches. Single-step healing implementations look correct on "userNa" and quietly misbehave on punctuation runs, emoji, and any multi-byte UTF-8 character split across pieces.

One related detail: for byte-level BPE, always compare bytes, not decoded str. A dropped token can end in the middle of a UTF-8 sequence, and decode() on it produces a replacement character that will never prefix-match anything. In HuggingFace, tokenizer.decode([i]) is also not always byte-identical to the piece — SentencePiece tokenizers carry a ▁ marker and add a prefix space. Build the table from the raw vocabulary pieces, not from per-id decode.

Does token healing break prefix caching or constrained decoding?

Prefix caching: no. Healing only alters the tail of the token sequence, so every KV block before the last one is byte-identical and still hits. You lose at most one block.

Constrained decoding: yes, this one bites. If you run a JSON grammar or a regex mask on top, the healing mask and the grammar mask must be intersected, and the intersection can be empty. Example: prompt prefilled with json\n{"naand a schema whose only legal next field isname. The grammar allows tokens continuing"na, and healing allows tokens starting with the dropped bytes, and if the two disagree on where the piece boundary falls you get an all-inf` logit vector and a NaN.

Handle it explicitly: if the intersection is empty, drop the healing constraint and keep the grammar. The grammar is a correctness requirement; healing is a quality optimization. Never let the optimization win that argument.

What do you do when the API won't let you heal?

Hosted chat APIs do not expose the sampler, so you cannot heal — you have to design the boundary so it never needs healing.

Three rules that cover almost everything:

1. Never end a prompt with whitespace. Move the space into the expected output. Write Answer: and let the model produce " 42", not Answer: followed by "42". This was in OpenAI's legacy completions guidance for exactly this reason, and it still applies to any raw-completion endpoint.
2. End on a strong token boundary. A newline, a closing brace, or a chat template's turn token are all real boundaries the model saw constantly in training. Mid-word and mid-punctuation-run are not.
3. For assistant prefill (supported on Claude Opus 4.x and Sonnet 4.x), prefill up to a structural token: json plus a newline, or {"`. Do not prefill a partial word. The Anthropic API rejects a prefill ending in trailing whitespace outright — that error is the platform stopping you from stepping on this exact rake.

If you serve your own weights with vLLM, TGI, or llama.cpp, healing is worth wiring in for completion and FIM endpoints. If you only call hosted chat models, boundary hygiene is the whole fix.

When is it not worth doing?

Free-form chat. The prompt ends at a turn boundary the chat template controls, so there is no partial piece to heal, and you would be adding a mask for zero benefit.

The payoff is concentrated in: code completion and FIM, raw-completion endpoints, few-shot templates with tight separators, any pipeline that programmatically concatenates a partial string. If that describes your workload, A/B it on prompts that end mid-token specifically — averaging over a corpus of clean-boundary prompts will dilute a real effect into noise.

So why does a trailing space break your LLM output?

Because BPE merges a space rightward into the following word, so a prompt ending with a trailing space forces the model into a token state it never saw in training and makes the correct merged continuation unreachable. Token healing fixes it by deleting the last prompt token, keeping its bytes, and masking the first decode steps to tokens that start with those bytes — recursively, until the dropped text is fully re-emitted. It costs one sorted vocabulary array, two binary searches per step, and zero extra forward passes. If you cannot reach the sampler, get the same result for free by ending prompts on real token boundaries and never on whitespace.