GDN-340M-Kestrel

A pure Gated DeltaNet language model (340M parameters, 24 layers, no attention over past tokens) fine-tuned for multi-hop fact lookup over long fact lists. It accompanies a course lab on memory models, where it is compared with the recurrent memory Transformer irodkin/ARMT-Gemma-3-1B-IT-Kestrel on the "Kestrel handbook" task.

Contents

Path Model
repository root the fine-tuned model (this card)
base/ the base model linear-moe-hub/Gated-Deltanet-340M, converted to the weight layout of flash-linear-attention 0.5 (fused gate_proj split into gate_proj + up_proj; the unused attn.D parameter dropped), used as the starting point for fine-tuning in the lab

Both are released under Apache-2.0, like the base model.

Task format

Question: In which town is the building where the leader of Project X works?
Facts:
<facts, one per line>
Question: In which town is the building where the leader of Project X works?
Answer: <leader>; <building>; <town>

Training

Full fine-tuning for 8,000 examples (batch 8, AdamW, learning rate 1e-4 with cosine decay, bf16 autocast), on freshly generated random worlds with the same name pools as the Kestrel handbook but different links. Facts in random order, the question's chain facts in hop order, contexts of 20โ€“615 facts (uniformly sampled), mostly 3-hop questions. The maximum length matches the training length of the ARMT model above, for a fair comparison of extrapolation.

Usage

Needs flash-linear-attention (Triton โ‰ฅ 3.3) and transformers 5.x:

import torch, fla
from fla.models.gated_deltanet.modeling_gated_deltanet import GatedDeltaNetForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer

# fla declares tied weights in the transformers 4 format; transformers 5 expects a mapping.
GatedDeltaNetForCausalLM._tied_weights_keys = {"lm_head.weight": "model.embeddings.weight"}
tokenizer = AutoTokenizer.from_pretrained("irodkin/GDN-340M-Kestrel")
model = AutoModelForCausalLM.from_pretrained("irodkin/GDN-340M-Kestrel").cuda().eval()

With recent transformers, generate() mishandles fla's recurrent-state cache; decode greedily by hand, passing past_key_values from one step to the next:

@torch.no_grad()
def greedy(prompt, max_new_tokens=32):
    ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
    out = model(input_ids=ids, use_cache=True, logits_to_keep=1)
    new = []
    for _ in range(max_new_tokens):
        token = out.logits[:, -1].argmax(-1, keepdim=True)
        if token.item() == tokenizer.eos_token_id:
            break
        new.append(token.item())
        out = model(input_ids=token, past_key_values=out.past_key_values, use_cache=True)
    return tokenizer.decode(new)

Results

Exact match on the whole answer chain, greedy (teacher-forced), on the Kestrel handbook (seed 7, never seen in training). Excerpts keep the gold facts and are sorted into sections like the full handbook.

Context 1-hop 2-hop 3-hop
300 facts 1.00 1.00 0.97
600 facts (training maximum) 1.00 1.00 0.97
1,200 facts 1.00 1.00 0.77
Whole handbook, 4,500 facts (~66k tokens) 0.72 0.52 0.10

30 questions per cell (50 for the whole handbook). For comparison, ARMT trained on the same maximum length scores 0.92 / 0.88 / 0.68 on the whole handbook. In our runs a 1.3B Gated DeltaNet and LoRA fine-tuning did not extrapolate better.

Downloads last month
44
Safetensors
Model size
0.4B params
Tensor type
F32
ยท
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support

Model tree for irodkin/GDN-340M-Kestrel

Finetuned
(1)
this model