File size: 7,534 Bytes
354d3f9
8bc871c
 
 
 
 
 
 
354d3f9
 
 
8bc871c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
---
library_name: transformers
tags:
- gated-deltanet
- ternary
- bitnet
- looped-transformer
- babylm
---
# GDN-Bit

GDN-Bit is a Gated DeltaNet whose linear layers are all ternary BitLinear (absmean ternary weights, int8 per-token activations). It optionally loops the whole layer stack T times. This repository holds the canonical 30M checkpoints trained on BabyLM 2026 Strict-Small:
- GDN-Bit
- full-precision GDN with the same architecture
- HGRN-Bit (the MatMul-free LM mixer)

Each comes without (T1) and with (T4) the loop, three seeds each. Folders are named `<arch>-30m[-k<dk>]-t<T>/babylm-strict-small-100Mw-s<seed>/`.

| Folder | Mixer | Linear weights | dk | T | Params | Peak LR |
|---|---|---|---:|---:|---:|---:|
| `gdnbit-30m-t1` | Gated DeltaNet | ternary | 64 | 1 | 30,470,472 | 2e-3 |
| `gdnbit-30m-t4` | Gated DeltaNet | ternary | 64 | 4 | 30,472,520 | 2e-3 |
| `gdn-30m-t1` | Gated DeltaNet | full precision | 64 | 1 | 30,453,064 | 1e-2 |
| `gdn-30m-t4` | Gated DeltaNet | full precision | 64 | 4 | 30,455,112 | 1e-2 |
| `hgrnbit-30m-t1` | HGRN (MatMul-free LM) | ternary | – | 1 | 30,566,304 | 2e-3 |
| `hgrnbit-30m-t4` | HGRN (MatMul-free LM) | ternary | – | 4 | 30,568,352 | 2e-3 |
| `gdnbit-30m-k16-t1` / `-t4` | Gated DeltaNet | ternary | 16 | 1 / 4 | 28,701,000 / 28,703,048 | 2e-3 |
| `gdnbit-30m-k32-t1` / `-t4` | Gated DeltaNet | ternary | 32 | 1 / 4 | 29,290,824 / 29,292,872 | 2e-3 |

**Architecture** (GDN and GDN-Bit):
- d=512, 6 layers, 6 heads, dk=64, dv=128, ReLU² FFN of width 2048, no short convolution.
- One fused input projection for q/k/v/g/a/b, sharing a single RMSNorm and int8 quantization.
- Tied float embedding and output head; the head is ternary in the Bit models.
- The K ablation keeps dv=128 and changes only dk.

**HGRN-Bit:** d=512, 7 layers, SwiGLU 1376, matched to ~30.5M parameters.

**Loop:** the T4 models apply the whole stack 4 times with weights shared across passes, as h ← Stack(h) + r_t ⊙ h_in. The per-pass vector r_t is fp32 and zero-initialized. Every pass starts its recurrent state from zero.

## Loading

The weights load with Transformers once the model classes are registered. This needs the `custom_models` package of the project code at commit `1bf7e65` (not public yet), plus `flash-linear-attention==0.5.1` and `transformers==4.57.3`; trained with torch 2.13. Plain Transformers does not recognize the model types `gdnbit` and `mmfloop`.

```python
import custom_models  # registers gdnbit (GDN, GDN-Bit) and mmfloop (HGRN-Bit)
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "ruhai-lin/GDNBit"
run = "gdnbit-30m-t4/babylm-strict-small-100Mw-s42"
tokenizer = AutoTokenizer.from_pretrained(repo, subfolder=run)
model = AutoModelForCausalLM.from_pretrained(repo, subfolder=run, dtype="float32")
```

**Warning for T4 models:** the HF recurrent cache is keyed by layer index only, so the 4 passes would overwrite each other's state. Score full sequences without a cache, or keep one state per (pass, layer) yourself. Do not use `generate()` with `use_cache=True`.

## Training, evaluation and budgets

**Training**
- **Data:** BabyLM 2026 Strict-Small, 10M words, tokenized with the official BabyLM 16,384-token tokenizer, which is included in every folder.
- **Budget:** 10 epochs = 100M words processed, 10,320 steps at batch 32 × 512 tokens.
- **Optimizer:** AdamW (β 0.9/0.95, weight decay 0.1, gradient clip 1.0), cosine schedule to 0 with 500 warmup steps, bf16 autocast.
- **Learning rate:** the peak LR was swept separately for ternary and full-precision models (see `eval/all_runs.csv`).

**Checkpoint selection:** a checkpoint is saved every 10M words. Each folder holds the checkpoint with the lowest full dev NLL; `eval/final_summary.json` gives `best_words_m`.

