Shard-2

A 49.4M-parameter language model that can take back what it just wrote. It has one token its tokenizer does not, <bs> at id 4096, and it emits that token to delete the token before it. The retracted text stays in the context the model reads, and the reader sees only what survived, so the model can look at an attempt it abandoned and try something else.

Shard-2 starts from a 49.4M base trained on 50B tokens of RedPajama v2, and adds the deletion token through three rounds of repair training described below. The base model is published in this repository as Base50M.

Quick start

pip install -r requirements.txt

python infer.py --info
python infer.py "The capital of France is"
python infer.py "A good recipe for tomato soup starts with" --raw --stats
echo "In the year 3000, the last job for humans is" | python infer.py --tokens 120

Everything needed to run the model is in this directory. infer.py imports torch and tokenizers, carries its own copy of the architecture, and reads its shape from config.json, so there is no project to clone and no package to install.

from infer import Writer

writer = Writer()
stream = writer.generate("The capital of France is", tokens=60, seed=7)

print(writer.visible(stream))    # what a reader sees, deletions applied
print(writer.raw(stream))        # the stream as emitted, retractions bracketed
print(writer.deletions(stream))  # how many deletions it used

generate returns the raw token stream with the deletions still in it. Apply them with visible_tokens(stream, 4096) if you want the ids rather than the text. A deletion at the very start of a turn is prevented from reaching into the prompt by the protected argument on both helpers.

What the deletion token does

--raw puts the retracted attempt in brackets and marks the deletion, which is where the behaviour becomes readable. Three generations at temperature 0.9:

The capital of France is Nice. However, Saudi Arabia was a key player in the[ 198]⌫ development of this country.
The capital of France is situated less than one hour away from the capital and is[ known]⌫ also the world's second busiest stop.
The capital of France is the village of Kochi, a village in the same region as the[ N]⌫ capital of China.

At temperature 1.0 with no top-k filter the repairs get longer and take back whole phrases:

Ambassador of the State of L[a][ Russ][ia]⌫⌫⌫os Angeles
the "Territorial[ K][ill]⌫⌫ Crisis of the EU"
a little bite[ together]⌫ of pizza, then tasting the[ hot][ p][iss]⌫⌫⌫ fruit
he[ says]⌫ adds        because[ they]⌫ they're        the[ answer]⌫ backdrop

Do the deletions help

They are scored by the base model, which has no deletion token and no stake in the answer. For each generation the base scores two versions of the same stream: the visible text with the deletions applied, and the same stream with the deleted tokens left in place. The comparison is restricted to the tokens present in both versions, so both sides score identical targets and only the surrounding context differs. Averaging over all of each version instead would favour the visible one for a mechanical reason, because the visible stream is the other one with its lowest-probability tokens removed.

The prompts are six written by hand and eight cut from a corpus region no training phase touched. Earlier rounds of the same recipe score +0.2002 and +0.3048 on the same prompts, so the deletions improved with each round of training.

Deletions fire after tokens the model rated unlikely, so the decoder settings matter. At top-k 40 and temperature 0.8 the sampler rarely produces such a token and the model rarely deletes anything. The defaults in infer.py sit where the behaviour is active.

Benchmarks

bench_one.py from the MythosMini repository, batch 64, against the base model this was trained from.

The deletion ability costs base capability, and this table is the price. Two things are worth knowing about how to read it. Likelihood benchmarks score text that is fixed in advance, so the model never generates during them and the deletion token can never help, while the probability it reserves for that token still counts against it. Scoring the same wiki text with the deletion column masked out gives 2.318 rather than 2.403, which recovers about a third of the gap and leaves the rest as real drift in the weights.

An untrained deletion row costs nothing at all: the base with one mean-initialised row appended scores 2.182, 80.37 and 39.60, identical to the base itself. The cost above comes from the repair training, not from the extra token.

How it was trained

The corpus is rewritten against the model's own mistakes. At every position the model's next-token draw is taken with teacher forcing, and where that draw differs from the corpus token, the draw is spliced into the sequence followed by <bs> and then the real token. The spliced wrong tokens are masked out of the loss, so they are context the model reads and never targets it is trained to produce. What the model learns from this is to notice a mistake and take it back.

Three rounds, each from the previous one:

Round 1 taught a flat deletion habit, because most disagreements are fluent alternatives that nothing in the context marks as wrong. Filtering to the draws the model itself rated unlikely is what turned the deletions into aborted words and wrong nouns, and the judge score rose with each round.

The deletion row is initialised to the mean of the existing embedding rows. The readout is tied to the embedding, so that initialisation makes the new token's logit the mean logit and leaves every prediction unchanged at the start. A random row of the same norm costs between 1.7 and 4.7 nats, because this readout's logits sit around mean -19 with the 99th percentile near +5.8, and anything close to zero outranks most of the vocabulary.

Training the same way from scratch is supported in the repository through python3 train.py <billions> --backspace, which splices repairs from the live weights every step for any model size.

Files

The architecture in infer.py reproduces the training implementation exactly. Checked against it on the same input: maximum absolute logit difference 0.0, and one-token-at-a-time decoding through the KV cache matches a full forward pass to 4e-5.

Limits

At 49.4M parameters the content is unreliable and often invented. The samples above are typical, including "The capital of France is Nice". What the deletion token improves is the local shape of the text, and a wrong fact stays wrong after a repair.

The deletions are also not an error detector in any strong sense. Measured on held-out text, the probability of a deletion after a token the model got wrong is about 1.11 times the probability after the corpus's own token, so the policy is closer to "delete where the last token looked unlikely" than to "delete where it was wrong". The correlation with the model's own uncertainty before the token is +0.31.

Two more things to keep in mind. The 4097th token is meaningless to any other model, so a checkpoint merge or a tokenizer swap has to account for it. And the context is 8192 while training used 4096, so anything past 4096 is untested.