Raven

Raven is a linear-time sequence model built on top of Flash Linear Attention. It introduces a routing memory mechanism that selectively writes to a fixed set of persistent memory slots using a learned sparse router — achieving sub-quadratic complexity while maintaining strong associative recall.

---

SelectingMemory Update: Routed RWKV-7 Memory

This fork adds RWKV-7 sequence mixers to Raven so the original Raven routing-memory mixer can be compared against RWKV variants inside the same RavenForCausalLM model API.

Implemented in this repo:

• sequence_mixer="rwkv7": original dense RWKV-7 TimeMix adapted from HRM-Text, using the LT2 CUDA backend when available.
• sequence_mixer="routed_rwkv7": Raven-style top-k router gating dense RWKV channel groups.
• sequence_mixer="slot_rwkv7": explicit Raven-level routed RWKV memory slots with full num_slots x D x D per-head state.
• sequence_mixer="low_rank_slot_rwkv7": explicit routed RWKV slots with lower-cost num_slots x rank x D per-head state.
• low_rank_slot_rwkv7_backend="triton_fused": train-ready Triton forward plus reverse-scan backward for the low-rank slot RWKV path.
• examples/compare_mixers.py, examples/run_compare_matrix.py, and examples/compare_hrm_text_losses.py: speed and held-out loss comparison scripts.

Local HRM-Text held-out loss comparison on RTX 4090, trained from scratch on /home/xiaol/X/hrm_text_subset_1B for 500 steps with hidden=256, layers=2, batch=4, seq=128, slots=64, topk=16, rank=8:

mixer	params M	initial val	final val	best val	final train	tok/s	peak GB
raven	35.39	11.1281	2.9195	2.9195	2.5000	49,627	0.448
low_rank_slot_rwkv7	35.56	11.0938	2.6273	2.6273	1.9531	52,854	0.991

In this local train-ready comparison, low-rank slot RWKV reaches lower validation loss and slightly higher throughput than Raven, while using more VRAM because the fused backward caches low-rank slot states. Full commands and implementation details are in the RWKV-7 sequence mixers section.

---

Sparse Memory Routing in Raven. Unlike SSMs that update the entire state densely, or SWA that enforces strict FIFO overwriting, Raven uses an input-dependent router. At each step, only a specific subset of memory slots (highlighted) is selected to undergo decay and receive new information. Unselected slots remain completely untouched, preventing interference and preserving long-range recall.

---

Architecture

Raven replaces the SSM block with an RSM (Routing State Model) layer. Unlike GLA/Mamba2 which write to all memory slots uniformly, Raven learns a per-token sparse router R that selects which slots to update.

---

How Raven Works

Each Raven layer maintains a matrix memory state H ∈ R^(slots × d_v). At each timestep the router selects the top-k most relevant slots and performs a gated update:

route_scores = TopK( sigmoid(r_proj(x)) )
decay         = exp( route_scores * f )     # sparse forgetting gate
H             = H * decay + (1 - decay) * k ⊗ v
o             = q · H                       # read

The table below places Raven in the broader landscape of linear models:

Key design choices:

• Sparse top-k routing — each token writes to a small subset of memory slots
• Gumbel noise during training for exploration (optional)
• Mamba2 or GLA decay for the forgetting gate
• Chunked Triton kernels for training, fused recurrent kernels for generation

---

Results

In-Context Recall Benchmarks

Table 2: In-context recall benchmarks and NIAH accuracy vs. context length. Accuracy (%) on SWDE/FDA/SQuAD and single NIAH-1/2/3 across context lengths. Bold = best, underline = second best.

~400M parameter models

