ARM: Attention with Routed-Memory for Learnable Sparse Control
Qiuhao Zeng, Jerry Huang, Peng Lu, Ruiyi Fang, Ruiyi Fang, Gezheng Xu, Zihao Jing, Yufei Cui, et al.
cs.LG · cs.AI
Quick verdict
Differentiable routed-memory KV cache with gated writes and adaptive retrieval; efficiency gains look real but baseline confounds weaken attribution.
Summary
The paper introduces ARM, a KV-cache structure for long-context LLM inference that replaces token eviction with a fixed-size, tree-structured routed memory. Incoming tokens descend a 4-level hierarchy (256 leaf buckets) via Gumbel-Softmax-trained routers; within a bucket, a learned slot-selection policy picks a slot whose content is softly updated through a sigmoid gate rather than overwritten, and a re-parameterized gate formulation enables prefix-scan parallelization during prefill. A second learned policy — framed as an MDP with an LSTM stopping rule — adaptively chooses how many buckets to retrieve per input. A theoretical section argues gated mixing is Bayes-optimal versus one-hot eviction under uncertain future access. On Llama-3.1-8B, continued-pretrained on 10B FineWeb-Edu tokens, ARM matches full attention on short-context benchmarks (71.89 vs 72.27 average), edges it on LongBench (18.02 vs 17.09), scores far higher on RULER (15.12 vs 8.01), and is the only method still usable at 128K, with a flat ~0.1 GB cache and near sliding-window latency. Main limitations: baselines did not receive the same continued pretraining, RULER baseline scores look implausibly low, no variance is reported, and recent learned-eviction baselines are absent from the comparison.
Review complete
Assessment at a glance
5
risk flags identified
● high 2
● medium 2
● low 1
4 strengths · 5 weaknesses · 5 author questions
AI text signal = Binoculars zero-shot detector score over 20 sampled passages (Fast-DetectGPT co-reported) — absence does not prove human authorship, and a flag does not imply misconduct.
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- A fixed-size hierarchical routed memory with Gumbel-Softmax slot selection and sigmoid-gated soft updates sustains quality comparable to full KV-cache attention without hard eviction (Tables 1-2).
- The learned MDP stopping policy adaptively sizes retrieval per input, beating fixed sparsity ratios (Table 9: 25.40 vs 24.42/24.94) and FIFO eviction (17.41).
- Gated writes are Bayes-optimal for shared-slot storage under uncertain future access (§4: w* = π), making them strictly more expressive than eviction-based updates.
- ARM holds a constant ~0.1 GB cache with near sliding-window latency and is the only evaluated method producing results at 128K context (Table 4: 5.78; Fig. 2).
Strengths
Coherent, well-positioned design
ARM unifies differentiable neural memory (NTM/ABC lineage) with KV-cache management; Fig. 1 and §2 clearly delimit it from quantization, eviction, and offloading, and the write path (Eqs. 3-6) is fully specified.
Honest parallelization analysis
§3.3 explicitly proves state-dependent gates admit no exact associative scan, then derives a workable re-parameterization (gates depend only on last sync point t0, Eq. 8) recovering O(log Ty) prefix-scan prefill (Eqs. 9-10) — a useful contribution beyond the cache structure.
Clean component ablation
Strong, concrete efficiency results
Weaknesses
Training confound in headline comparisons
Implausible RULER baseline numbers
No uncertainty quantification
Missing modern eviction baselines
Internal inconsistencies
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
Report preview
Scope & Claims
ARM targets KV-cache management during long-context LLM inference. The stated contributions (§1) are: (i) learnable "soft eviction" — Gumbel-Softmax routing into a fixed set of memory slots with sigmoid-gated updates instead of hard replacement; (ii) adaptive Top-M retrieval — a learned stopping policy that chooses how many buckets to attend to per input; (iii) scalable long-context inference. A fixed 4-level, 4-ary tree yields 256 leaf buckets; bucket count and slots per bucket are constant, so cache memory no longer grows with sequence length. §4 argues gated mixing is Bayes-optimal versus one-hot eviction when future access is uncertain.
The remaining sections cover Methodology Assessment, Statistical Evidence, Results & Reproducibility, Limitations & Risks, and Verdict & Recommendations.
- All 4 strengths & 5 weaknesses, analyzed in detail
- 5 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
What is ARM (Attention with Routed-Memory)?
ARM is a KV-cache structure for LLM inference presented in this 2026 paper. Instead of letting the cache grow unboundedly or evicting tokens, it keeps a fixed-size tree of 256 memory buckets. New tokens are routed to a bucket and merged into an existing slot through a learned, sigmoid-gated soft update, and a learned policy decides per input how many buckets to read from during attention.
Does ARM beat full attention?
Not cleanly. On short-context benchmarks ARM is slightly below full attention (71.89 vs 72.27 average in Table 1). On LongBench it scores a bit higher (18.02 vs 17.09) and on RULER much higher (15.12 vs 8.01), but the ARM model received 10 billion tokens of additional pretraining that the baselines did not, so those gains cannot yet be attributed to the cache mechanism itself. Its clear advantage is efficiency: a flat ~0.1 GB cache and the only working configuration at 128K context.
How much memory does ARM's KV cache use?
Per Figure 2, ARM's cache stays at roughly 0.1 GB regardless of generation length (up to ~8K output tokens tested), whereas a 4-bit quantized cache (KIVI) grows from 0.1 GB to 8.1 GB and offloading/SWA configurations grow or transfer comparably large volumes. ARM's decoding latency (~16-17 ms/token) is on par with sliding-window and sink caches and about 10x faster than CPU offloading.
Is code for ARM available?
The paper does not mention a code release. It does provide a detailed training recipe (10B FineWeb-Edu tokens, Llama-3.1-8B, AdamW with lr 4e-4, 2048-token context, 8×A100-40GB training, H100-80GB inference, lm-eval/RULER harness versions in Appendix B), but reproducing the results would require substantial multi-GPU compute.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- Figures 1 and 2 were only partially recoverable from text extraction; the latency/memory values cited were read from embedded plot labels and may omit plot detail.
- The supplied metadata omits Gang Niu and Boyu Wang, who appear in the paper byline (corresponding author: Boyu Wang); the review reflects the byline in the text.
- The ablation (Table 9) covers only short-context recall tasks; no long-context ablation exists in the provided text.
- Absolute RULER numbers could not be cross-checked against published results within this offline review; the implausibility judgment rests on general knowledge of Llama-3.1-8B RULER performance.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.