GatedDeltaNet-360M-15B-SlimPajama

This model is for research purposes only and is not intended for production use.

A GatedDeltaNet language model (357.8M parameters) pretrained from scratch on 15B tokens from SlimPajama.

GatedDeltaNet combines scalar decay gating with the delta rule for selective memory writing. See Yang et al., 2024.

Trained with flash-linear-attention and Flame.

Usage

Requires: pip install flash-linear-attention

import torch
import fla.models  # registers GatedDeltaNet with HuggingFace Auto classes

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "puigde/gated-deltanet-360M-15B-slimpajama",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).cuda()
tokenizer = AutoTokenizer.from_pretrained(
    "puigde/gated-deltanet-360M-15B-slimpajama"
)

inputs = tokenizer("The capital of France is", return_tensors="pt").to("cuda")
output = model.generate(**inputs, max_new_tokens=50, do_sample=False)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The import fla.models line registers the GatedDeltaNet architecture with HuggingFace's Auto classes. Without it, from_pretrained will fail with an unknown model type error.

Architecture

Based on the Flame gated_deltanet_340M.json config with two modifications: expand_v=1 (vs 2) and num_heads=4 (vs 6), which reduces the parameter count from ~460M to 358M.

Training

Tokenizer: LlamaTokenizer (from fla-hub/gla-1.3B-100B), vocab 32,000.

Evaluation

Zero-shot, lm-evaluation-harness:

RULER (2k context): S1=1.00, S2=1.00, S3=0.66, MK1=0.32, avg=0.75.

Citation

@article{yang2024gateddeltanet,
  title={Gated Delta Networks: Improving Mamba2 with Delta Rule},
  author={Yang, Songlin and Keller, Jan and Wang, Bailin and others},
  journal={arXiv preprint arXiv:2412.06464},
  year={2024}
}