CharlieH512/km_tokenizer_ctx14_8x1000

KM tokenizer ctx14 (8 × 1000)

A discrete tokenizer for keyboard and mouse input. It turns 100 Hz key and mouse frames into 8 codes per 80 ms tick (each code in [0, 1000)), and turns codes back into key and mouse frames. It was trained on Minecraft gameplay recordings.

• Weights: model.safetensors (10.76 M parameters, fp32)
• Config: config.json
• Code: km_tokenizer_ctx14.py (PyTorch only, plus huggingface_hub and safetensors for loading)

Architecture

Input frames. Time runs at 100 Hz: one frame every 10 ms. Each frame has

• key_press — 79 binary values: 1 if that key or button is held during the frame (key order below);
• mouse_movement — 2 values (dx, dy): the cursor displacement in pixels accumulated during the frame (positive dy is downward), clamped to ±1280 (dx) and ±720 (dy).

Each frame is turned into 244 features: the 79 held states, 79 "pressed this frame" flags, 79 "released this frame" flags, and 7 mouse features: dx/50 and dy/50 clipped to ±4; log(1+|dx|)/log(1+1280) and log(1+|dy|)/log(1+720); sign(dx) and sign(dy); and log(1+√(dx²+dy²)) / log(1+√(1280²+720²)).

Ticks. 8 consecutive frames form one tick (80 ms). Inputs are padded at the end to a whole number of ticks.

Encoder.

1. A linear layer maps each frame's 244 features to 256 dimensions; a learned position embedding for the frame's slot within its tick (0–7) is added.
2. 2 transformer blocks attend over the 8 frames of each tick (each tick on its own).
3. The tick's 8 frame vectors are concatenated (8 × 256) and mapped by a linear layer to one 256-dimensional tick vector.
4. 4 transformer blocks attend across ticks, bidirectionally, with rotary position encoding on the tick axis. Each block only lets a tick attend to ticks within a fixed distance (4, 4, 3 and 3 ticks), so every tick's latent depends on exactly the 14 ticks before and the 14 after it (1.12 s each way), the same for every tick. Padded ticks are masked out as keys.
5. A layer norm and a linear layer give one 32-dimensional latent vector per tick.

All transformer blocks are pre-norm: layer norm → multi-head self-attention (8 heads, width 256) → residual; layer norm → MLP (256 → 1024 → 256, GELU) → residual.

Quantizer. The 32-dimensional latent is cut into 8 fixed slices of 4 dimensions, one per code group. Each slice is quantized on its own: each of its 4 values is bounded with a shifted, scaled tanh and rounded to one of 8, 5, 5 and 5 levels respectively, giving 8 × 5 × 5 × 5 = 1000 possible values per group. The four level indices are packed into one integer code by mixed radix (the first dimension is the least significant digit). Rounding uses a straight-through gradient; there is no learned codebook. No group ever sees another group while being computed: the slices of one latent vector are quantized independently.

Decoder.

1. The 8 quantized slices of a tick are concatenated back into a 32-dimensional vector (each group enters as its own input channels) and mapped by a linear layer to 256 dimensions.
2. 4 transformer blocks attend across ticks causally (tick n reads ticks ≤ n), with rotary position encoding on the tick axis.
3. A linear layer expands each tick vector to 8 frame vectors (8 × 256); a learned frame-slot position embedding is added.
4. 2 transformer blocks attend over the 8 frames of each tick.
5. A layer norm and a linear layer (in fp32) give, per frame, 79 key logits and 2 × 241 mouse logits.

Mouse bins. Each mouse axis is predicted as a distribution over 241 bins: every integer from −96 to 96 (193 bins), plus 24 bins on each side spaced geometrically out to the clamp (±1280 for dx, ±720 for dy), with integer centers. The decoded movement of a frame is the expected bin center under the softmax, rounded to a whole pixel.

Keys. A key is decoded as held in a frame when the sigmoid of its logit is ≥ 0.5.

Training objective. Keys: per-frame sigmoid focal loss (α = 0.25, γ = 2), normalized by the number of held key-frames. Mouse: per-frame, per-axis cross-entropy against a Gaussian soft target over bin index (σ = 0.5 bins); frames at the clamp are excluded. Plus the mean squared excess of the encoder latent over |x| = 2, which keeps the quantizer's tanh out of saturation. Trained 100,000 steps at 512 windows of 10 s per step with Adam (lr 1e-3, 2,000-step warmup, cosine decay to 1e-4), bf16 mixed precision.

Rare-pattern fine-tune (current weights). The model above reconstructed rare key patterns noticeably worse than common ones, on training data as much as on held-out data. It was then fine-tuned for 40,000 more steps (512 windows per step, Adam lr 3e-4, 1,000-step warmup, cosine decay to 1e-5) with the key loss of each tick multiplied by clamp((0.01 / f) ** 0.5, 1, 10), where f is the training-set frequency of that tick's full 8-frame × 79-key held pattern (counted over 3.4 M training ticks); the weights are rescaled to mean 1 in every batch. The mouse loss is unchanged. Architecture and config are unchanged.

