dragonling-gate-adapter

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.

---

How to use

Requires bdh_lt2/ from the dragonling repo.

from bdh_lt2.student_model import register_student_model
from bdh_lt2.gating import apply_bdh_gating
from transformers import AutoModelForCausalLM, AutoTokenizer
from safetensors.torch import load_file

# 1. Register the custom model class and load the base model
register_student_model()
model = AutoModelForCausalLM.from_pretrained(
    "chili-lab/Ouro-hybrid-1.4B",
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(
    "chili-lab/Ouro-hybrid-1.4B", trust_remote_code=True
)

# 2. Inject gate architecture (identity init — bit-identical to base at this point)
apply_bdh_gating(model, init_identity=True)

# 3. Load trained gate weights from this adapter
gate_sd = load_file("adapter_weights.safetensors")
model.load_state_dict(gate_sd, strict=False)

# 4. Use the model
inputs = tokenizer("Q: What is 2 + 2?\nA:", return_tensors="pt").to(model.device)
out = model.generate(
    **inputs,
    max_new_tokens=200,
    do_sample=False,
    repetition_penalty=1.3,
    use_cache=False,   # GDN has no KV cache in the pure-PyTorch implementation
)
print(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.

---

Paper

Carter, J. & Knox, S. (2026). Multiplicative Gate Adapters for Looped Transformers. Zenodo. https://zenodo.org/records/21311787

Citation

@misc{dragonling2026,
  title  = {Multiplicative Gate Adapters for Looped Transformers},
  author = {Carter, Jackie and Knox, Samuel},
  year   = {2026},
  url    = {https://zenodo.org/records/21311787}
}

---

Upstreams

• LT2 — chili-lab/LT2, arXiv 2605.20670
• BDH — pathwaycom/bdh, arXiv 2509.26507
• Ouro — ByteDance/Ouro-1.4B, arXiv 2510.25741