Model	Params	Mem (M)	SWDE	FDA	SQuAD	N1-1K	N1-2K	N1-4K	N1-8K	N1-16K	N1-32K	N2-1K	N2-2K	N2-4K	N2-8K	N2-16K	N2-32K	N3-1K	N3-2K	N3-4K	N3-8K	N3-16K	N3-32K
Transformer
w. RoPE	340	∞ / 0	42.3	34.5	22.1	100	100	0	0	0	0	100	100	0	0	0	0	71.6	47.6	0	0	0	0
w. Gate (FoX)	376	∞ / 0	52.5	64.3	30.1	100	100	32.2	8.0	4.2	0	100	100	100	24.0	11.6	3.2	95.4	85.6	64.2	11.6	7.2	0
SSM
GLA	475	12.5 / 0.4	29.0	11.4	30.3	74.6	25.1	8.2	2.2	0	0	91.2	37.2	21.4	3.6	0	0	84.2	57.1	20.8	10.2	2.3	0
GSA	399	12.5 / 0	23.8	14.5	24.9	99.2	97.1	90.0	67.4	29.6	11.0	96.6	98.8	28.0	5.1	1.0	0	60.0	30.1	13.5	1.0	0	0
GDN	475	12.5 / 0.4	29.5	8.3	31.3	99.2	100	99.8	92.0	41.8	22.1	99.2	92.0	43.6	17.8	6.2	4.0	92.6	80.6	37.8	5.2	6.8	2.5
Mamba-2	382	12.5 / 0.4	25.7	14.9	31.9	99.2	95.6	52.2	12.8	5.4	2.8	99.8	98.0	68.2	15.4	4.4	3.8	53.4	53.6	17.4	1.8	2.2	3.2
SWA	374	12.5 / 0	10.0	14.4	29.7	29.8	11.0	6.2	3.4	1.2	0	36.2	14.4	10.2	3.8	3.2	0	26.2	9.2	7.4	1.4	1.8	0
Raven	424	12.5 / 0	34.1	22.7	35.4	99.8	100	99.8	99.8	99.4	91.4	98.8	98.0	98.8	81.6	23.0	8.8	76.8	43.6	13.4	1.0	0	0

Language Modeling & Zero-Shot Evaluation

Table 3: Language modeling and zero-shot evaluation results. Perplexity on Lambada (LMB.) and zero-shot accuracy across standard benchmarks. Bold = best, underline = second best.

~400M parameter models

Model	Params	LMB. ppl↓	LMB. acc↑	PIQA↑	Hella.↑	Wino.↑	ARC-e↑	ARC-c↑	Avg.↑
Transformer
w. RoPE	340	42.0	31.0	64.4	30.2	51.0	44.3	18.7	39.9
w. Gate (FoX)	376	48.1	30.6	64.9	30.7	51.1	44.7	18.9	40.1
SSM
GLA	400	42.1	30.7	64.4	30.1	52.7	43.8	19.6	40.2
GSA	399	44.1	30.3	64.9	30.7	51.5	45.6	20.5	40.5
GDN	475	40.1	31.6	65.6	31.4	50.2	45.7	19.3	40.6
Mamba-2	382	43.0	29.9	65.0	31.5	51.2	47.5	20.5	40.1
SWA	374	40.7	30.5	64.5	30.4	51.6	44.9	18.6	40.0
Raven	424	41.0	32.7	64.1	30.3	51.7	43.9	18.4	40.2

Hybrid Models Retrieval Ability

Table 4: Hybrid-Raven vs. other hybrid architectures on retrieval tasks. ✓ = no convolutional memory needed.

~400M parameter models

Model	No Conv.	SWDE	FDA	SQuAD	N1-1K	N1-2K	N1-4K	N1-8K	N1-16K	N1-32K	N2-1K	N2-2K	N2-4K	N2-8K	N2-16K	N2-32K	N3-1K	N3-2K	N3-4K	N3-8K	N3-16K	N3-32K
GDN	✗	54.6	67.2	34.5	100	100	100	100	93.2	70.5	100	100	100	8.0	0	0	93.2	70.2	50.0	0	0	0
Mamba-2	✗	56.3	68.8	36.0	100	100	16.4	0	0	0	100	100	85.8	0	0	0	76.9	80.6	60.8	0	0	0
SWA-RoPE	✓	51.0	68.1	34.1	100	100	100	100	98.2	60.4	100	100	100	98.2	3.1	0	93.4	78.2	12.8	60.0	4.4	0
Raven	✓	51.4	64.2	31.4	100	100	100	100	98.4	78.6	100	100	100	100	95.4	65.4	90.0	67.0	73.8	60.0	10.2	14.4

---

Model Structure

