CruxTok (T1) and NCM baseline tokenizers trained on XSum

README.md

@@ -0,0 +1,63 @@
+---
+license: mit
+language:
+- en
+tags:
+- tokenizer
+- vector-quantization
+- residual-vq
+datasets:
+- EdinburghNLP/xsum
+---
+
+# CruxTok tokenizers (XSum)
+
+Residual VQ text tokenizers from the paper *What Does the First Level of a Residual Text Tokenizer Know? Learned-Query Pooling Puts the Gist in Level 1* (arXiv link to follow). Code: [github.com/Zilch123/CruxTok](https://github.com/Zilch123/CruxTok).
+
+CruxTok (T1) replaces the fixed-window pooling of [Next-Concept Modeling](https://github.com/HazyResearch/ncm) (NCM) with learned-query pooling. Against an NCM baseline with the same parameter count, its first level (4 codes) recovers much more of a sentence.
+
+## Files
+| Folder | Model | Params | Seed, precision | Text accuracy, k=1 / k=3 / full |
+|---|---|---|---|---|
+| `t1_seed42` | CruxTok (T1) | 73.2M | 42, bf16 | 23.5 / 60.1 / 98.8 |
+| `t1_seed2` | CruxTok (T1) | 73.2M | 2, bf16 | 24.4 / 60.7 / 98.8 |
+| `t0cap_seed42` | NCM, same budget (T0-cap) | 60.6M | 42, fp16 | 23.5 / 47.1 / 95.8 |
+| `t0cap_seed2` | NCM, same budget (T0-cap) | 60.6M | 2, bf16 | 23.0 / 46.5 / 95.7 |
+| `t0cap_matched` | NCM with 5 encoder and 3 decoder blocks | 73.1M | 42, bf16 | 19.3 / 45.6 / 97.8 |
+
+Each folder holds `tokenizer.pt` (weights and config), `config.yaml` and `eval.jsonl` (held-out reconstruction every 5,000 updates). Parameter counts include codebooks. "k" means levels 1..k together: k=1 is 4 codes, k=3 is 20, and full is all 9 levels over a 48-token window (GPT-2 tokenizer).
+
+All models were trained on XSum training summaries for 30,000 updates at batch 256, about 16x below NCM's tokenizer budget.
+
+## Usage
+```bash
+git clone --recurse-submodules https://github.com/Zilch123/CruxTok && cd CruxTok && pip install -e .
+```
+```python
+from cruxtok.model.cruxtok import CruxTok
+
+tok = CruxTok.from_pretrained("TimothDevA/cruxtok", filename="t1_seed42/tokenizer.pt")
+codes = tok.encode("West Ham striker Andy Carroll is facing up to six weeks out of action.")
+print(tok.decode(codes, upto_level=1))   # from the first 4 codes only
+print(tok.decode(codes))                 # all levels
+```
+The NCM baselines load through the repository's NCM wrapper:
+```python
+from huggingface_hub import snapshot_download
+from cruxtok.eval.semantic import load_run
+
+path = snapshot_download("TimothDevA/cruxtok", allow_patterns="t0cap_matched/*")
+tokenizer, cfg = load_run(f"{path}/t0cap_matched", "cpu")
+```
+
+## Results at level 1
+| | Embedding cosine to original | LLM judge, CruxTok win rate |
+|---|---|---|
+| T0-cap | 0.359 | 0.700 [0.678, 0.722] |
+| T0-cap, matched | 0.342 | 0.740 (seed 42), 0.772 (seed 2) |
+| CruxTok (seeds 42 / 2) | 0.449 / 0.455 | |
+
+On 1,000 XSum test summaries. The judge (GLM-5.3, validated on STS-B at Spearman 0.90) compares level-1 reconstructions in both orders. CruxTok does not make similar sentences get similar codes: on STS-B at k=3, code similarity reaches Spearman 0.28 for the matched baseline against 0.02-0.03 for CruxTok.
+
+## Attribution and license
+MIT. Built on [Next-Concept Modeling](https://github.com/HazyResearch/ncm) (Kumbong et al., arXiv:2610.14162, MIT).

t0cap_matched/config.yaml

@@ -0,0 +1,71 @@
+run_name: t0cap_matched
+seed: 42
+data:
+  dataset: xsum
+  root: data
+model:
+  arch: ncm
+  d: 384
+  in_vocab_size: 50304
+  text_tokenizer: gpt2
+  context_length: 48
+  pad_id: 50257
+  eos_id: 50256
+  encoder_blocks: 5
+  decoder_blocks: 3
+  n_heads: 6
+  dropout: 0.1
+  levels:
+  - 4
+  - 6
+  - 10
+  - 14
+  - 20
+  - 26
+  - 32
+  - 40
+  - 48
+  codebook_size: 4096
+  pool_self_attn: false
+  attn_prior: true
+  prior_sharpness: 4.0
+  qk_norm: true
+  query_init_std: 0.02
+  final_level_identity: true
+  ste: per_level
+  m: 3
+  vq:
+    decay: 0.99
+    epsilon: 1.0e-05
+    commitment: 0.25
+    restart_threshold: 0.0
+  nested_dropout:
+    prob: 0.5
+    min_keep: 1
+  pre_quant_dropout: 0.25
+  cascade_norm: true
+train:
+  steps: 30000
+  micro_batch: 256
+  grad_accum: 1
+  lr: 0.0005
+  lr_floor: 5.0e-05
+  warmup: 1000
+  betas:
+  - 0.9
+  - 0.95
+  weight_decay: 0.01
+  grad_clip: 1.0
+  max_nonfinite_steps: 200
+runtime:
+  precision: bf16
+  compile: false
+  ckpt_every: 1000
+  keep_ckpts: 2
+  log_every: 50
+  eval_every: 5000
+  eval_examples: 2000
+  eval_batch_size: 64
+  wandb: false
+  wandb_project: cruxtok
+  output_root: /workspace/runs/t0match

t0cap_matched/eval.jsonl

@@ -0,0 +1,6 @@
+{"step": 5000, "accuracy": [53.212501525878906, 57.42499923706055, 61.85729217529297, 66.28020477294922, 72.75, 79.13749694824219, 87.21979522705078, 91.26354217529297, 96.57604217529297], "accuracy_text": [17.762380599975586, 24.85587501525879, 32.44826126098633, 40.21064376831055, 51.664817810058594, 62.99334716796875, 77.33000946044922, 84.50296020507812, 93.92646026611328], "ce": [3.503490924835205, 2.8185951709747314, 2.3135547637939453, 1.8987091779708862, 1.4300987720489502, 1.0507631301879883, 0.6727535724639893, 0.502348005771637, 0.27380239963531494], "utilization": [28.6376953125, 50.1708984375, 40.33203125, 21.19140625, 21.8017578125, 17.96875, 15.4296875, 6.2255859375, 31.6650390625], "perplexity": [936.9683089475537, 1344.1670879501537, 925.1250105943946, 135.37318233196393, 155.09484999805014, 116.18998023565207, 89.03961497265546, 40.22143538704933, 153.62241451274926]}
+{"step": 10000, "accuracy": [54.03853988647461, 58.83333206176758, 64.26145935058594, 69.38541412353516, 76.75729370117188, 83.40625, 90.60729217529297, 93.98332977294922, 98.15625], "accuracy_text": [18.75831413269043, 27.213600158691406, 36.69068908691406, 45.726165771484375, 58.77124786376953, 70.5654067993164, 83.33887481689453, 89.32742309570312, 96.7294921875], "ce": [3.422107696533203, 2.6940462589263916, 2.123539924621582, 1.6633046865463257, 1.1756455898284912, 0.8036051392555237, 0.4717015326023102, 0.3293425440788269, 0.1590157300233841], "utilization": [32.7392578125, 58.7158203125, 43.84765625, 25.9765625, 28.61328125, 24.51171875, 24.169921875, 8.984375, 55.5908203125], "perplexity": [1020.0326273258261, 1748.5534958309022, 1030.3590384171284, 255.76402873799273, 221.10817838771123, 208.12533860637004, 182.87976701832994, 73.7914468367675, 251.62545791675865]}
+{"step": 15000, "accuracy": [54.568748474121094, 60.68541717529297, 66.67292022705078, 72.18854522705078, 79.27291870117188, 85.41874694824219, 91.97291564941406, 95.12812805175781, 98.45625305175781], "accuracy_text": [19.508499145507812, 30.35291862487793, 40.90354919433594, 50.66888427734375, 63.23540115356445, 74.13525390625, 85.76126861572266, 91.35809326171875, 97.26164245605469], "ce": [3.443415641784668, 2.616473913192749, 2.024202585220337, 1.5482820272445679, 1.0750923156738281, 0.7210527062416077, 0.41320720314979553, 0.27610570192337036, 0.13386695086956024], "utilization": [27.001953125, 57.03125, 46.5576171875, 29.638671875, 33.3740234375, 27.685546875, 29.3701171875, 11.572265625, 68.9208984375], "perplexity": [843.8422620312075, 1715.9488661503397, 1066.1287329441939, 327.1443329685047, 289.5716238634691, 261.6600560956132, 257.4115963019083, 121.41742356250575, 338.19243715241373]}
+{"step": 20000, "accuracy": [54.37187576293945, 61.248958587646484, 68.01041412353516, 74.19999694824219, 81.48750305175781, 87.63021087646484, 93.6312484741211, 96.11458587646484, 98.62083435058594], "accuracy_text": [19.164819717407227, 31.387657165527344, 43.287139892578125, 54.236881256103516, 67.161865234375, 78.05802154541016, 88.702880859375, 93.10791015625, 97.55358123779297], "ce": [3.4246578216552734, 2.5127182006835938, 1.8760426044464111, 1.3933480978012085, 0.9295699596405029, 0.6015011668205261, 0.3386856019496918, 0.23001256585121155, 0.12257106602191925], "utilization": [23.4130859375, 57.177734375, 49.0234375, 32.7392578125, 38.427734375, 31.8115234375, 34.27734375, 13.5986328125, 76.0009765625], "perplexity": [781.2966820113413, 1757.7449285141902, 1084.9180355494564, 386.28056183935485, 374.1097122966993, 329.6584370895928, 331.3581345873344, 161.4191506528769, 389.46282182534964]}
+{"step": 25000, "accuracy": [54.712501525878906, 62.01874923706055, 69.13957977294922, 75.48958587646484, 82.85832977294922, 88.75208282470703, 94.24270629882812, 96.6656265258789, 98.7593765258789], "accuracy_text": [19.8115291595459, 32.74390411376953, 45.30118179321289, 56.524391174316406, 69.59349822998047, 80.04804229736328, 89.78750610351562, 94.08536529541016, 97.79933166503906], "ce": [3.377192735671997, 2.4383018016815186, 1.7972086668014526, 1.3069918155670166, 0.8591858148574829, 0.5518490672111511, 0.3081727623939514, 0.2099113017320633, 0.11539354175329208], "utilization": [21.875, 56.640625, 51.46484375, 35.8154296875, 42.822265625, 34.619140625, 37.5732421875, 15.0146484375, 80.078125], "perplexity": [737.6826493238182, 1782.457467788626, 1156.5689732132591, 481.827457823869, 499.3417301766279, 433.61576743572493, 442.51335549606955, 225.29719274115868, 461.14357208087625]}
+{"step": 30000, "accuracy": [54.42499923706055, 61.759376525878906, 69.31458282470703, 75.921875, 83.54479217529297, 89.49270629882812, 94.76979064941406, 96.8531265258789, 98.78229522705078], "accuracy_text": [19.327421188354492, 32.313377380371094, 45.60421371459961, 57.29120635986328, 70.8111572265625, 81.36178588867188, 90.72246551513672, 94.41796112060547, 97.8399887084961], "ce": [3.3955657482147217, 2.450507879257202, 1.7756669521331787, 1.2720072269439697, 0.8213236331939697, 0.5197571516036987, 0.2890416979789734, 0.19958478212356567, 0.1130252555012703], "utilization": [21.6552734375, 56.103515625, 53.1982421875, 38.916015625, 46.142578125, 36.71875, 39.9658203125, 15.673828125, 82.763671875], "perplexity": [735.4516851069714, 1760.5046327899727, 1141.1894710955467, 537.4602178484665, 522.2748436525364, 484.24891789658665, 497.8270670713057, 269.0723331349587, 530.6274898378966]}

t0cap_matched/tokenizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec8fb5a39dfac60d8f0a540ec035e290ea09e02c3f9a528f83c06e844914fd40
+size 349363531

t0cap_seed2/config.yaml

@@ -0,0 +1,71 @@
+run_name: t0cap_seed2
+seed: 2
+data:
+  dataset: xsum
+  root: data
+model:
+  arch: ncm
+  d: 384
+  in_vocab_size: 50304
+  text_tokenizer: gpt2
+  context_length: 48
+  pad_id: 50257
+  eos_id: 50256
+  encoder_blocks: 0
+  decoder_blocks: 1
+  n_heads: 6
+  dropout: 0.1
+  levels:
+  - 4
+  - 6
+  - 10
+  - 14
+  - 20
+  - 26
+  - 32
+  - 40
+  - 48
+  codebook_size: 4096
+  pool_self_attn: false
+  attn_prior: true
+  prior_sharpness: 4.0
+  qk_norm: true
+  query_init_std: 0.02
+  final_level_identity: true
+  ste: per_level
+  m: 3
+  vq:
+    decay: 0.99
+    epsilon: 1.0e-05
+    commitment: 0.25
+    restart_threshold: 0.0
+  nested_dropout:
+    prob: 0.5
+    min_keep: 1
+  pre_quant_dropout: 0.25
+  cascade_norm: true
+train:
+  steps: 30000
+  micro_batch: 256
+  grad_accum: 1
+  lr: 0.0005
+  lr_floor: 5.0e-05
+  warmup: 1000
+  betas:
+  - 0.9
+  - 0.95
+  weight_decay: 0.01
+  grad_clip: 1.0
+  max_nonfinite_steps: 200
+runtime:
+  precision: bf16
+  compile: false
+  ckpt_every: 1000
+  keep_ckpts: 4
+  log_every: 50
+  eval_every: 5000
+  eval_examples: 2000
+  eval_batch_size: 64
+  wandb: false
+  wandb_project: cruxtok
+  output_root: /workspace/runs/t0cap_s2

t0cap_seed2/eval.jsonl

@@ -0,0 +1,6 @@
+{"step": 5000, "accuracy": [54.06770706176758, 59.66770935058594, 64.77395629882812, 69.11770629882812, 74.49478912353516, 79.54166412353516, 85.48229217529297, 89.43645477294922, 94.97604370117188], "accuracy_text": [18.588321685791016, 28.482999801635742, 37.529563903808594, 45.2309684753418, 54.7579460144043, 63.71027374267578, 74.24797058105469, 81.26200866699219, 91.08832550048828], "ce": [3.607410430908203, 3.0456554889678955, 2.510672092437744, 2.060492515563965, 1.631625771522522, 1.2697405815124512, 0.9334864616394043, 0.6983961462974548, 0.42934003472328186], "utilization": [5.810546875, 11.3525390625, 20.9716796875, 27.9296875, 39.9658203125, 49.8779296875, 57.8125, 47.998046875, 74.072265625], "perplexity": [133.43796704117676, 69.98283569532767, 101.66345140742592, 102.78259641251253, 149.7559511960893, 161.76619899637163, 186.51413776685283, 195.15556997090508, 168.8576372643109]}
+{"step": 10000, "accuracy": [55.657291412353516, 62.05521011352539, 67.70520782470703, 72.57917022705078, 77.81354522705078, 82.63854217529297, 87.90729522705078, 91.73124694824219, 96.63749694824219], "accuracy_text": [21.40983009338379, 32.790096282958984, 42.75313949584961, 51.37657165527344, 60.6485595703125, 69.20362091064453, 78.54952239990234, 85.33259582519531, 94.03547668457031], "ce": [3.373622417449951, 2.716947317123413, 2.1619253158569336, 1.6997838020324707, 1.309045433998108, 1.0027559995651245, 0.7158554792404175, 0.5058043599128723, 0.2927784025669098], "utilization": [13.7451171875, 20.4345703125, 36.71875, 48.388671875, 62.4267578125, 63.7939453125, 65.6005859375, 52.7099609375, 74.90234375], "perplexity": [426.8032146803929, 141.02980789058003, 183.32705854379046, 163.3409118476242, 143.63668112708413, 132.61082977767686, 133.733719956568, 160.64752279545573, 116.18427379468825]}
+{"step": 15000, "accuracy": [55.890625, 62.639583587646484, 68.75, 73.7281265258789, 79.34479522705078, 84.16354370117188, 89.46979522705078, 92.88333129882812, 97.11145782470703], "accuracy_text": [21.849594116210938, 33.865482330322266, 44.6082763671875, 53.412784576416016, 63.36289596557617, 71.91056823730469, 81.32113647460938, 87.3761978149414, 94.8761978149414], "ce": [3.3286664485931396, 2.6200225353240967, 2.040570020675659, 1.5765327215194702, 1.1897622346878052, 0.8818261027336121, 0.613118588924408, 0.4318864047527313, 0.24756917357444763], "utilization": [18.896484375, 26.0009765625, 46.240234375, 57.421875, 68.9697265625, 65.52734375, 65.91796875, 52.099609375, 74.8291015625], "perplexity": [624.37334867757, 172.51539334433417, 234.2229772225131, 178.69873982524487, 182.69490111328017, 133.7779831096999, 142.58550630698443, 140.22894457632668, 99.8436784331231]}
+{"step": 20000, "accuracy": [56.115623474121094, 63.15833282470703, 69.2281265258789, 74.44271087646484, 80.01249694824219, 84.82395935058594, 90.12916564941406, 93.37291717529297, 97.3812484741211], "accuracy_text": [22.29120445251465, 34.80044174194336, 45.463783264160156, 54.682186126708984, 64.54730224609375, 73.08019256591797, 82.49076080322266, 88.24464416503906, 95.35476684570312], "ce": [3.2981956005096436, 2.5627315044403076, 1.9732420444488525, 1.5060522556304932, 1.1193089485168457, 0.8261856436729431, 0.5667018890380859, 0.3978751599788666, 0.22641001641750336], "utilization": [22.94921875, 29.248046875, 50.9765625, 62.4267578125, 72.119140625, 65.2099609375, 65.478515625, 51.46484375, 74.4140625], "perplexity": [732.9527101984335, 191.8198232620225, 249.98393029165425, 182.88813878130912, 182.51579193679342, 129.47609940905355, 138.7944260125042, 133.5627365717088, 98.07097507780051]}
+{"step": 25000, "accuracy": [56.446876525878906, 63.53333282470703, 69.8499984741211, 75.19479370117188, 80.68333435058594, 85.39479064941406, 90.61978912353516, 93.72916412353516, 97.515625], "accuracy_text": [22.89726448059082, 35.46932601928711, 46.57058334350586, 56.02180480957031, 65.74279022216797, 74.09645080566406, 83.36289978027344, 88.87841796875, 95.59312438964844], "ce": [3.2691566944122314, 2.5241336822509766, 1.9279544353485107, 1.451008915901184, 1.0709283351898193, 0.78147292137146, 0.5368722081184387, 0.3779560625553131, 0.21578150987625122], "utilization": [26.6357421875, 31.34765625, 54.1748046875, 64.697265625, 73.3642578125, 66.1376953125, 65.3076171875, 50.9033203125, 74.072265625], "perplexity": [842.2507607249916, 201.95943187711416, 240.35836672756608, 175.48137656517125, 183.56726327472245, 130.4789623361532, 140.12286778623536, 127.10695244103789, 99.03986082718376]}
+{"step": 30000, "accuracy": [56.515625, 63.68333435058594, 69.82499694824219, 75.27916717529297, 80.84896087646484, 85.8218765258789, 90.78229522705078, 93.7906265258789, 97.5999984741211], "accuracy_text": [23.046932220458984, 35.71138381958008, 46.52623748779297, 56.1622314453125, 66.03289031982422, 74.85218048095703, 83.6511459350586, 88.98743438720703, 95.74464416503906], "ce": [3.2579843997955322, 2.5054564476013184, 1.91356360912323, 1.4327224493026733, 1.0509898662567139, 0.759772539138794, 0.5240222215652466, 0.3678230941295624, 0.2089213877916336], "utilization": [29.931640625, 33.0810546875, 55.8349609375, 66.69921875, 73.73046875, 66.0400390625, 65.1123046875, 51.3427734375, 73.9990234375], "perplexity": [920.9595246422899, 210.34261824088722, 245.59177159740554, 178.00706733472964, 184.36948845025094, 131.89028537608547, 137.77632310969395, 129.1095807115925, 99.23608961784494]}

t0cap_seed2/tokenizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df8a8bd682aa86f61d67a1b431210ca04c4fb00a8ba354d38f6119076659774f
+size 299635871

t0cap_seed42/config.yaml

@@ -0,0 +1,69 @@
+run_name: t0cap_seed42
+seed: 42
+data:
+  dataset: xsum
+  root: data
+model:
+  arch: ncm
+  d: 384
+  in_vocab_size: 50304
+  text_tokenizer: gpt2
+  context_length: 48
+  pad_id: 50257
+  eos_id: 50256
+  encoder_blocks: 0
+  decoder_blocks: 1
+  n_heads: 6
+  dropout: 0.1
+  levels:
+  - 4
+  - 6
+  - 10
+  - 14
+  - 20
+  - 26
+  - 32
+  - 40
+  - 48
+  codebook_size: 4096
+  pool_self_attn: false
+  attn_prior: true
+  prior_sharpness: 4.0
+  query_init_std: 0.02
+  final_level_identity: true
+  ste: per_level
+  m: 3
+  vq:
+    decay: 0.99
+    epsilon: 1.0e-05
+    commitment: 0.25
+    restart_threshold: 0.0
+  nested_dropout:
+    prob: 0.5
+    min_keep: 1
+  pre_quant_dropout: 0.25
+  cascade_norm: true
+train:
+  steps: 30000
+  micro_batch: 256
+  grad_accum: 1
+  lr: 0.0005
+  lr_floor: 5.0e-05
+  warmup: 1000
+  betas:
+  - 0.9
+  - 0.95
+  weight_decay: 0.01
+  grad_clip: 1.0
+runtime:
+  precision: fp16
+  compile: false
+  ckpt_every: 1000
+  keep_ckpts: 2
+  log_every: 50
+  eval_every: 5000
+  eval_examples: 2000
+  eval_batch_size: 64
+  wandb: false
+  wandb_project: cruxtok
+  output_root: /kaggle/working

t0cap_seed42/eval.jsonl

@@ -0,0 +1,6 @@
+{"step": 5000, "accuracy": [54.34687423706055, 60.140625, 65.1968765258789, 69.5218734741211, 74.58541870117188, 79.609375, 85.63541412353516, 89.6187515258789, 95.05937194824219], "accuracy_text": [19.17036247253418, 29.369918823242188, 38.29637908935547, 45.95343780517578, 54.918697357177734, 63.83037567138672, 74.51958465576172, 81.58536529541016, 91.23614501953125], "ce": [3.5326426029205322, 2.9980578422546387, 2.4724016189575195, 2.0435757637023926, 1.6318587064743042, 1.268642544746399, 0.9179883599281311, 0.682036280632019, 0.4159570634365082], "utilization": [7.3486328125, 10.8154296875, 21.484375, 28.4912109375, 40.4541015625, 50.68359375, 58.4228515625, 46.5576171875, 72.36328125], "perplexity": [130.9318166882287, 60.93954789871098, 93.1573349044893, 112.80948003747488, 126.39419267768109, 134.70961630531963, 132.3233649094053, 156.3945517802934, 68.08224614434756]}
+{"step": 10000, "accuracy": [55.84166717529297, 62.12187576293945, 67.73542022705078, 72.53958129882812, 77.88021087646484, 82.83958435058594, 88.484375, 91.85832977294922, 96.70625305175781], "accuracy_text": [21.856985092163086, 32.97856521606445, 42.830745697021484, 51.33591842651367, 60.76496505737305, 69.56208801269531, 79.57317352294922, 85.55986785888672, 94.1592788696289], "ce": [3.34218692779541, 2.711482048034668, 2.162088632583618, 1.705936312675476, 1.3034499883651733, 0.9804188013076782, 0.6839263439178467, 0.49466997385025024, 0.2837505638599396], "utilization": [16.6015625, 20.1416015625, 37.646484375, 48.974609375, 62.1826171875, 64.2578125, 65.966796875, 51.3671875, 72.75390625], "perplexity": [511.7939955434662, 112.3994731170562, 166.19694234383218, 170.2669072612785, 165.1594976208472, 142.0599621397199, 130.465150844885, 150.65661952103477, 89.12860766259386]}
+{"step": 15000, "accuracy": [56.24583435058594, 62.86979293823242, 68.93645477294922, 74.08229064941406, 79.45104217529297, 84.37604522705078, 89.73645782470703, 92.95104217529297, 97.2406234741211], "accuracy_text": [22.498151779174805, 34.2960090637207, 44.95380783081055, 54.04841232299805, 63.55691146850586, 72.28750610351562, 81.79785919189453, 87.5, 95.10901641845703], "ce": [3.2963075637817383, 2.609851360321045, 2.035344362258911, 1.5629961490631104, 1.1664491891860962, 0.8629749417304993, 0.5945923328399658, 0.4247514307498932, 0.2376238852739334], "utilization": [21.58203125, 25.048828125, 45.80078125, 57.6904296875, 69.7509765625, 65.3564453125, 66.162109375, 50.4150390625, 72.6806640625], "perplexity": [695.4064788502327, 132.61772243786194, 190.33187963127475, 187.45854002830737, 191.57466604876723, 140.1668394739967, 125.01020408727094, 151.43837777731662, 93.04816797735988]}
+{"step": 20000, "accuracy": [56.485416412353516, 63.35729217529297, 69.62395477294922, 74.82499694824219, 80.39167022705078, 85.484375, 90.5250015258789, 93.6343765258789, 97.45625305175781], "accuracy_text": [22.99334716796875, 35.223575592041016, 46.199188232421875, 55.36585235595703, 65.22173309326172, 74.25350952148438, 83.19475555419922, 88.70842742919922, 95.48780822753906], "ce": [3.2694475650787354, 2.5511386394500732, 1.9699078798294067, 1.4940296411514282, 1.0963995456695557, 0.7958168387413025, 0.5467096567153931, 0.3858109414577484, 0.22150662541389465], "utilization": [26.0009765625, 28.5888671875, 51.611328125, 62.451171875, 72.3876953125, 65.625, 65.7958984375, 49.951171875, 72.4609375], "perplexity": [827.3374099673713, 147.82848875514574, 214.84706568805413, 194.2287852938266, 194.63245208082347, 141.79210217128258, 126.01109911704962, 151.55367111005665, 94.31519425789206]}
+{"step": 25000, "accuracy": [56.657291412353516, 63.77396011352539, 69.94895935058594, 75.38021087646484, 81.00416564941406, 85.97604370117188, 90.78646087646484, 93.86145782470703, 97.62395477294922], "accuracy_text": [23.29637908935547, 35.90909194946289, 46.76644515991211, 56.345157623291016, 66.30451202392578, 75.1238021850586, 83.65669250488281, 89.11123657226562, 95.78529357910156], "ce": [3.246438503265381, 2.5143277645111084, 1.9254882335662842, 1.450392723083496, 1.0542837381362915, 0.7612700462341309, 0.5246148109436035, 0.3701670467853546, 0.20957987010478973], "utilization": [29.4189453125, 30.9326171875, 54.39453125, 65.5029296875, 73.291015625, 65.52734375, 65.087890625, 49.9755859375, 72.3876953125], "perplexity": [926.9544351301709, 153.23561433344636, 231.3666269464052, 199.34353446078825, 191.20870286591926, 137.90416279576615, 126.70067217708873, 149.06736918590545, 96.1642973025234]}
+{"step": 30000, "accuracy": [56.75728988647461, 63.96041488647461, 70.14895629882812, 75.73958587646484, 81.16458129882812, 86.04270935058594, 90.99166870117188, 94.015625, 97.66041564941406], "accuracy_text": [23.458980560302734, 36.27494430541992, 47.11936569213867, 56.977088928222656, 66.59091186523438, 75.24205780029297, 84.02069854736328, 89.38469696044922, 95.84996032714844], "ce": [3.23435378074646, 2.495004415512085, 1.9029996395111084, 1.4246848821640015, 1.0380394458770752, 0.7538253664970398, 0.5143020153045654, 0.3601054549217224, 0.20502860844135284], "utilization": [31.640625, 32.4951171875, 56.4697265625, 67.0654296875, 74.072265625, 65.4541015625, 64.8681640625, 50.1220703125, 72.4365234375], "perplexity": [995.534502474921, 161.337582380827, 237.87095598663637, 206.06728145241632, 191.96256458392781, 135.93205310494784, 127.57912533833478, 146.26251585017687, 97.97118498733744]}

t0cap_seed42/tokenizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9aa5f80641de37b882545c93cef3acddcf516953d37a029a071dd757fca126c
+size 299635807

t1_seed2/config.yaml

@@ -0,0 +1,69 @@
+run_name: t1_seed2
+seed: 2
+data:
+  dataset: xsum
+  root: data
+model:
+  arch: cruxtok
+  d: 384
+  in_vocab_size: 50304
+  text_tokenizer: gpt2
+  context_length: 48
+  pad_id: 50257
+  eos_id: 50256
+  encoder_blocks: 4
+  decoder_blocks: 2
+  n_heads: 6
+  dropout: 0.1
+  levels:
+  - 4
+  - 6
+  - 10
+  - 14
+  - 20
+  - 26
+  - 32
+  - 40
+  - 48
+  codebook_size: 4096
+  pool_self_attn: false
+  attn_prior: true
+  prior_sharpness: 1.0
+  qk_norm: true
+  query_init_std: 0.02
+  final_level_identity: true
+  ste: per_level
+  m: 3
+  vq:
+    decay: 0.99
+    epsilon: 1.0e-05
+    commitment: 0.25
+    restart_threshold: 0.01
+  nested_dropout:
+    prob: 0.5
+    min_keep: 1
+train:
+  steps: 30000
+  micro_batch: 128
+  grad_accum: 2
+  lr: 0.0005
+  lr_floor: 5.0e-05
+  warmup: 1000
+  betas:
+  - 0.9
+  - 0.95
+  weight_decay: 0.01
+  grad_clip: 1.0
+  max_nonfinite_steps: 200
+runtime:
+  precision: bf16
+  compile: false
+  ckpt_every: 1000
+  keep_ckpts: 2
+  log_every: 50
+  eval_every: 5000
+  eval_examples: 2000
+  eval_batch_size: 64
+  wandb: false
+  wandb_project: cruxtok
+  output_root: /workspace/runs/t1s2

t1_seed2/eval.jsonl

@@ -0,0 +1,6 @@
+{"step": 5000, "accuracy": [55.140625, 62.751041412353516, 73.11250305175781, 83.48645782470703, 93.05416870117188, 96.64895629882812, 98.5406265258789, 99.09166717529297, 99.16041564941406], "accuracy_text": [20.436067581176758, 33.932003021240234, 52.317073822021484, 70.71138000488281, 87.68107604980469, 94.05764770507812, 97.41500091552734, 98.39246368408203, 98.5107192993164], "ce": [3.1804537773132324, 2.5852463245391846, 1.7850604057312012, 1.032672643661499, 0.38938450813293457, 0.20286588370800018, 0.11698425561189651, 0.08413901925086975, 0.07818673551082611], "utilization": [56.103515625, 69.6044921875, 84.3994140625, 90.869140625, 95.263671875, 91.943359375, 98.0712890625, 99.7802734375, 65.869140625], "perplexity": [1821.493333938759, 1796.9275016530246, 1679.7554939525205, 1689.3302395659646, 1634.955959946045, 1592.2415810882194, 2093.276468114015, 2444.1801641288594, 1654.9849332453493]}
+{"step": 10000, "accuracy": [56.45833206176758, 64.21353912353516, 74.67292022705078, 84.63021087646484, 93.98229217529297, 97.09791564941406, 98.7406234741211, 99.18645477294922, 99.27812194824219], "accuracy_text": [22.777162551879883, 36.5280876159668, 55.07945251464844, 72.73835754394531, 89.33111572265625, 94.85218048095703, 97.76607513427734, 98.5569076538086, 98.71951293945312], "ce": [3.112482786178589, 2.4984405040740967, 1.6711896657943726, 0.9625691771507263, 0.33915868401527405, 0.17722320556640625, 0.10017817467451096, 0.07250573486089706, 0.06723276525735855], "utilization": [41.69921875, 65.9423828125, 85.791015625, 88.818359375, 95.4345703125, 91.3330078125, 97.7294921875, 99.1455078125, 57.6171875], "perplexity": [1294.9579979358834, 1741.600331064031, 1631.9906855015229, 1502.2585012059355, 1521.141240711649, 1635.6499580976742, 2170.922738635056, 1934.7959151004673, 1313.3795281436785]}
+{"step": 15000, "accuracy": [56.82604217529297, 65.45625305175781, 76.13229370117188, 86.23332977294922, 94.03125, 97.45417022705078, 98.84687805175781, 99.25312805175781, 99.30000305175781], "accuracy_text": [23.455284118652344, 38.74168395996094, 57.675537109375, 75.58203887939453, 89.41426086425781, 95.4841079711914, 97.95454406738281, 98.67516326904297, 98.75831604003906], "ce": [3.053471088409424, 2.3636202812194824, 1.5416793823242188, 0.8347063660621643, 0.3407808542251587, 0.16278861463069916, 0.09493208676576614, 0.07106860727071762, 0.06734021008014679], "utilization": [42.28515625, 66.2841796875, 87.9150390625, 89.4287109375, 93.8720703125, 93.45703125, 97.802734375, 99.658203125, 57.2021484375], "perplexity": [1317.9944456856117, 1724.5568716445262, 1745.8999370891952, 1528.6597415816566, 1458.4354543892434, 1721.8321707922082, 2142.8608895148604, 2302.4647495579416, 1251.507866888947]}
+{"step": 20000, "accuracy": [57.11771011352539, 65.88749694824219, 77.01979064941406, 87.1875, 94.58958435058594, 97.7093734741211, 98.8531265258789, 99.1937484741211, 99.25208282470703], "accuracy_text": [23.939393997192383, 39.493717193603516, 59.24427032470703, 77.27642059326172, 90.4028091430664, 95.93680572509766, 97.96562957763672, 98.5698471069336, 98.6733169555664], "ce": [3.018975019454956, 2.3048031330108643, 1.4664212465286255, 0.7771292328834534, 0.31753838062286377, 0.15883898735046387, 0.09886088967323303, 0.07723783701658249, 0.07355011254549026], "utilization": [42.8466796875, 66.748046875, 88.427734375, 90.5517578125, 93.5302734375, 93.1396484375, 97.4609375, 99.5849609375, 55.7861328125], "perplexity": [1345.4842417352743, 1751.1868715983926, 1790.1745994479193, 1607.462225279443, 1558.899051772292, 1819.1480076238074, 2197.494502273359, 2304.12757143997, 1203.860098226307]}
+{"step": 25000, "accuracy": [57.43229293823242, 66.40833282470703, 77.59791564941406, 87.37916564941406, 94.97396087646484, 97.80833435058594, 98.97708129882812, 99.23332977294922, 99.28333282470703], "accuracy_text": [24.493717193603516, 40.415740966796875, 60.26422882080078, 77.61456298828125, 91.08832550048828, 96.11603546142578, 98.189208984375, 98.64375305175781, 98.72875213623047], "ce": [2.9927451610565186, 2.258744716644287, 1.4235345125198364, 0.7586280703544617, 0.3002151846885681, 0.1534961760044098, 0.09437588602304459, 0.07536539435386658, 0.07284251600503922], "utilization": [41.1865234375, 63.818359375, 88.7451171875, 89.6240234375, 91.845703125, 91.064453125, 96.09375, 99.4384765625, 51.513671875], "perplexity": [1296.0149443227683, 1662.2381519992875, 1707.3146260739395, 1540.0652413934336, 1405.1790850524008, 1523.7198196369857, 1927.3658992422415, 2135.064500536427, 985.9759226837346]}
+{"step": 30000, "accuracy": [57.381248474121094, 66.39375305175781, 77.82083129882812, 87.73542022705078, 95.19479370117188, 97.9749984741211, 98.99583435058594, 99.26041412353516, 99.29582977294922], "accuracy_text": [24.408721923828125, 40.39356994628906, 60.66518783569336, 78.24833679199219, 91.47634887695312, 96.4079818725586, 98.2187728881836, 98.68810272216797, 98.75092315673828], "ce": [2.988381862640381, 2.2457072734832764, 1.404836893081665, 0.7258455157279968, 0.2905394434928894, 0.15000738203525543, 0.09283383935689926, 0.07463284581899643, 0.07256363332271576], "utilization": [40.966796875, 62.4267578125, 88.1591796875, 90.1123046875, 91.6748046875, 90.0146484375, 94.482421875, 99.12109375, 50.09765625], "perplexity": [1287.5675062103066, 1646.9904020881731, 1704.5123158046617, 1562.4458436521609, 1438.3441912377446, 1525.114011111846, 1839.7939263993092, 2044.051705285007, 853.8413203380503]}

t1_seed2/tokenizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3bac3be79a386f6fa022770b95456367d7d9a2a883da70f5b3ec9f91f94b3299
+size 349802315

t1_seed42/config.yaml

@@ -0,0 +1,69 @@
+run_name: t1_seed42
+seed: 42
+data:
+  dataset: xsum
+  root: data
+model:
+  arch: cruxtok
+  d: 384
+  in_vocab_size: 50304
+  text_tokenizer: gpt2
+  context_length: 48
+  pad_id: 50257
+  eos_id: 50256
+  encoder_blocks: 4
+  decoder_blocks: 2
+  n_heads: 6
+  dropout: 0.1
+  levels:
+  - 4
+  - 6
+  - 10
+  - 14
+  - 20
+  - 26
+  - 32
+  - 40
+  - 48
+  codebook_size: 4096
+  pool_self_attn: false
+  attn_prior: true
+  prior_sharpness: 1.0
+  qk_norm: true
+  query_init_std: 0.02
+  final_level_identity: true
+  ste: per_level
+  m: 3
+  vq:
+    decay: 0.99
+    epsilon: 1.0e-05
+    commitment: 0.25
+    restart_threshold: 0.01
+  nested_dropout:
+    prob: 0.5
+    min_keep: 1
+train:
+  steps: 30000
+  micro_batch: 128
+  grad_accum: 2
+  lr: 0.0005
+  lr_floor: 5.0e-05
+  warmup: 1000
+  betas:
+  - 0.9
+  - 0.95
+  weight_decay: 0.01
+  grad_clip: 1.0
+  max_nonfinite_steps: 200
+runtime:
+  precision: bf16
+  compile: false
+  ckpt_every: 1000
+  keep_ckpts: 10
+  log_every: 50
+  eval_every: 5000
+  eval_examples: 2000
+  eval_batch_size: 64
+  wandb: false
+  wandb_project: cruxtok
+  output_root: /workspace/runs/qknorm

t1_seed42/eval.jsonl

@@ -0,0 +1,6 @@
+{"step": 5000, "accuracy": [54.79375076293945, 62.29166793823242, 72.00729370117188, 82.6812515258789, 92.62604522705078, 96.72187805175781, 98.47083282470703, 99.03854370117188, 99.12604522705078], "accuracy_text": [19.826311111450195, 33.12084197998047, 50.34922409057617, 69.28307342529297, 86.91980743408203, 94.1869888305664, 97.28750610351562, 98.2945327758789, 98.44973754882812], "ce": [3.2460713386535645, 2.63155460357666, 1.8734214305877686, 1.0819611549377441, 0.41279634833335876, 0.20239666104316711, 0.12097768485546112, 0.08711159974336624, 0.08110690861940384], "utilization": [47.314453125, 68.8232421875, 84.716796875, 91.4306640625, 95.2392578125, 94.7509765625, 96.0693359375, 99.21875, 65.7958984375], "perplexity": [1523.93489841884, 1833.5748810035898, 1859.381016057292, 1811.11856386743, 1606.1336675590878, 1795.7360298881833, 1951.8247166159993, 2115.511783240834, 1642.6594246845148]}
+{"step": 10000, "accuracy": [55.98125076293945, 64.78020477294922, 74.56562805175781, 84.87604522705078, 93.53020477294922, 97.3343734741211, 98.703125, 99.1624984741211, 99.2406234741211], "accuracy_text": [21.925350189208984, 37.531410217285156, 54.89282989501953, 73.18366241455078, 88.5291976928711, 95.27531433105469, 97.70140075683594, 98.5162582397461, 98.65299224853516], "ce": [3.1202988624572754, 2.4176151752471924, 1.680463433265686, 0.9490679502487183, 0.3663715124130249, 0.17349804937839508, 0.10628468543291092, 0.07730516791343689, 0.07276992499828339], "utilization": [42.7490234375, 63.6474609375, 84.9853515625, 88.720703125, 93.5791015625, 93.4814453125, 96.484375, 98.7060546875, 56.591796875], "perplexity": [1378.6484962820055, 1592.9030164311419, 1641.056830632925, 1466.89763724492, 1386.4593059188473, 1735.4334849232175, 1939.0231868938245, 1895.5098250269916, 1177.0408571095247]}
+{"step": 15000, "accuracy": [56.25728988647461, 66.01145935058594, 76.3187484741211, 86.35520935058594, 94.23542022705078, 97.31354522705078, 98.69062805175781, 99.16458129882812, 99.24791717529297], "accuracy_text": [22.42793846130371, 39.72283935546875, 58.006282806396484, 75.80377197265625, 89.77826690673828, 95.23651123046875, 97.67922973632812, 98.51995849609375, 98.66592407226562], "ce": [3.0898919105529785, 2.309372663497925, 1.5377250909805298, 0.8394669890403748, 0.33260130882263184, 0.174981489777565, 0.10495343059301376, 0.07660169154405594, 0.07079881429672241], "utilization": [42.1142578125, 63.671875, 87.20703125, 88.916015625, 93.1640625, 92.7734375, 97.3876953125, 99.267578125, 56.4453125], "perplexity": [1354.1385987788587, 1675.4302282706426, 1778.044092379772, 1504.512340163417, 1471.4456407581142, 1605.9628889194328, 2046.688939298212, 2186.8563960395845, 1131.7122021635212]}
+{"step": 20000, "accuracy": [56.561458587646484, 66.60624694824219, 77.36145782470703, 86.94166564941406, 94.5374984741211, 97.62916564941406, 98.88021087646484, 99.2281265258789, 99.29166412353516], "accuracy_text": [22.96563148498535, 40.77605438232422, 59.85218048095703, 76.8403549194336, 90.3104248046875, 95.7945327758789, 98.013671875, 98.63082122802734, 98.7435302734375], "ce": [3.058846950531006, 2.244004726409912, 1.4540947675704956, 0.7876423597335815, 0.32247376441955566, 0.16117192804813385, 0.09811964631080627, 0.07403852045536041, 0.07015985995531082], "utilization": [40.13671875, 63.6474609375, 89.501953125, 90.234375, 91.7236328125, 93.65234375, 97.4853515625, 99.51171875, 53.3935546875], "perplexity": [1272.9619831915318, 1718.5528054983597, 1898.5944590349425, 1683.992351956181, 1425.0761553132572, 1844.5355950471774, 2192.9588683071197, 2318.587616647525, 979.8258756106333]}
+{"step": 25000, "accuracy": [56.907291412353516, 66.86666870117188, 77.44271087646484, 87.43228912353516, 95.10832977294922, 97.94166564941406, 98.99791717529297, 99.2718734741211, 99.3125], "accuracy_text": [23.562454223632812, 41.22875213623047, 59.988914489746094, 77.70879364013672, 91.32298278808594, 96.3488540649414, 98.22246551513672, 98.70842742919922, 98.78048706054688], "ce": [3.023197889328003, 2.226142644882202, 1.4279580116271973, 0.7442720532417297, 0.2919045686721802, 0.14680874347686768, 0.09102652221918106, 0.07105451822280884, 0.06896264106035233], "utilization": [36.474609375, 57.6904296875, 87.2802734375, 90.4296875, 91.6015625, 91.259765625, 94.775390625, 98.8037109375, 44.7021484375], "perplexity": [1146.9842177240826, 1423.66411399845, 1681.6251963731768, 1519.6354622330805, 1329.361274648958, 1447.713187674207, 1756.7335627615232, 1906.7680339542387, 611.3863270546876]}
+{"step": 30000, "accuracy": [56.84166717529297, 66.74583435058594, 77.51249694824219, 87.80416870117188, 95.1312484741211, 98.02291870117188, 99.015625, 99.30208587646484, 99.32707977294922], "accuracy_text": [23.451589584350586, 41.019954681396484, 60.114559173583984, 78.37028503417969, 91.36548614501953, 96.49298095703125, 98.2538833618164, 98.76200866699219, 98.80635833740234], "ce": [3.0156617164611816, 2.2187554836273193, 1.4125587940216064, 0.7218500375747681, 0.29246342182159424, 0.14466850459575653, 0.08965033292770386, 0.06952819228172302, 0.06726808845996857], "utilization": [36.42578125, 57.2998046875, 88.18359375, 90.380859375, 91.5283203125, 90.966796875, 94.6044921875, 99.365234375, 44.384765625], "perplexity": [1152.170911136179, 1488.6486009826322, 1749.2406616511873, 1583.9389963130232, 1354.5228469281071, 1497.7616615306704, 1799.9907485907074, 2010.8431966535475, 665.9026279054835]}

t1_seed42/tokenizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f23b7dbc2a83f54cb3a850199e9b52fce4fb8bf01187d9fa2c57ea89861320a9
+size 349802315