GitHub: https://github.com/YG3-ai/dragonling | Model: https://huggingface.co/YG3-ai/dragonling-gate-adapter | Paper: https://zenodo.org/records/21311787
Early release. Experiments are ongoing; results and code will grow here. Issues, questions, and pull requests are welcome.
A novel gating method. We adapt BDH's sparse elementwise-product primitive as a multiplicative gate on the FFN residual of looped transformers:
# original (LT2 / Ouro)
h = h + feed_forward(ffn_norm(h))
# dragonling — BDH multiplicative gate
h = h + relu(gate_proj(ffn_norm(h))) * feed_forward(ffn_norm(h))
The gate is initialised to 1 everywhere (gate_proj.weight = 0,
gate_proj.bias = 1) so a pretrained model is bit-identical at step 0 — any
improvement has to be earned. This identity-safe initialisation is what makes
the gating compatible with existing pretrained weights.
A novel implementation. bdh_lt2/student_model.py is a pure-PyTorch
StudentForCausalLM for chili-lab/Ouro-hybrid-1.4B. This class does not
exist in the LT2 repo or the HF model card; we derived it from the checkpoint
weight structure and implemented the GatedDeltaNet recurrence without Triton,
so it runs on any CUDA GPU including Blackwell where kernel compilation is
unreliable.
Honest empirical findings. Early results on both from-scratch and pretrained settings — mixed, and reported as-is. See Results.
LT2 (chili-lab) builds looped transformers
where the same parameters are reused across multiple forward passes. BDH
(pathwaycom) proposes that sparse, non-negative elementwise products —
relu(a) * relu(b) — are a powerful primitive for gating information flow.
We test whether the BDH multiplicative gate improves a looped transformer, both when training from scratch and when fine-tuning a pretrained 1.4B model.
A small looped transformer (2 layers × 4 loops, ~1.7 M params) trained from scratch on byte-level tiny_shakespeare. Control and experimental arms share identical seed / data / optimiser — the only difference is the gate.
| Loop count | Control (BPC) | Gated (BPC) | Δ (nats) |
|---|---|---|---|
| 1 | 2.915 | 2.927 | −0.008 |
| 2 | 3.008 | 3.051 | −0.030 |
| 4 | 2.481* | 2.572* | −0.064* |
| 8 | 2.708 | 2.800 | −0.064 |
* 4000 steps; others 2000 steps. BPC = bits per character = loss / ln 2.
Finding: additive residual wins from scratch, and the gap grows with loop count — the opposite of the hypothesis. We attribute this to the random gate initialisation compounding noise through more loop iterations, rather than an architectural inferiority (see Experiment 2).
chili-lab/Ouro-hybrid-1.4B is a 1.4 B-parameter looped transformer
(24 layers, total_ut_steps = 4) distilled from ByteDance/Ouro-1.4B.
All pretrained parameters are frozen; only the injected gate projections
(100.71 M parameters, 6.25 % of total) are trained. The baseline is the
pretrained model with the identity-init gate — equivalent to no gate.
| Method | Params trained | Steps | Val loss (nats) | Δ from baseline |
|---|---|---|---|---|
| Baseline (pretrained, no fine-tune) | — | 0 | 9.76 | — |
| lm_head only | 100.66 M (6.24 %) | 200 | 7.22 | +2.67 |
| Gate only | 100.71 M (6.25 %) | 200 | 5.07 | +4.69 |
| Gate only | 100.71 M (6.25 %) | 500 | 4.43 | +5.42 |
Data: tiny_shakespeare (tokenised with the Ouro BPE tokeniser, 90 / 10 split).
Finding: at an identical parameter budget, gate-only fine-tuning achieves 75 % more improvement than lm_head-only fine-tuning at 200 steps, and the gap widens at 500 steps. The gate weight RMS at step 500 is 0.020 — tiny perturbations from identity producing large loss reductions. The gate's distributed position across all 24 blocks (vs. the lm_head touching only the final projection) is the likely driver; the result is still being analysed.
Note on the baseline loss: 9.76 nats reflects the pretrained model being out-of-distribution for Shakespeare (it was trained on mathematical reasoning data). The lm_head and gate experiments are both domain adaptation; their comparison isolates architectural efficiency, not absolute quality.
The fine-tuned gate weights (100.71 M params, ~200 MB bfloat16) are released as a standalone adapter — you need the base model separately:
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
register_student_model()
model = AutoModelForCausalLM.from_pretrained(
"chili-lab/Ouro-hybrid-1.4B", torch_dtype="auto", device_map="auto"
)
apply_bdh_gating(model, init_identity=True) # inject gate architecture
gate_sd = load_file("adapter_weights.safetensors") # download from HF
model.load_state_dict(gate_sd, strict=False) # apply trained gatesSee MODEL_CARD.md for full details.
bdh_lt2/
gating.py ← GatedFeedForward · apply_bdh_gating · gate_parameters
micro_model.py ← tiny from-scratch looped transformer (no CUDA kernels)
student_model.py ← pure-PyTorch StudentForCausalLM for Ouro-hybrid-1.4B
(no Triton; sequential GDN recurrence)
experiments/
train_micro.py ← Experiment 1: control vs gated, writes results/
finetune_ouro_gate.py ← Experiment 2a: gate-only fine-tune
finetune_ouro_lmhead.py ← Experiment 2b: lm_head comparison baseline
scripts/
load_ouro_gated.py ← identity-init sanity check on the real model
eval_gsm8k.py ← GSM8K accuracy eval (base model or gate adapter)
download_models.sh ← fill the models/ slot (Linux/Mac)
data/
download_tiny_shakespeare.py ← stdlib-only data fetch
download_gsm8k.py ← GSM8K train + test splits (needs: pip install datasets)
models/
Ouro-hybrid-1.4B/ ← download separately (huggingface-cli)
dragonling-gate-adapter/ ← written by finetune_ouro_gate.py --save-adapter
MODEL_CARD.md ← HuggingFace model card for the gate adapter
results/ ← JSON output from each run (gitignored; commit selectively)
LT2/ ← vendored chili-lab/LT2 (upstream)
bdh/ ← vendored pathwaycom/bdh (upstream reference)
- CUDA 12.8+ GPU (tested on Blackwell / sm_120; should work on Ampere / Ada)
- Python 3.11+
- PyTorch 2.7+ with matching CUDA wheels
# Install with CUDA 12.8 wheels (adjust for your CUDA version):
pip install torch --index-url https://download.pytorch.org/whl/cu128
pip install -r requirements.txtNo Triton required.
bdh_lt2/student_model.pyimplements GatedDeltaNet with a pure-PyTorch sequential recurrence so the real-model path runs on any CUDA GPU, including Blackwell where Triton compilation is unreliable.
python data/download_tiny_shakespeare.py
# Full sweep: both arms, all loop counts
python experiments/train_micro.py --steps 2000 --loops 1 --tag loops1
python experiments/train_micro.py --steps 2000 --loops 2 --tag loops2
python experiments/train_micro.py --steps 2000 --loops 4 --tag loops4
python experiments/train_micro.py --steps 2000 --loops 8 --tag loops8
# Or just the canonical run:
python experiments/train_micro.py --steps 4000 --tag long# Download weights (~3 GB)
huggingface-cli download chili-lab/Ouro-hybrid-1.4B \
--local-dir models/Ouro-hybrid-1.4B
# Verify identity-init (should print delta ≈ 0)
python scripts/load_ouro_gated.py
# Gate-only fine-tune on tiny_shakespeare
python experiments/finetune_ouro_gate.py --steps 500 --data shakespeare
# lm_head comparison (matched parameter budget)
python experiments/finetune_ouro_lmhead.py --steps 200 --data shakespeareOuro was trained on mathematical reasoning. Fine-tuning on GSM8K keeps the model in-domain, separating gate architectural benefit from domain-adaptation benefit.
| Method | Params trained | Steps | Val loss (nats) | Δ from baseline |
|---|---|---|---|---|
| Baseline (pretrained) | — | 0 | ~8.9 | — |
| lm_head only | 100.66 M | 500 | 5.66 | +3.41 |
| Gate only | 100.71 M | 500 | 2.94 | +5.88 |
Finding: the gate achieves 72 % more improvement than lm_head at identical parameter budget and steps — nearly identical to the Shakespeare result (75 %). The gate's advantage is domain-invariant, ruling out the hypothesis that it mainly suppresses out-of-distribution FFN activations. Accuracy: 0 % (model generates unrelated math problems rather than answering the given question — full-sequence loss teaches text distribution, not Q→A mapping).
Answer tokens (66 % of the stream) are the only positions that receive loss. This teaches the model to produce answers given questions, not just to model the Q+A text distribution.
| Method | Params trained | Steps | Val loss (nats) | Δ from baseline |
|---|---|---|---|---|
| Baseline (pretrained) | — | 0 | ~8.93 | — |
| lm_head only | 100.66 M | 1 000 | 4.27 | +4.65 |
| Gate only | 100.71 M | 1 000 | 2.14 | +6.81 |
Finding: the gate achieves 46 % more improvement than lm_head (+2.16 nats
absolute gap, the largest yet). Answer-only masking fixed format collapse — the
model now outputs #### N rather than generating unrelated problems — but 1 000
steps is insufficient for arithmetic reasoning (2 % GSM8K accuracy). The model
collapses toward common training answers rather than computing. Gate weights had
not converged at step 1 000 (rms still rising); longer training warranted.
Same setup as Experiment 4 with 5× more steps and a cosine LR schedule (1e-3 → 5e-5), which stabilises the late-training oscillation seen at 1 000 steps.
| Method | Params trained | Steps | Val loss (nats) | Δ from baseline |
|---|---|---|---|---|
| Baseline (pretrained) | — | 0 | ~8.68 | — |
| Gate only | 100.71 M | 5 000 | 1.57 | +7.11 |
Gate weight RMS plateaued at 0.034 around step 3 000 — weights converged. Best single val reading: 1.32 nats at step 3 950. Train loss at step 5 000: 1.07 nats (some overfitting, expected on a 1.4 B model with batch size 1). GSM8K accuracy: 0 % (0/50).
Key finding — format vs. reasoning. At 5 000 steps the model generates
coherent GSM8K-style chain-of-thought — <<expr=result>> annotations, multi-step
reasoning, #### N answers — but about a hallucinated problem, not the actual
question asked. The gate learned the genre and format of math reasoning; it did
not restructure the attention mechanism that grounds reasoning in the specific
input question. This is consistent with what 100 M gate-only parameters can
achieve: style and format adaptation, not new reasoning capability. Full-parameter
fine-tuning or a much larger adapter budget would be needed to close the gap.
# Download GSM8K (~7 MB)
python data/download_gsm8k.py
# Experiments 3 & 4
python experiments/finetune_ouro_gate.py --steps 500 --data gsm8k --tag gate_gsm8k
python experiments/finetune_ouro_lmhead.py --steps 500 --data gsm8k --tag lmhead_gsm8k
python experiments/finetune_ouro_gate.py --steps 1000 --data gsm8k --answer-only --tag gate_gsm8k_ao
python experiments/finetune_ouro_lmhead.py --steps 1000 --data gsm8k --answer-only --tag lmhead_gsm8k_ao
# Experiment 5 — extended run with cosine LR decay
python experiments/finetune_ouro_gate.py --steps 5000 --data gsm8k --answer-only --tag gate_gsm8k_5k
# Accuracy eval (reads models/dragonling-gate-adapter/ written by gate fine-tune)
python scripts/eval_gsm8k.py --adapter models/dragonling-gate-adapter --tag gate_5kSpeed note: GSM8K generation is slow without a KV cache for the GDN layers (~1–3 min per example). The default
--max-examples 50takes roughly 1–2 hours. Chunked parallel GDN (on the roadmap) would cut this to minutes.
Planned next experiments, in rough priority order:
- Identity-init from-scratch micro — does initialising the gate at 1 (not random) close the gap in Experiment 1? This isolates whether optimisation noise is the true cause of the micro regression.
- Bigger micro model (dim=512, n_layers=4, ~25 M params) — more realistic scale before the real model.
- Longer gate fine-tune — 5 000 steps with cosine LR decay completed (Experiment 5); val loss 1.57 nats, gate weights converged. Accuracy eval pending.
- LR sweep (1e-4, 3e-4, 1e-3, 3e-3) — gate learning-rate sensitivity.
- LoRA comparison — another matched-budget baseline beyond lm_head.
- Gate activation analysis — histogram of trained gate values across the corpus: do gates stay near 1 or learn strongly bimodal patterns?
- Multiple seeds — statistical significance for all key comparisons.
- Different domain — GSM8K in-domain experiment added (Experiment 3 above).
- Chunked parallel GDN — replace the sequential Python loop in
student_model.pywith a parallel prefix-scan for faster training.
- LT2: chili-lab/LT2 — arXiv 2605.20670
- BDH: pathwaycom/bdh — arXiv 2509.26507
- Ouro-hybrid-1.4B: chili-lab/Ouro-hybrid-1.4B (distilled from ByteDance/Ouro-1.4B, arXiv 2510.25741)