ternary-composer-charlm

A strict all-ternary character-level language model — every linear weight is {-1, 0, +1} — built on a Composer lane-FFN (4 conditional ternary "lanes" gated by a ternary controller) with ternary attention, trained with the alive-init recipe. Banked here as a checkpoint of ongoing research (TernaryTest project). Not a chat model — a small research char-LM.

Author: Cody Dixon (Claude / Anthropic as engineering collaborator) Status: research checkpoint; the accompanying paper is unfinished and will ship with the finished model.

Results (held-out test, real corpus)

Checkpoint	Context	Bits / char	Next-char acc	−1 / 0 / +1
best_model.pt (K=16, learned pos)	16	1.374	0.711	21% / 58% / 21%
ctx2048/best_model.pt (RoPE)	2048	1.219	0.734	12% / 75% / 13%
ctx16k/best_model.pt (RoPE, batch 1)	16384	1.440	0.700	—
live1/best_model.pt (RoPE, streaming 12-source mix, per-source held-out)	2048	2.052 mean over 12 sources (TinyStories 1.59, arithmetic 1.06)	—	—
matched float baseline (K=16 arch, same budget)	16	4.789 (overfits)	0.682	—

The 16k run used one window per step (same tokens/step as the 2048 run) and lost to it at every step: one document per step collapses gradient diversity. Long context needs packed multi-document windows (see train_live.py --pack) and more distinct data.

Both rows scored with eval_compare.py on the same 64 held-out windows × 2048 chars, all positions (unigram baseline ≈ 0.166 acc). The K=16 model's originally reported 0.723 / 1.305 used last-position-only sampling of the same held-out text. Ternary params: 1,325,056 → ~323 KB @ 2 bits (vs ~2.6 MB fp16) for either model.

Evaluation caveat. The corpus is concatenated (wiki → math reasoning → arithmetic traces) and split 90/10 contiguously, so the held-out 10% is entirely arithmetic / unit-conversion traces (19.5% digits, no User: lines). Absolute bits/char are measured on formulaic text and are flattering relative to prose; comparisons between checkpoints are fair. A document-interleaved split is planned.

At matched training the float model edges ternary on raw accuracy (+0.02) but badly overfits the small corpus (b/c 4.79); the strict-ternary model generalizes far better — its quantization acts as regularization. It generates coherent-ish text and reproduces the corpus's <think>…</think> The answer is… reasoning format and arithmetic/unit-conversion patterns.

What's in this repo

File	Purpose
best_model.pt	{"state_dict", "config", "vocab"} — fp32 latent weights (see note)
charset.json	the 96-char vocabulary (vocab, stoi)
ternary_lm.py	the model definition (ternary Lin, composer LaneFFN, ternary attention; learned-position or RoPE) + load_checkpoint
generate.py	CLI sampler: python generate.py [--ckpt f] [--prompt "..."] [--temp t] [--greedy] [--info] — hard-rounds the latents to {-1,0,+1} and verifies before sampling
train_best.py	the frozen K=16 recipe that produced best_model.pt
train_ctx.py	the same recipe at long context (RoPE, --K 2048), with periodic held-out eval, best-checkpoint saving and immediate upload
ctx2048/	the 2048-context run: best_model.pt (step 30k), progress.json (per-500-step val curve), log.txt
ctx16k/	the 16384-context run (batch 1) — a negative result, kept for the record
live1/	the live streaming-mix char run (train_live.py, 12 sources, per-source held-out)
tokenizer.json (+.stats.json)	32,768-piece byte-level BPE for the v2 token models (make_tokenizer.py)
corpus_stream.py, train_live.py, run_live.sh, make_tokenizer.py	the live streaming-mix training mode (control.json knobs, code reload, KD, int8 emb/KV, packing, QK-norm)
big1/	the big token model (d512×8, 32k BPE, int8 emb/KV, QK-norm): best_model.pt (step 22k, 1.283 mean b/c over 36 held-out sources), latest.pt (full resumable state incl. Adam), logs
flip_lm.py, fixed_corpus.py	"Adam without Adam": low-rank-Adam trit flips — trains the ternary codes with NO fp32 latents and no per-weight optimizer state (train_live.py --flip); memory-light fixed corpus for side arms
launch_big1.sh, rollover_big1.sh, relaunch_big1.sh	launch / snapshot-to-HF / restore-on-a-fresh-kernel for the big run
eval_compare.py	scores any checkpoints on identical held-out windows (python eval_compare.py a.pt b.pt)

Important: best_model.pt stores the fp32 latent weights used by straight-through-estimator training. The deployed ternary weights are recovered deterministically as round(clamp(w, -1, 1)) (see det_tern in generate.py) — that is exactly what the forward pass uses. Pack those {-1,0,+1} values at 2 bits/weight for the ~323 KB deployable form.

Recipe (the part that matters)

• Alive-init: ternary weights are 0 unless |latent| > 0.5; the default 1/√d init leaves 0% alive → the net starts dead and can't train. Init latents at std ≈ 0.30 (10% alive). This single change unlocked strict-ternary attention.
• Composer lane-FFN: 4 strict-ternary lanes + a ternary sigmoid gate (conditional capacity). The biggest jump (0.45 → 0.604). Sweet spot = 4 lanes.
• Training: Adam 3e-3, cosine LR, 40–60k steps, batch 128, context 16, d=128, 4 blocks, 4 heads. QAT: fp latents + STE; deployed weights ternary.
• Data: 4.65 MB char corpus (wiki + reasoning chat records + arithmetic traces), 90/10 contiguous split.

Ablations at this scale: more lanes (8) worse; deeper (6 blocks) unstable; wider (d192/d256) worse; more data helped (0.663 → 0.723) up to the local limit.

Usage

pip install torch
python generate.py --info                                   # arch, metrics, packed size
python generate.py --prompt "User: What is 2+2?" --temp 0.6  # sample
python generate.py --ckpt ctx2048/best_model.pt --prompt "The " --n 500   # a long-context checkpoint

from ternary_lm import load_checkpoint
m, ckpt, stats = load_checkpoint("best_model.pt")   # latents hard-rounded to {-1,0,+1} and verified
print(m.arch(), stats)                               # e.g. {'d':128,'blocks':4,'K':16,'lanes':4,'heads':4,'pos':'learned'}

What it sounds like. Both models reproduce the corpus's surface form — spelling, punctuation, the User: / Assistant: <think>…</think> The answer is… skeleton and unit-conversion arithmetic. The K=16 model cannot hold a question in view while answering it. The 2048-context model produces cleaner multi-line reasoning traces but still neither computes correctly nor answers the question asked: at 1.3M weights this is a format-and-spelling model, and capacity (not context) is what limits it. Greedy decoding loops for both; temperature 0.6 is the most readable.

Honest limitations

• Tiny char-level model on a small corpus; a research artifact, not a product.
• Training used floating point (Adam / LayerNorm / softmax); only the weights are strictly ternary. Fully integer-only training is an open goal.
• Embeddings, LayerNorm and biases remain fp32.

Provenance

Trained on a Kaggle T4 in an autonomous improvement loop; progression 0.308 → 0.45 → 0.604 → 0.624 → 0.663 → 0.703 → 0.723 (dense → longer → lanes → +longer → +cosine → +data → +data/steps), then the 2048-context run (train_ctx.py --K 2048 --batch 8 --steps 30000, 35 min on one T4). License: to be set by the author.