DAMI: Dual-Axis Memory Initialization for Enhanced Long-Context Retrieval in Gated Linear Attention

Zhangchi Feng · Harry Jiang · Yangyifei Luo · Ray Zhang · Richong Zhang

Project page · Code, evaluation data and scripts

DAMI allocates memory timescales across heads and input responses across layers to improve long-context retrieval. This repository contains the 15 paper checkpoints: five architectures (GDN, KDA, GDN2, Mamba-2, RWKV-7) with three initialization conditions (Random, PoST, DAMI). Each model has approximately 100M parameters and was trained for approximately 5B tokens with a 4096-token training context.

Checkpoints

All checkpoints are stored under checkpoints/, using the directory names below. Each directory contains the full model weights (model.safetensors), config.json, generation_config.json, tokenizer.json, and tokenizer_config.json.

Architecture	Random	PoST	DAMI
GDN	main_gdn_random	main_gdn_PoST	main_gdn_DAMI
KDA	main_kda_random	main_kda_PoST	main_kda_DAMI
GDN2	main_gdn2_random	main_gdn2_PoST	main_gdn2_DAMI
Mamba-2	main_mamba2_random	main_mamba2_PoST	main_mamba2_DAMI
RWKV-7	main_rwkv7_random	main_rwkv7_PoST	main_rwkv7_DAMI

Download and evaluation

Clone the code repository and install its dependencies as described in its README. From the code repository root, download the checkpoints:

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="SpeedAI-Lab/DAMI",
    local_dir=".",
    allow_patterns=["checkpoints/**"],
)

To download only GDN + DAMI, use allow_patterns=["checkpoints/main_gdn_DAMI/*"]. Access to this repository requires authentication while it is private.

Run all 15 checkpoints on the fixed RULER evaluation data in Linux/WSL with CUDA:

python evaluate_ruler.py --preflight --all --models-root checkpoints
python evaluate_ruler.py --all --models-root checkpoints \
  --output-dir outputs/ruler-retest --skip-done

For a single checkpoint, replace --all with --architecture gdn --initialization dami. The evaluator includes the architecture-specific FLA and RWKV adaptations required by these checkpoints; use the provided loading path for evaluation.

The code repository includes 5,500 fixed samples across 11 RULER task/length settings and 82,500 original predictions. Evaluation uses greedy decoding with up to 128 new tokens. Scores use case-insensitive answer containment; the average weights all 11 settings equally.

RULER results

Average accuracy (%), as reported in the paper:

Architecture	Random	PoST	DAMI
GDN	30.42	33.49	44.36
KDA	26.38	33.67	43.09
GDN2	30.27	33.47	43.71
Mamba-2	9.95	13.73	24.42
RWKV-7	23.87	20.69	29.16

Citation

@article{feng2026dami,
  title={DAMI: Dual-Axis Memory Initialization for Enhanced Long-Context Retrieval in Gated Linear Attention},
  author={Feng, Zhangchi and Jiang, Harry and Luo, Yangyifei and Zhang, Ray and Zhang, Richong},
  journal={arXiv preprint arXiv:2610.41994},
  year={2026}
}

Acknowledgement

We thank the authors and developers of FLA, RWKV, GatedDeltaNet, GatedDeltaNet-2, Kimi-Linear, and Mamba for their contributions to the architectures and implementations used in this work. We also thank Transformers and FLAME for their training infrastructure.

Our retrieval evaluation uses RULER. Benchmark attribution and data notices are provided in data/ruler/NOTICE.md. We thank the authors and developers for making their work available. Please cite the corresponding papers when using these resources.

License

These checkpoints are released under the MIT License. See the LICENSE file. Benchmark attribution and data notices are provided in the code repository.