MKA: Memory-Keyed Attention for Efficient Long-Context Reasoning

2026-03-21 · Dong Liu, Yanxuan Yu, Ben Lengerich, Ying Nian Wu arxiv

As long-context language modeling becomes increasingly important, the cost of maintaining and attending to large Key/Value (KV) caches grows rapidly, becoming a major bottleneck in both training and inference. While prior works such as Multi-Query Attention (MQA) and Multi-Latent Attention (MLA) reduce memory by sharing or compressing KV features, they often trade off representation quality or incur runtime overhead. We propose Memory-Keyed Attention (MKA), a hierarchical attention mechanism that integrates multi-level KV caches (local, session, and long-term) and learns to route attention across them dynamically. We further introduce Route-Fused MKA (FastMKA), a broadcast-routed variant that fuses memory sources before attention computation for improved efficiency. Experiments on different sequence lengths show that FastMKA achieves a favorable accuracy-efficiency trade-off: comparable perplexity to MLA while achieving up to 5x faster training throughput and 1.8x lower evaluation latency. These results highlight MKA as a practical and extensible framework for efficient long-context attention.

📄 PDF Abstract BibTeX arXiv:2603.20586

Code (0)

등록된 구현이 없습니다.

MSA: Memory Sparse Attention for Efficient End-to-End Memory Model Scaling to 100M Tokens

2026-03-06 · Yu Chen, Runkai Chen, Sheng Yi, Xinda Zhao 외 arxiv

Long-term memory is a cornerstone of human intelligence. Enabling AI to process lifetime-scale information remains a long-standing pursuit in the field. Due to the constraints of full-attention architectures, the effecti…

CoMIC: Collaborative Memory and Insights Circulation for Long-Horizon LLM Agents in Cloud-Edge Systems

2026-05-30 · Yannan Wang, Longli Yang, Zhen Liu, Abhishek Kumar 외 arxiv

Deploying lightweight Large Language Model (LLM) agents on edge servers can reduce latency and move agentic services closer to users, but resource-constrained edge models often struggle with long-horizon tasks that requi…

CogMem: A Cognitive Memory Architecture for Sustained Multi-Turn Reasoning in Large Language Models

2025-12-16 · Yiran Zhang, Jincheng Hu, Mark Dras, Usman Naseem arxiv

Large language models (LLMs) excel at single-turn reasoning but often lose accuracy and coherence over extended, multi-turn interactions. Recent evaluations such as TurnBench highlight recurring failure modes-reasoning b…

ReContext: Recursive Evidence Replay as LLM Harness for Long-Context Reasoning

2026-07-02 · Yanjun Zhao, Ruizhong Qiu, Tianxin Wei, Yuanchen Bei 외 arxiv

Understanding and reasoning over long contexts has become a key requirement for deploying large language models (LLMs) in realistic applications. Although recent LLMs support increasingly long context windows, they often…

Answer Generation

Event-Keyed Summarization

2024-02-10 · William Gantt, Alexander Martin, Pavlo Kuchmiichuk, Aaron Steven White

We introduce event-keyed summarization (EKS), a novel task that marries traditional summarization and document-level event extraction, with the goal of generating a contextualized summary for a specific event, given a do…

Document-level Event Extraction Event Extraction