**Metrics**
- Dev/test NLL is per token, FP32, on consecutive BOS-prefixed 512-token windows over the full BabyLM dev/test sets. bpw is bits per word.
- BLiMP, BLiMP supplement and entity tracking come from the official BabyLM evaluation pipeline (zero-shot, causal, filtered sets).
- Entity tracking varies strongly across seeds; we do not draw conclusions from it.

## Results (mean ± std over seeds 7, 42, 1337)

| Model | Dev NLL | Test NLL | BLiMP | BLiMP suppl. |
|---|---:|---:|---:|---:|
| gdnbit-30m-t1 | 3.554 ± 0.009 | 3.483 ± 0.013 | 69.8 ± 1.3 | 58.9 ± 1.0 |
| **gdnbit-30m-t4** | **3.531 ± 0.050** | **3.454 ± 0.063** | **71.2 ± 0.0** | **60.6 ± 0.5** |
| gdn-30m-t1 | 3.431 ± 0.040 | 3.357 ± 0.054 | 70.6 ± 0.7 | 62.1 ± 2.4 |
| gdn-30m-t4 | 3.365 ± 0.010 | 3.274 ± 0.016 | 71.7 ± 0.4 | 60.7 ± 0.8 |
| hgrnbit-30m-t1 | 3.571 ± 0.032 | 3.476 ± 0.037 | 70.3 ± 0.7 | 58.7 ± 0.3 |
| hgrnbit-30m-t4 | 3.562 ± 0.018 | 3.468 ± 0.025 | 70.2 ± 0.8 | 60.2 ± 0.6 |
| gdnbit-30m-k16-t1 | 3.624 ± 0.019 | 3.562 ± 0.027 | 69.5 ± 0.5 | 60.6 ± 1.5 |
| gdnbit-30m-k16-t4 | 3.535 ± 0.019 | 3.454 ± 0.027 | 70.3 ± 0.0 | 59.8 ± 0.9 |
| gdnbit-30m-k32-t1 | 3.574 ± 0.021 | 3.499 ± 0.030 | 69.9 ± 0.6 | 59.5 ± 0.8 |
| gdnbit-30m-k32-t4 | 3.524 ± 0.015 | 3.444 ± 0.018 | 70.8 ± 0.4 | 60.1 ± 2.9 |

**Result files**
- `eval/comparison.csv`: all metrics with mean and std, including bpw and entity tracking.
- `eval/all_runs.csv`: every run-seed, including the full-precision LR sweep points, whose weights are not published.
- `eval/state_quant.csv`: dev NLL when the recurrent state is stored at reduced precision after every token (seed 1337). An int8 state with one scale per row costs at most 0.004 nats for every model.

## Files in each run folder

- **HF checkpoint:** `config.json`, `model.safetensors` (fp32 latent weights), the tokenizer files and `generation_config.json`.
- **`eval/`:**
  - `final_summary.json`: dev/test NLL and bpw, BLiMP, supplement, entity tracking, best checkpoint
  - `babylm_eval_summary.json` and per-task reports: `blimp_report.txt`, `blimp_supplement_report.txt`, `entity_tracking_report.txt` (per-field and per-phenomenon accuracy)
  - `args.json`: training arguments
  - `dev_curve.log`: dev NLL at every epoch
- **`ternary/`** (`gdnbit-30m-t1` and `gdnbit-30m-t4` only): deployment weights. See below.

## Ternary deployment weights (`ternary/`)

The latent fp32 weights must not be re-quantized independently. About a fifth of them sit exactly on the rounding threshold |w| = 0.5·mean|w|, so a different floating-point summation order (for example CPU vs GPU) flips their ternary level and changes the model.

`ternary/ternary.safetensors` stores, for every BitLinear `<module>`, the levels and scale as computed on CUDA by the training kernel:
- `<module>.levels`: int8 values in {-1, 0, 1}, shape [out, in]
- `<module>.scale`: fp32 scalar, mean|W|

The forward pass is

    y = Q8(RMSNorm(x; eps 1e-6) * norm.weight) @ (levels * scale)^T,

where Q8 is per-token absmax int8 quantization (127 / max|x|). The RMSNorm weights, embedding, A_log, dt_bias, output norm and loop weights r_t are the fp32 tensors in `model.safetensors`.

The fused attention input projection `model.layers.<i>.attn.in_proj` has 2316 output rows. In order: q 384 | k 384 | v 768 | g 768 | a 6 | b 6.
- a (decay) and b (write strength) are its last 12 rows. They share the matrix's single scale.
- After the projection: decay = exp(−exp(A_log)·softplus(a + dt_bias)), β = sigmoid(b), and q and k are L2-normalized.

The tied `lm_head` levels are quantized from the float embedding.

`ternary/verify.json` records the check: dev NLL of the model rebuilt from these files minus the native model's.