Component	Details
Layer type	RavenAttention (replaces standard attention)
Memory	Fixed-size slot matrix per head (num_slots slots)
Router	Linear or MLP projection → top-k sigmoid/softmax
Decay	Mamba2 (A_log + dt_bias) or GLA (logsigmoid)
Feature map	Swish, ReLU, or T2R
Computation	Chunked (training) / Fused recurrent (inference)
Hybrid layers	Optional standard attention layers at specified indices

---

Installation

Install the FLA dependency first, following the official FLA guide:

pip install flash-linear-attention

Then clone this repo:

git clone https://github.com/AvivBick/RoutingMemory
cd RoutingMemory
pip install -e .

Requirements: PyTorch ≥ 2.5, Triton ≥ 3.0, einops, transformers ≥ 4.45.0

---

Usage

As a layer

from raven.layers import RavenAttention

attn = RavenAttention(
    hidden_size=1024,
    num_heads=4,
    num_slots=256,
    topk=32,
    decay_type='Mamba2',    # or 'GLA'
    feature_map='swish',
    router_type='lin',      # or 'mlp'
    router_score='sigmoid', # or 'softmax'
).cuda()

x = torch.randn(1, 2048, 1024).cuda()
y, _, _ = attn(x)  # (batch, seq_len, hidden_size)

As a full causal LM

from raven.models import RavenConfig, RavenForCausalLM
from transformers import AutoModelForCausalLM

config = RavenConfig(
    hidden_size=1024,
    num_hidden_layers=24,
    num_heads=4,
    num_slots=256,
    topk=32,
    decay_type='Mamba2',
    feature_map='swish',
    vocab_size=32000,
)
model = AutoModelForCausalLM.from_config(config).cuda()

RWKV-7 sequence mixers

This fork can also replace the default routing-memory layer with RWKV-7 mixers adapted from HRM-Text. The implementation keeps Raven's embedding, normalization, MLP, LM head, and Hugging Face model API, while swapping RavenAttention for an RWKV-7 mixer inside each non-attention block.

config = RavenConfig(
    hidden_size=1024,
    num_hidden_layers=24,
    sequence_mixer="rwkv7",  # or "routed_rwkv7" / "slot_rwkv7" / "low_rank_slot_rwkv7"
    rwkv7_head_size=64,
    rwkv7_backend="cuda",
    rwkv7_chunk_len=16,
    low_rank_slot_rwkv7_rank=8,
    low_rank_slot_rwkv7_backend="triton_fused",
    vocab_size=32000,
)
model = RavenForCausalLM(config)

Implemented options:

sequence_mixer	Memory/routing granularity	State size per head	Kernel path	Training status
"rwkv7"	Dense RWKV-7 state, no router	D x D	LT2 CUDA when available	Trainable baseline
"routed_rwkv7"	Raven-style top-k router mapped to per-head channel groups	D x D	LT2 CUDA when available	Trainable router adapter
"slot_rwkv7"	Explicit per-head recurrent slot states, closest to Raven memory slots	num_slots x D x D	PyTorch CUDA recurrence	Correct but slow
"low_rank_slot_rwkv7"	Explicit routed slots with low-rank per-slot state	num_slots x rank x D	Triton forward/backward with triton_fused	Ready-to-train comparison path

Use sequence_mixer="slot_rwkv7" when you want RWKV to have Raven-level routed memory slots. It creates num_slots independent RWKV state matrices per head and applies the router inside the recurrent update. This is semantically closer to Raven, but slower until a dedicated slot-aware CUDA kernel is written.

Use sequence_mixer="low_rank_slot_rwkv7" to keep explicit routed slots but reduce each slot state from head_dim x head_dim to rank x head_dim; configure the rank with low_rank_slot_rwkv7_rank. Forward/eval can use the Triton kernel via low_rank_slot_rwkv7_backend="auto" or "triton". For training, use "triton_fused" for a Triton reverse-scan backward, or "triton_autograd" to keep the older Triton-forward/PyTorch-recompute backward path for debugging.

To compare Raven vs. RWKV-7 with the same model shape:

PYTHONPATH=/home/xiaol/X/HRM-Text:$PYTHONPATH \
LT2_RWKV7_CUDA_DIR=/home/xiaol/X/LT2_upstream/apps/LT2/cuda/rwkv7 \
python examples/compare_mixers.py \
  --device cuda \
  --dtype bf16 \
  --rwkv-backend cuda \
  --hidden-size 512 \
  --num-layers 4 \
  --seq-len 512 \
  --batch-size 2 \
  --steps 10

