docs/how-it-works.md

How tiny-autocomplete was built

This model was compiled, not trained. No gradient descent, no optimizer, no pretrained weights, and no other model's outputs were used.

• build_weights.py turns text into a standard LlamaForCausalLM checkpoint using closed-form rules.
• export_web.py writes the same kind of network in sparse form for the browser.
• Rebuilds are byte-identical.

The network computes an interpolated Kneser-Ney 5-gram language model exactly. It predicts the next token from the current token and the 3 before it. Witten-Bell smoothing and other context lengths are supported too; the choice was made by measurement (section 9).

1. Source text and tokenizer

Source	Lines	Weight
phrases/*.txt: everyday English, written for this project by an AI coding assistant (Claude)	1,139	1
data/prompts.txt: what people type to an assistant (OASST user turns, MBPP, CoNaLa)	7,133	1
data/replies.txt: human-written assistant replies from OASST coding conversations	12,834	1

scripts/fetch_coding_prompts.py downloads the data, removes code blocks, code-looking lines and markdown, and keeps short prose sentences. Licenses and attribution are in data/SOURCES.md.

The tokenizer is word-level with GPT-2 style byte-level pre-tokenization:

• NFKC normalization, then lowercase.
• Split on the GPT-2 regex. Words keep their leading space as Ġ, and punctuation and contractions ('s, 't, ...) become their own tokens.
• Vocabulary:
  • every hand-written word, then the most frequent data words seen ≥ 3 times;
  • capped at 3,000 for the dense model, where each word costs one hidden dimension;
  • capped at 8,000 for the browser export, where 6,612 words qualify;
  • two specials: <|endoftext|> (id 0, used as BOS, EOS and pad) and <unk> (id 1);
  • rarer words become <unk>, which can appear inside contexts but is never predicted.
• A post-processor prepends BOS, so every prompt starts at position 0 with <|endoftext|>.

2. The reference n-gram

Each line becomes [BOS × 4, w1, ..., wn, EOS]. For every position, the compiler counts which token follows each context of length k = 0..4 that ends there, weighted by source.

• Order 0: add-one-smoothed unigram, with <unk> pinned at logit -30.

• Order k ≥ 1: interpolate with the next-shorter context ctx[1:]. Both smoothing methods have the form p_k(j) = A_j + γ · p_{k-1}(j):

  Witten-Bell:  A_j = c_j / (c + u)                  γ = u / (c + u)
  Kneser-Ney:   A_j = (c_j - d_j) / c,  d_j = min(0.75, c_j)    γ = Σ d_j / c
  

  Here c_j is the count of token j after the context, c the total, and u the number of distinct next tokens. Kneser-Ney uses continuation counts (how many distinct tokens precede ctx + j) for every order below the top one, as KenLM does. An unseen or pruned context backs off: p_k = p_{k-1}.

The model must reproduce p_K(· | longest kept context).

3. The residual stream as named slots (dense model)

Dims	Slot	Written by
0	constant 1 (Llama has no biases; this plays that role)	embedding
1–384	code of the current token	embedding
385–768	code of the token 1 back	attention head 1
769–1152	code of the token 2 back	attention head 2
1153–1536	code of the token 3 back	attention head 3
1537–4536	logits, one dim per vocab entry, scaled by 1e-5	MLP
4537–4559	unused padding (transformers wants hidden % heads == 0)	—

Token codes are seeded Gaussian vectors in 384 dimensions, normalized to unit length. With 3,000 of them, the largest |cosine| between two different codes is 0.256. That's far enough below 1 that a "match" (dot = 1) and a "non-match" (dot ≤ 0.26) can't be confused.

4. Attention: previous-token heads from RoPE alone

Llama has no absolute positions. The only position signal is RoPE, which rotates query and key pairs by position * inv_freq[i].

• The key for every token is the same constant (read from the constant dim): 1 on the first half of each RoPE pair.

• The query for head a is also constant, with each pair's phase set to -a * inv_freq[i]. Then

  score(t, s) = A * Σ_i cos((t - s - a) * inv_freq[i]) / sqrt(head_dim)
  

  which peaks exactly when s = t - a.

• When t < a there is no position a back, and the closest position is 0, which is BOS. The model then "sees" BOS in that slot, matching the BOS padding used when counting.

• rope_theta is picked to maximize the worst-case gap between the target position and every other position, for all t < 128 and every head. It came out to 2. A scales that gap to 30 nats.

• Value = the current-token code. Output projection = identity into slot a.

All heads share one key/value head (num_key_value_heads = 1).

5. MLP: one neuron per kept context

Llama's MLP is down(silu(gate(x)) * up(x)). The compiler creates one neuron for each kept context (lengths 1–4) and one always-on unigram neuron. For a context (c_{k-1}, ..., c_0) (c_0 = current token):

• gate = 100 * (code(c_0)·cur + code(c_1)·prev1 + ... - (k - 0.5) * const).
  • If all k tokens match, this is +50.
  • If any one fails, it's at most 100 * (0.26 - 0.5) = -24, and silu(-24) ≈ -1e-9, effectively off.
• up = 1 (the constant dim).
• down writes, for each token j seen after the context, log p_k(j) − log(γ · p_{k-1}(j)). Every unseen token would get the same log γ, which softmax ignores, so the column is sparse.

Why this is exact. If a context of length k is kept, so are its shorter suffixes, and their neurons fire too. The unigram neuron writes log p_0, and the neurons for lengths 1..K add their deltas. The sum telescopes to log p_K(· | longest kept context) plus a constant.

Pruning. A context is kept if it was seen at least min_count times (weighted) or comes from the hand-written phrases. A suffix is always seen at least as often as the contexts ending in it, so pruning never breaks a backoff chain. The reference n-gram uses the same kept set, so the model still matches it exactly.

• Dense model: min_count is the smallest value that fits 400M parameters. That's 8, keeping 24,773 of 434,284 contexts. Each neuron costs 3 × hidden ≈ 13.7k parameters.
• Browser export: min_count = 2, keeping 107,864 of 520,490 contexts, because a sparse neuron costs only a few bytes.

6. Norms and output head (dense model)

• The residual has a known fixed norm at each point: |x|² = 2 before attention and 5 after. Each RMSNorm weight is sqrt(norm² / hidden), so the norms leave the stream unchanged.
• The MLP writes logits scaled by 1e-5, so they barely change the final norm (< 0.001%). At 1e-4, the 3,000-entry logit slot shifted the norm enough to cause 1e-2-nat errors on extreme logits.
• lm_head is 1e5 times the identity on the logit slot.

7. Browser export (web/model.bin, web/autocomplete.ts)

93% of the dense weights are zero. The attention and embedding weights are fixed by construction, and each MLP neuron is fully described by:

• which tokens its gate matches (its context);
• which logits its down-projection changes, and by how much.

export_web.py stores exactly that, plus the unigram neuron and the vocabulary:

• contents: per context its token ids, its count (support) and the logit changes it writes (as float32);
• size: 4.5 MB for 107,864 contexts.

web/autocomplete.ts looks up the neurons whose context matches the last 1–4 tokens (hash maps per context length), adds their deltas to the unigram logits, and applies softmax.

• Speed: 0.5 ms median per keystroke.
• Exactness: a test exports the dense model's own tables and checks that the runtime's top-5 next tokens and log-probs equal PyTorch's.

8. Confidence policy (suggest.py, mirrored by web/autocomplete.ts)

Nothing here is learned; the policy reads the compiled probabilities.

1. Candidates. The last pre-token is the word being typed (just Ġ after a space). Only vocabulary tokens starting with it are candidates (fun → function).
2. Probability. Candidate probabilities are renormalized over the candidates, then scaled by 1 − P(OOV | prefix): the weighted share of source words with this prefix that are outside the vocabulary. Without that, a prefix whose only in-vocabulary completion is function would get 100% even when the user is typing funky.
3. Show the top candidate only if its probability ≥ τ (and the matched context's support ≥ min_support; the sweep found 0 best).
4. Extend greedily word by word while the product of probabilities stays ≥ τ, up to 8 words, stopping at end-of-text or after . ? ! :.

τ = 0.72 by default. It's a runtime argument (ac.suggest(text, {...ac.policy, tau})).

9. Choosing the settings (scripts/sweep.py)

The compile rules are fixed. A handful of compile settings and the threshold were picked by benchmarking:

1. Hold out 10% of data/prompts.txt (711 prompts) and compile without them.
2. Type each held-out prompt one character at a time. After each keystroke, record the suggestion chain.
3. Score every threshold: precision (shown suggestions that exactly match the next characters), coverage, and keystrokes saved (accepting each correct suggestion with one Tab).
4. Keep the setting with the most keystrokes saved at precision ≥ 90%, and an export ≤ 50 MB.

Highlights from docs/sweep.json (keystrokes saved at ≥ 90% precision, 0.01 threshold grid):

Setting	Keystrokes saved
Witten-Bell, 4-token context (phrases ×5)	21.7%
Kneser-Ney, 4-token context (phrases ×5)	22.5%
Kneser-Ney, 5 / 6 / 8 / 10-token context (phrases ×5)	22.2% / 22.0% / 21.6% / 21.6%
no assistant replies (prompts + phrases only)	20.1–20.9%
assistant replies ×0.3	21.8–22.4%
final: KN, 4 tokens, all sources ×1, vocab 8,000, min_count 2	23.0%
final settings with 5 / 6 / 8 / 10-token context	22.8% / 22.7% / 22.4% / 22.4%

• Longer contexts: they lose slightly. With ~21k lines, 5+-token contexts rarely repeat, and when they don't match the model is back to the same 4-token prediction.
• Topic head (scripts/experiment_topic.py): a content-addressed head that would carry the most recent language or tool name (python, sql, react, …) from anywhere earlier in the prompt.
  • Tested as a closed-form log-ratio bias λ · log p(w | topic) / p(w).
  • At best it left keystrokes saved at 23.0% (22.97% at λ = 0.25), and λ = 1 was worse.
  • It wasn't compiled.

10. Verification (test_model.py)

• Hugging Face loading with AutoModelForCausalLM, AutoTokenizer and pipeline on transformers 4.52 and 5.22.
• Exactness: dense log-probs match the reference n-gram within 1e-3 nats (observed ~2e-5) on 200 phrase prefixes and 50 random sequences of up to 80 tokens.
• Completions: a set of greedy completions; unknown words back off cleanly.
• Dense TypeScript runtime (model.ts): same top-3 tokens and log-probs as PyTorch.
• Browser runtime:
  • the same suggestion chains (text, probability, support) as suggest.py on 300 random prefixes;
  • the same distribution as PyTorch when given the dense model's tables;
  • model.bin ≤ 50 MB.

11. Changing it

Edit phrases/*.txt or data/*.txt (one sentence per line), or the settings in compile_settings.json. Then run python build_weights.py && python export_web.py && python test_model.py. Run python scripts/sweep.py to re-check the settings.

Setting	Meaning
n_prev	previous tokens attended to (context = n_prev + 1)
smoothing, discount	kn or wb; Kneser-Ney discount
phrase_weight, prompt_weight, reply_weight	how much each source counts
max_vocab, min_word_count	vocabulary cap and minimum word frequency
min_count / param_target	pruning threshold, or the dense budget to fit it to
policy.tau, policy.min_support, policy.max_words	confidence policy shipped in model.bin

History

1. 120 parameters (hidden size 2): one fixed next word for each of 22 words, with no context.
2. 91M parameters: the compiled 5-gram over 1,139 hand-written phrases.
3. 380M parameters: adds ~20k human-written coding sentences, with a vocabulary cap and count pruning to fit the budget.
4. Current: Kneser-Ney smoothing; a confidence-gated ghost-text policy; settings chosen by held-out precision and keystrokes saved; a 4.5 MB browser export with a zero-dependency runtime.