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

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

Paper · Project page · Checkpoints

Table of Contents

• Overview
• Requirements
• Checkpoints and Datasets
• Training
• Evaluation
• Results
• Citation
• Acknowledgement
• License

Overview

DAMI allocates memory timescalesacross heads and input responses across layers to improve long-context retrieval. This repository supports all 15 combinations of GDN / KDA / GDN2 / Mamba-2 / RWKV-7 and Random / PoST / DAMI through one training entry point.

Requirements

Clone the repository and enter its root directory:

git clone https://github.com/SpeedAI-Lab/DAMI.git
cd DAMI

Use Python 3.12 and Linux/WSL with an NVIDIA CUDA GPU. Install the appropriate PyTorch 2.9 CUDA wheel first, then:

pip install setuptools wheel packaging ninja
pip install --no-build-isolation -r requirements.txt

Model implementations and kernels come from FLA and Mamba; the training loop uses Transformers Trainer. CUDA build tools may be needed when compatible wheels are unavailable.

Checkpoints and Datasets

Included: 5,500 fixed samples across 11 task/length settings and 82,500 original predictions from the 15 paper checkpoints.

Model weights: SpeedAI-Lab/DAMI on Hugging Face. Access requires authentication and repository permission while the weight repository is private. Weights are stored under checkpoints/; the checkpoint directory names are:

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

From the code repository root, download the checkpoints while preserving their directory structure:

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/*"]. The resulting layout is:

checkpoints/
├── main_gdn_random/
├── main_gdn_PoST/
├── main_gdn_DAMI/
└── ...

Each model directory should contain config.json, model.safetensors, and tokenizer files. Run the following commands from the repository root; replace checkpoints with your own checkpoint directory if needed.

The fixed evaluation samples are included in data/ruler/, and the original predictions are in reference/ruler_predictions/. Training data is supplied separately through --train-files.

Training

Select an architecture and initialization:

python train.py --architecture gdn --initialization dami \
  --train-files '/data/train-*.parquet' --tokenizer TinyLlama/TinyLlama_v1.1 \
  --output-dir outputs/gdn/dami

• --architecture: gdn, kda, gdn2, mamba2, rwkv7.
• --initialization: random, post, dami.

Run all 15 paper presets sequentially, with separate output directories:

python train.py --all \
  --train-files '/data/train-*.parquet' --tokenizer TinyLlama/TinyLlama_v1.1 \
  --output-dir outputs/main

Defaults: approximately 100M parameters, 4096-token context, global batch 128, and 9537 steps (approximately 5B tokens per run). The training seed is 42 except for RWKV-7 + DAMI, which uses 43. --all --dry-run previews the matrix. Set CUDA_VISIBLE_DEVICES to select a GPU; use torchrun for distributed training of an individual combination. PoST denotes the paper's spectral-only control.

The original experiments used FLAME; these training entry points use Transformers Trainer. Paper initialization parameters, fixed evaluation data, and recorded predictions are included.

Evaluation

Check the downloaded models without running inference (Windows/CPU supported):

python evaluate_ruler.py --preflight --all --models-root checkpoints

Recompute the original main table without models or a GPU:

python evaluate_ruler.py --rescore-reference --all --output-dir outputs/ruler-reference

Re-run all 15 checkpoints in Linux/WSL with CUDA:

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

Use --architecture and --initialization to evaluate one combination. The evaluator preserves the original prompts and uses greedy bf16 decoding with 128 new tokens and the original batch sizes. It saves predictions, ruler_results.csv, and comparison.json. Rescoring the bundled predictions reproduces all 165/165 recorded cells. Optional --check-weight-hashes verifies the complete checkpoint SHA-256 hashes.

Results

All models are approximately 100M parameters and trained for 5B tokens. DAMI improves average retrieval accuracy over Random by 13.95 / 16.71 / 13.44 / 14.47 / 5.29 percentage points for GDN / KDA / GDN2 / Mamba-2 / RWKV-7.

NIAH-1/2 use 1K, 2K, 4K, and 8K contexts; MK-NIAH uses 1K, 2K, and 4K. Each setting has 500 samples. Scores use case-insensitive answer containment, and Avg. weights all 11 settings equally. Numeric results are available in assets/retrieval.csv.

Citation

If you use DAMI in your research, please cite:

@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

Original repository code is released under the MIT License. Benchmark data retains its original attribution and applicable terms.