Use --backward to compare training-step cost instead of inference-only forward cost. The LT2 CUDA path currently requires BF16, rwkv7_head_size=64, and rwkv7_chunk_len=16.

For held-out loss comparisons on the already-prepared HRM-Text token subset, use:

PYTHONPATH=/home/xiaol/X/HRM-Text:/home/xiaol/X/raven:$PYTHONPATH \
LT2_RWKV7_CUDA_DIR=/home/xiaol/X/LT2_upstream/apps/LT2/cuda/rwkv7 \
/home/xiaol/X/HRM-Text/.venv/bin/python examples/compare_hrm_text_losses.py \
  --dataset-path /home/xiaol/X/hrm_text_subset_1B \
  --out-dir outputs/hrm_text_loss_compare_raven_vs_low_rank_slot_rwkv7_s500 \
  --steps 500 \
  --eval-every 100 \
  --eval-batches 20 \
  --batch-size 4 \
  --seq-len 128 \
  --hidden-size 256 \
  --num-layers 2 \
  --num-heads 4 \
  --num-slots 64 \
  --topk 16 \
  --low-rank-slot-rwkv-rank 8 \
  --low-rank-slot-rwkv-backend triton_fused \
  --lr 3e-4 \
  --warmup-steps 50 \
  --dtype bf16 \
  --device cuda

Local HRM-Text causal-LM loss result on RTX 4090, using the command above:

mixer	params M	initial val	final val	best val	final train	tok/s	peak GB
raven	35.39	11.1281	2.9195	2.9195	2.5000	49,627	0.448
low_rank_slot_rwkv7	35.56	11.0938	2.6273	2.6273	1.9531	52,854	0.991

Validation checkpoints:

step	raven	low_rank_slot_rwkv7
0	11.1281	11.0938
100	3.8078	3.2961
200	3.2234	2.8469
300	2.9813	2.6836
400	2.9367	2.6406
500	2.9195	2.6273

In this short local run, the train-ready low-rank slot RWKV path reaches lower held-out loss and slightly higher throughput than the Raven mixer, with higher VRAM use from cached low-rank slot states during fused backward. Full results are saved in outputs/hrm_text_loss_compare_raven_vs_low_rank_slot_rwkv7_s500/.

---

Training

Raven uses the flame training framework. Add a config from configs/raven_340M_*.json to flame/configs/, then:

CUDA_VISIBLE_DEVICES=0,1,2,3 NGPU=4 bash train.sh \
    --job.config_file flame/models/fla.toml \
    --job.dump_folder exp/raven-340M \
    --model.config configs/raven_340M_1.json \
    --optimizer.name AdamW \
    --optimizer.lr 3e-4 \
    --lr_scheduler.warmup_steps 1024 \
    --lr_scheduler.decay_type cosine \
    --training.batch_size 16 \
    --training.seq_len 2048 \
    --training.gradient_accumulation_steps 4 \
    --training.steps 30720 \
    --training.dataset /path/to/SlimPajama-627B \
    --training.streaming \
    --training.compile \
    --checkpoint.interval 3072

The configs/ directory contains 12 ablation configurations varying the router design (linear vs. MLP, sigmoid vs. softmax, with/without Gumbel noise and bias).

---

Repository Structure

raven/
├── layers/
│   └── raven.py                # RavenAttention layer
└── models/
    └── raven/
        ├── configuration_raven.py
        └── modeling_raven.py

configs/
└── raven_340M_*.json           # 12 ablation configs (340M scale)

assets/img/                     # figures used in this README

---

Upstream: Flash Linear Attention

This repo builds on [fla-org/flash-linear-attention] and depends on it for hardware-efficient Triton kernels. In particular, Raven currently reuses FLA’s GSA chunked and fused recurrent kernels rather than vendoring separate Raven ops in this repository.

---

Citation

@article{afzalbick2026raven,
  title={Raven: High-Recall Sequence Modeling with Sparse Memory Routing},
  author={Arshia Afzal, Aviv Bick, Eric P. Xing, Volkan Cevher, Albert Gu},
  year={2026},
  publisher={MDPI}
}