ruhai-lin/GDNBit

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.

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.