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.
| parameters | 49,409,537 |
| vocabulary | 4097 (4096 BPE plus the deletion token) |
| layers | 20, 8 heads of 64 |
| attention | sliding window 256, every sixth layer global and position-free |
| context | 8192 (trained at 4096) |
| model file | 198 MB |
| deletion rate | 5.4% of emitted tokens at temperature 0.9, 12.5% at temperature 1.0 |
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.
| decoder | deletion rate | generations with deletions | judge gain per token | helped |
|---|---|---|---|---|
| temperature 1.0, no top-k | 12.50% | 14 of 14 | +0.5447 | 14 of 14 |
| temperature 0.9, top-k 100 | 5.36% | 12 of 14 | +0.3854 | 10 of 12 |
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.
| Base50M | Shard-2 | |
|---|---|---|
| wiki byte perplexity | 2.182 | 2.403 |
| BLiMP | 80.37 | 79.19 |
| ARC-Easy | 39.60 / 41.40 | 37.60 / 39.00 |
| ARC-Challenge | 18.80 / 21.00 | 18.40 / 22.80 |
| HellaSwag (acc_norm) | 39.40 | 38.40 |
| PIQA (acc_norm) | 57.40 | 56.20 |
| overall | 72.37 | 70.74 |
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 | corpus | repairs | what changed |
|---|---|---|---|
| 1 | 50M tokens | 4.50% of positions | every model and corpus disagreement is eligible |
| 2 | 50M tokens | 5.45% of positions | only draws the model itself rated below p=0.02 |
| 3 | 8M tokens | 15.10% of positions | repairs may take back a run of up to three tokens |
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
| file | what it is |
|---|---|
Shard-2.pt |
the checkpoint, 198 MB, with its architecture config inside |
infer.py |
standalone inference, the architecture, and the deletion rendering |
config.json |
the model's shape, including backspace_token_id |
tokenizer.json |
the 4096-token byte-level BPE the model was trained on |
requirements.txt |
torch and tokenizers |
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.
- Downloads last month
- 26