Gate adapter weights for chili-lab/Ouro-hybrid-1.4B, trained using the
dragonling BDH multiplicative gating method.
This adapter replaces the additive FFN residual in a looped transformer with a
learned sparse multiplicative gate:
# original (Ouro-hybrid / LT2)
h = h + feed_forward(ffn_norm(h))
# dragonling gate
h = h + relu(gate_proj(ffn_norm(h))) * feed_forward(ffn_norm(h))
The gate is initialised to identity (gate_proj.weight = 0, gate_proj.bias = 1),
so at step 0 the patched model is bit-identical to the pretrained baseline.
Adapter contents
File
Description
adapter_weights.safetensors
Gate projection weights for all 24 transformer blocks (100.71 M params, ~200 MB bfloat16)
adapter_config.json
Training metadata (steps, lr, data, baseline / final loss)
The adapter does not include the base model weights. Download
chili-lab/Ouro-hybrid-1.4B separately.
Results summary
Five experiments across two domains. The gate consistently outperforms an
lm_head-only baseline at identical parameter budget (100.66 M vs 100.71 M).
Experiment
Data
Steps
Gate Δ (nats)
lm_head Δ (nats)
Gate advantage
Exp 2
tiny_shakespeare
500
+6.34
+5.33
+19 %
Exp 3
GSM8K full-seq
500
+5.96
+3.24
+84 %
Exp 4
GSM8K answer-only
1 000
+6.81
+4.65
+46 %
Exp 5
GSM8K answer-only
5 000
+7.11
—
—
This adapter corresponds to Experiment 5 (5 000 steps, GSM8K answer-only
loss, cosine LR decay 1e-3 → 5e-5). Val loss: 8.68 → 1.57 nats.
Note on accuracy: GSM8K exact-match accuracy is 0 % across all adapter
configurations. The gate adapter learns the format and style of mathematical
reasoning but does not restructure the attention mechanism needed to ground
reasoning in the specific input question. See the paper for full discussion.
1from bdh_lt2.student_model import register_student_model
2from bdh_lt2.gating import apply_bdh_gating
3from transformers import AutoModelForCausalLM, AutoTokenizer
4from safetensors.torch import load_file
56# 1. Register the custom model class and load the base model7register_student_model()8model = AutoModelForCausalLM.from_pretrained(9"chili-lab/Ouro-hybrid-1.4B",10 torch_dtype="auto",11 device_map="auto",12 trust_remote_code=True,13)14tokenizer = AutoTokenizer.from_pretrained(15"chili-lab/Ouro-hybrid-1.4B", trust_remote_code=True16)1718# 2. Inject gate architecture (identity init — bit-identical to base at this point)19apply_bdh_gating(model, init_identity=True)2021# 3. Load trained gate weights from this adapter22gate_sd = load_file("adapter_weights.safetensors")23model.load_state_dict(gate_sd, strict=False)2425# 4. Use the model26inputs = tokenizer("Q: What is 2 + 2?\nA:", return_tensors="pt").to(model.device)27out = model.generate(28**inputs,29 max_new_tokens=200,30 do_sample=False,31 repetition_penalty=1.3,32 use_cache=False,# GDN has no KV cache in the pure-PyTorch implementation33)34print(tokenizer.decode(out[0]))
Training details
Value
Trainable params
100.71 M gate_proj (6.25 % of 1,612 M total)
Frozen params
all 1,512 M pretrained backbone parameters
Data
GSM8K train split (7 473 Q/A pairs, 1.42 M tokens)
Loss
Answer-only (question tokens masked with label = -100)
Optimiser
AdamW, lr 1e-3 → 5e-5 cosine, weight_decay 0.01, grad_clip 1.0
Steps
5 000
Batch size
1 × seq_len 256
Hardware
Single RTX 5080 (Blackwell sm_120, 16 GB VRAM)
GDN implementation
Pure PyTorch sequential recurrence — no Triton required
License
This adapter is released for research use only, consistent with the license of
chili-lab/Ouro-hybrid-1.4B. The adapter weights (gate projections only) are
original trained parameters; no base model weights are included.