Measured reconstruction, before → after the fine-tune, on the same windows (7,422 training-set windows and 758 held-out windows; ticks at least 15 ticks from a window edge; keys at 0.5). Share of ticks whose whole 8-frame key pattern is reconstructed exactly, by how often that pattern occurs in training:

pattern frequency	share of ticks	train	held-out
≥ 1e-2	75 %	0.999 → 0.999	0.999 → 0.999
1e-3 – 1e-2	12 %	0.980 → 0.988	0.979 → 0.988
1e-4 – 1e-3	7 %	0.954 → 0.982	0.950 → 0.981
1e-5 – 1e-4	3.4 %	0.925 → 0.972	0.920 → 0.970
< 1e-5	2.5 %	0.810 → 0.899	0.855 → 0.926
never seen in training	—	—	0.762 → 0.859

Key onset recall (±20 ms) 0.984 → 0.994 held-out; scroll up/down onset recall ~0.88 → 0.94 (training windows). Mouse, held-out: 10 ms frames with nonzero movement reconstructed to the exact pixel 0.896 → 0.899; mean per-axis error of the summed tick displacement 0.50 → 0.44 px.

Usage

import torch
from km_tokenizer_ctx14 import KmTokenizerCtx14   # km_tokenizer_ctx14.py from this repo

tok = KmTokenizerCtx14.from_pretrained("CharlieH512/km_tokenizer_ctx14_8x1000", device="cuda")

# key_press: [B, T, 79] in {0, 1}; mouse_movement: [B, T, 2] pixels, both at 100 Hz.
key_press = torch.zeros(1, 1000, 79, device="cuda")            # 10 s
mouse_movement = torch.zeros(1, 1000, 2, device="cuda")

with torch.no_grad():
    ids = tok.encode_tokens(key_press, mouse_movement)          # [B, 125, 8] long, values in [0, 1000)
    key_logits, mouse_logits, mouse_px = tok.decode_tokens(ids)
    keys_held = torch.sigmoid(key_logits) >= 0.5                # [B, 125, 8, 79]
    # mouse_px: [B, 125, 8, 2] whole-pixel (dx, dy) per 10 ms frame

• encode_tokens(key_press, mouse_movement, lengths=None) — lengths ([B], in frames) marks valid frames when a batch is padded. Returns ceil(T / 8) ticks.
• decode_tokens(ids) → (key_logits [B,N,8,79], mouse_logits [B,N,8,2,241], mouse_px [B,N,8,2]).
• forward(key_press, mouse_movement, lengths=None) encodes and decodes in one call and also returns token_ids, the pre-quantization latent, and frame/tick masks.

Context. Each tick's codes depend on the 14 ticks before and the 14 after it; the key features also read the frame just before each tick. To encode a long stream so that every tick gets that full context, encode overlapping windows (e.g. 125 ticks) and keep only the ticks at least 15 ticks from each window edge, except at the stream's own start and end. The decoder is causal across ticks: to decode tick n in a stream, pass the codes of the preceding ticks together with tick n (up to 124 previous ticks, as in training) and read the last tick's frames.

Key order (index: name (GLFW key code)): 0 space (32), 1 a (65), 2 b (66), 3 c (67), 4 d (68), 5 e (69), 6 f (70), 7 g (71), 8 h (72), 9 i (73), 10 j (74), 11 k (75), 12 l (76), 13 m (77), 14 n (78), 15 o (79), 16 p (80), 17 q (81), 18 r (82), 19 s (83), 20 t (84), 21 u (85), 22 v (86), 23 w (87), 24 x (88), 25 y (89), 26 z (90), 27 0 (48), 28 1 (49), 29 2 (50), 30 3 (51), 31 4 (52), 32 5 (53), 33 6 (54), 34 7 (55), 35 8 (56), 36 9 (57), 37 Escape (256), 38 Return (257), 39 Tab (258), 40 BackSpace (259), 41 Insert (260), 42 Delete (261), 43 Right (262), 44 Left (263), 45 Down (264), 46 Up (265), 47 Page_Up (266), 48 Page_Down (267), 49 Home (268), 50 End (269), 51 Shift_L (340), 52 Control_L (341), 53 Alt_L (342), 54 Super_L (343), 55 Shift_R (344), 56 Control_R (345), 57 Alt_R (346), 58 Super_R (347), 59 Menu (348), 60 bracketleft (91), 61 bracketright (93), 62 backslash (92), 63 semicolon (59), 64 apostrophe (39), 65 comma (44), 66 period (46), 67 slash (47), 68 minus (45), 69 equal (61), 70 grave (96), 71 F1 (290), 72 F3 (292), 73 F5 (294), 74 Caps_Lock (280), 75 mouse left button, 76 mouse right button, 77 scroll up, 78 scroll down.