Memory-Efficient Attention

Follow this topic Ask about this topic

Latest papers

24 papers1 paper in the last four weeks, against 2 the four weeks before. 0.0% of all new papers.

Sep 28, 2026

SANTA++: Sampling Attention through Representative Keys

Attention often concentrates on a small subset of tokens in the context, but which subset matters changes from one query to the next. To exploit this changing structure, we introduce SANTA++, a training-free stochastic attention method that uses representative keys for memory-efficient selection without scanning the entire key-value (KV) cache. Cached keys are organized into teams, and the query scores one representative from each team to decide which teams to sample. We compute exact attention scores within the sampled teams and reweight each team's contribution by the inverse of its inclusion probability. This importance sampling correction estimates attention over the full cache, with a sampling budget that lets us trade memory reads for accuracy. Remarkably, with 32 or 64 sampled teams, SANTA++ uses 16% to 22% of dense attention's KV reads and retains 94% to 99% of the dense-attention baseline's scores on LongBench v2 and HELMET's retrieval-augmented generation subset, and 85% to 91% on RULER, with Qwen2.5-7B-Instruct at 32K context. With 31 sampled teams, our GPU implementation delivers a

1.69\times

attention speedup over the dense FlashAttention baseline at 32K context. By reducing the number of cache entries read, SANTA++ in principle complements architectures with compressed KV representations, such as multi-head latent attention. Our kernels are available at: https://github.com/OPUSLab/santapp-kernel-demo.git.

Kyle Lee, Christian Z. Pratt, Ruoyu Fang +4

Department of Electrical and Computer Engineering University of California, Santa Barbara, CA, USA · Flucta, San Francisco, CA, USA

Memory-Efficient Inference LLM Inference Acceleration Efficient Attention Memory-Efficient Attention

Aug 28, 2026

Sliding-window beats linear attention

Due to the nature of quadratic attention, Large Language Models (LLMs) consume a lot of memory and energy: every new token costs more than the previous one, and its keys and values must be stored in memory indefinitely, which is unsustainable. Two main lines of work address this: compressing the KV cache, e.g., by evicting or quantizing keys and values, and retrofitting LLMs to use Linear Attention, which replaces the KV cache with a fixed-size state. Retrofitting has attracted a lot of attention, given its promise to solve the quadratic scaling problem with state-of-the-art performance at low cost. However, it has not been properly compared to the simplest form of KV-cache eviction: Sliding Window Attention (SWA) with attention sinks. In this work, we show that SWA with sinks performs as well or better than most retrofitted Linear Attention models across multiple LLMs and downstream tasks, with the largest gains on long-context and generative tasks. On long-context reasoning tasks (Needle-in-a-Haystack and BABILong), SWA achieves massively higher performance (2 to 10 times higher than linear attention). SWA requires no additional training, is extremely fast, and requires little memory, making it an extremely cheap and reliable solution. When the training budget is limited, switching to SWA is a much more effective way to reduce inference memory cost than retrofitting linear attention. Linear attention models have shown promise, but they require training from scratch or extensive retrofitting to reap their architectural benefits and come close to SWA.

Alexia Jolicoeur-Martineau, Rhea Sanjay Sukthanker, Pashmina Cameron +1

Microsoft Applied Sciences Group (ASG) · Independent

Linear Attention KV Caching LLM Inference Acceleration Long-Context Inference Attention Sinks Sliding-Window Attention+1

Aug 18, 2026

Q-Interference: Memory-Efficient Phase-Aware Quantum-Inspired Attention

GPT attention measures token compatibility through dot-product similarity. This mechanism is simple, effective, and memory-efficient. But it does not explicitly model whether strong token features should reinforce or suppress one another. We introduce Q-Interference, a fully classical quantum-inspired attention mechanism for autoregressive language modeling that augments each query and key feature with an amplitude and a learned phase. The resulting attention score is phase-aware which aligned phases contribute constructively while conflicting phases contribute destructively. Although Q-Interference yields a richer interaction rule than similarity alone, a naive implementation of Q-Interference requires a large token-pair-feature interaction tensor, making it memory-intensive and often impractical. To address this limitation, we propose an exact trigonometric factorization that computes the same score using two standard matrix multiplications avoiding materialization of the large intermediate tensor. Q-Interference fits directly into a Transformer block in GPT and leaves the remainder of the model architecture and next-token prediction objective unchanged. Experiments on public benchmark datasets and baseline models show that the proposed reformulation trains stably in a controlled GPT-style setting and provides a consistent memory advantage over naive phase-aware interference attention. These results support the specific contribution of this work: an exact memory-efficient reformulation that makes phase-aware interference attention practical within a standard GPT pipeline.

Emama Nahid, Tahmid Imtiaz Imu, Huayue Gu +3

Kennesaw State University Marietta, GA, USA · Miami University Oxford, OH, USA · Georgia State University Atlanta, GA, USA

Memory-Efficient Attention Transformer Attention Efficient Attention Autoregressive Language Modeling

Aug 11, 2026

Efficient Reinforcement Learning for Long-Horizon Tool-Use Agentic Tasks

Long-horizon tool-using agents must reason over user goals, domain policies, tool calls, simulator state, and delayed verifiable rewards. Reinforcement learning (RL) is a natural fit for this setting, but multi-turn on-policy rollouts create long contexts, while model-specific attention layers may require custom masks and learned sink normalization. We present SINKFLEX-RL, a modular training system for RL in dual-control tool-use environments. The system combines a Gymnasium-compatible environment wrapper, a VERL-style rollout dataflow, group-relative policy optimization without a separate value model, and a sink-aware FlexAttention path designed to preserve model-specific sink scaling under causal and sliding-window masks. In a preliminary Tau2Bench retail run, validation reward (mean@1) rises from 0.25 early in training to

0.44

later in the observed training window, while training-score and trajectory-reward proxies also trend upward. In a fixed-configuration memory benchmark, the optimized attention path reduces peak VRAM from 28.06GB to 22.52GB at 4096 tokens, a

19.7\%

reduction, and runs the measured 8192-token configuration using

25.53

~GB where the eager baseline runs out of memory. These results illustrate the value of integrating environment interfaces, RL dataflow, and attention-kernel design for memory-feasible long-horizon agent training.

Zelei Cheng, Amritansh Mishra, Sambit Sahu +1

AI Foundations, Capital One

Memory-Efficient Attention Efficient Attention Long-Horizon Agent Tasks RL for LLM Agents RL for Tool Use

Jul 21, 2026

MoA-Structured Decode Attention DNF Derivation, KV-Cache Accumulation, GQA/MQA, and OpenACC Kernel

We derive four memory-optimal inference artifacts for transformer attention using the Mathematics of Arrays (MoA), each following directly from the forward-pass Denotational Normal Form (DNF) of with the query-row index fixed to the current decode step. The artifacts are: (1)~a single-query decode DNF in which the

ψ

-reduction eliminates the

K^\top

buffer algebraically, achieving

(d_k + nd_k+ nd_v+ d_v)\times4\,{B}

Dynamic Random Access Memory (DRAM) traffic result numerically verified to

\|{err}\|_\leq2\times10^{-7}

; (2)~a C/OpenACC Graphics Processing Unit (GPU) kernel with Operational Normal Form (ONF) stride arithmetic and hardware-coalesced memory access, verified to

\|\mathrm{err}\|_\infty=0

(exact IEEE-754 floating-point arithmetic); (3)~a multi-step KV-cache with

O(d_k+d_v)

per-step append via MoA concatenation

\#

; and (4)~Grouped-Query Attention (GQA) and Multi-Query Attention (MQA) derived via

ψ

-selection, achieving a proven

\frac {h_q} { h_{kv} }

reduction in KV traffic. All programs are verified against PyTorch scaled_dot_product_attention.

Lenore Mulin, Gaetan Hains

University at Albany, SUNY · LACL, Universit´e Paris-Est Cr´eteil

Attention Mechanisms Efficient Transformer Inference Grouped-Query Attention Transformer Inference GPU Kernel Optimization KV Caching+3

Jul 2, 2026

A Hippocampus for Linear Attention: An Exact Memory for What the Recurrent State Forgets

Linear-attention and state-space language models compress the prefix into a fixed-size recurrent state, yielding O(1) memory at the cost of a lossy exact memory: when many key--value associations compete, earlier facts are overwritten and needle recall degrades. Inspired by Complementary Learning Systems, we give linear attention a hippocampal complement. HOLA (Hippocampal Linear Attention) keeps the usual delta-rule state as a compressive memory and adds a bounded exact KV cache, forming a semiparametric test-time memory: the state models linearly compressible structure, while the cache stores associations that should not be forced through that state. The cache writes without a learned eviction module, keeping tokens with large beta * ||e||, the prediction residual actually committed to the state; a decoupled RMSNorm-gamma cache read then turns these exact KV pairs into sharp retrieval rather than soft averaging. At 340M parameters trained on 15B SlimPajama tokens, HOLA lowers Wikitext perplexity from 27.32 to 22.92 (-16.1%), below a full-attention Transformer++ (26.88), and improves LAMBADA perplexity from 30.95 to 30.26. It also achieves the best linear in-context retrieval and remains much more robust than GDN or a matched HOLA+recency cache on RULER needle-in-a-haystack recall out to 32k tokens (16x its training length).

Wanyun Cui

Shanghai University of Finance and Economics

Hybrid Attention Long-Context Retrieval Memory-Augmented Language Models KV Caching Long-Context Language Modeling Persistent Memory for Language Models+2

Jun 23, 2026

Scalable Peptide Design via Memory-Efficient Equivariant Transformer

Target-specific peptide design requires sequence and structure co-design under full atom geometric constraints. Latent generative frameworks offer an effective route for this problem by compressing fine grained atomic structures into block level latent representations and performing conditional generation in a compact latent space. However, the scalability of such systems depends heavily on the geometric backbone used throughout their encoding, decoding, and denoising components. We introduce MEET (Memory Efficient Equivariant Transformer), an E(3) equivariant backbone for scalable atomistic peptide modeling. MEET maintains coupled invariant scalar and equivariant vector feature streams, while reformulating geometric computation around memory efficient attention. It initializes vector features through global coordinate aggregation, incorporates pairwise distances through augmented query and key dot products, and injects covalent bond information through sparse bond adaptation. Integrated into a VAE and latent diffusion pipeline for full atom peptide generation, MEET achieves linear memory scaling with atom count and improves generation quality over existing peptide design methods. Experiments on large scale AFDB derived datasets further show that the proposed backbone supports systematic model and data scaling, leading to better binding affinity, physical validity, and sample diversity.

Rui Jiao, Xiangzhe Kong, Yinjun Jia +4

Department of Computer Science and Technology, Tsinghua University · Institute for AI Industry Research, Tsinghua University · Department of Electronic Engineering, Tsinghua University +1

Memory-Efficient Attention Equivariant Neural Networks Efficient Attention Protein Design

Jun 22, 2026

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

Long contexts have become standard in pretrained LLMs, yet they remain expensive to run: prefill compute grows quadratically with sequence length, and every decode step re-reads a key-value cache that grows linearly with it. Sparse attention cuts these costs by attending only to a relevant subset of past tokens, but selecting that subset is itself expensive. We present SpotAttention, a lightweight selector that attaches to a frozen pretrained transformer and learns by KL distillation to estimate its attention distribution. The selector picks the top-K keys each query attends to, and because its estimate is a calibrated distribution, a dual top-p rule reads the per-query, per-layer budget directly from it. Across Qwen3 (dense, 4B-32B) and Qwen3.5 (hybrid linear/full attention, 4B-9B), SpotAttention matches dense accuracy at contexts up to 128K tokens, eight times the training length. Decode at L=128K runs 3.9x faster than FlashAttention and 1.8x faster than Twilight, the strongest training-free baseline. Quantizing the selector's K-cache to INT4 or FP4 microscale shrinks it 3.5x at no accuracy cost.

Huzama Ahmad, Se-Young Yun

KAIST

Sparse Attention Block-Sparse Attention Attention-Guided Distillation Long-Context Language Model Inference LLM Inference Acceleration Memory-Efficient Attention+2

Jun 20, 2026

Keyless Attention: Value-Space Routing and Value-Only Caching for Efficient Transformers

We propose Keyless Attention, an attention mechanism that eliminates the key projection entirely, operating over queries and values only. This yields a Value-Only Cache that reduces KV cache memory and access overhead by exactly 50% over standard attention, while matching or exceeding standard attention's decode throughput. Beyond efficiency, we introduce Depth-

m

Attention Factorization: standard attention computes a depth-2 factorization of the attention bilinear form, while Keyless Attention realizes a depth-

m

instance of this family. At m=3, Keyless Attention matches the projection matrix count of standard attention via a value-space routing matrix that replaces the key projection and introduces a coupling between routing and retrieval. Experiments across five models and four architectures (GPT-2 280M, GPT-2 557M, Pythia 410M, Qwen2 1.5B, and Llama 3.2 1B) show that Keyless Attention matches or outperforms standard QKV attention on perplexity in 4 out of 5 models. On downstream zero-shot evaluation (GPT-2 557M), Keyless Attention outperforms on 4 out of 5 commonsense reasoning benchmarks, while achieving 50% KV cache reduction throughout.

Xin Gao

Department of Mathematics and Statistics York University Toronto, M3J 1P3

Transformer Inference KV Caching Memory-Efficient Optimization Memory-Efficient Inference Efficient Attention Memory-Efficient Attention

Jun 9, 2026

Express Language Modeling

We introduce a new tool, Express, for converting a non-causal attention approximation into a causal approximation with matching approximation guarantees. When combined with the state-of-the-art Thinformer approximation, Express improves upon the best known causal attention guarantees, delivering

\log^{3/2}(n)/s

approximation error with only

O(s)

memory and

O(s^2 \log^2(n))

compression overhead for a sequence of length

n

. We pair these developments with an efficient I/O-aware Triton implementation, demonstrate substantial speedups over FlashAttention 2, and use Express to overcome four resource bottlenecks in the language modeling pipeline: long-context prefill, KV cache compression, long-form memory-constrained decoding, and long-form compute-constrained decoding.

Albert Gong, Annabelle Michael Carrell, Raaz Dwivedi +1

Cornell Tech · University of Cambridge · Microsoft Research

Long-Context Language Model Inference Attention Mechanisms Efficient Language Model Inference Attention Approximation Memory-Efficient Attention Causal Attention

Jun 5, 2026

Attention at the Theoretical Minimum: A Mathematics of Arrays Framework for Memory-Optimal Transformer Kernels

The attention mechanism is the dominant computational bottleneck in modern transformer-based AI. Its standard implementation incurs quadratic memory traffic in the sequence length~

n

, and DRAM accesses cost 100--1000

\times

more energy than arithmetic operations on contemporary hardware, so any analysis focused solely on FLOP counts fundamentally mischaracterises the bottleneck. We present a Mathematics of Arrays (MoA) reformulation of scaled dot-product attention and its numerically stable softmax, deriving a Denotational Normal Form (DNF) that eliminates all intermediate arrays -- including the implicit transposed-key buffer and every softmax temporary -- by algebraic construction rather than empirical tuning. The DNF achieves

O(n_{dk} + n{_{dv}})

data movement versus

O(n^2 + n_{dk} + n_{dv})

for the standard implementation, where

n

is the sequence length,

dk

is the key dimensionality and

dv

the value dimensionality, and is verified numerically against PyTorch at full double-precision floating-point on concrete inputs. Unlike hardware-specific accelerators or empirical tiling schemes such as FlashAttention, MoA simultaneously provides array fusion, shape-transformation correctness, and predictive cost models from a single algebraic framework. Memory minimality is a theorem established before any code is written. A predictive performance model projects

2

--

100\times

speedup and

2

--

50\times

energy reduction, with the advantage widening at exascale. The derivation establishes a formally verified pipeline from Python specification through (ONF) Operational Normal Form, and dimension-lifted hardware mapping, providing performance-portable AI kernels of direct relevance to DARPA edge-deployment and DOE exascale priorities.

Lenore Mullin, Gaetan Hains

University at Albany, SUNY, Albany, NY, USA · LACL, Université Paris-Est Créteil, Créteil, France

Memory-Efficient Attention GPU Kernel Optimization Energy-Efficient ML High-Performance Computing Memory-Efficient Optimization Transformer Attention+1

May 27, 2026

Augmenting Attention with Exponentially Decaying Memory Improves Query-Aware KV Sparsity

Efficient inference is critical for long-context language models, where attention computation and KV-cache access dominate the cost. Recent work RAT+, introduces a recurrence-augmented attention backbone that enables flexible dilated attention at inference time. In this paper, we investigate whether this exponentially decaying memory can also improve existing query-aware sparse inference methods. Using representative methods including Quest, MoBA, and SnapKV, we show that RAT+ consistently improves accuracy over standard attention across sparse budgets on eight needle-in-a-haystack tasks. We validate these gains both on the released checkpoints from the RAT+ paper and on OLMo2-7B, which we continue pretraining with the added memory module for 10B tokens. Finally, we propose two hypotheses explaining why this memory module benefits query-aware sparse inference and design targeted experiments to support them.

Xiuying Wei, Caglar Gulcehre

EPFL, CLAIRE Lausanne, Switzerland

KV Caching Memory-Augmented Language Models Self-Attention Sparse Attention Long-Context Language Model Inference LLM Inference Acceleration+1

May 22, 2026

Approaching I/O-optimality for Approximate Attention

We revisit the I/O complexity of attention in large language models. Given query-key-value matrices

Q,K,V\in\mathbb{R}^{n\times d}

, and a machine with fast memory size

M

, the goal is to compute the "attention matrix"

A=\text{softmax}(Q K ^{\top}/\sqrt{d}) V

with the minimal number of data transfers between fast and slow memory. Existing methods in the literature, most notably FlashAttention and its variants, incur an I/O cost that depends quadratically on

n

, while a trivial lower bound only requires

Ω(nd)

I/O's to read the inputs and write the output. In this work, we present a technique for computing attention where the I/O cost only depends almost-linearly on

n

in most parameter regimes. This is achieved by developing I/O-efficient algorithms inspired by the recent approximate attention framework of Alman and Song. We also prove corresponding lower bounds in each parameter regime to show that our algorithms are indeed close to I/O-optimal.

Pál András Papp, Aleksandros Sobczyk, Anastasios Zouzias

Computing Systems Lab, Huawei Technologies, Zurich, Switzerland

Self-Attention Efficient Attention Memory-Efficient Attention

May 11, 2026

Variational Linear Attention: Stable Associative Memory for Long-Context Transformers

Linear attention reduces the quadratic cost of softmax attention to

\mathcal{O}(T)

, but its memory state grows as

\mathcal{O}(T)

in Frobenius norm, causing progressive interference between stored associations. We introduce Variational Linear Attention (VLA), which reframes the memory update as an online regularised least-squares problem with an adaptive penalty matrix maintained via the Sherman-Morrison rank-1 formula. We prove that normalising the write direction to unit length gives the recurrence Jacobian spectral norm exactly

1

for all sequence lengths and head dimensions (Proposition 2), and that the state norm is self-limiting under bounded inputs (Proposition 1). Empirically, VLA reduces

\|S_t\|_F

by

109\times

relative to standard linear attention at

T{=}1{,}000

, achieves near-perfect exact-match accuracy on multi-query associative recall within the effective per-head memory regime (

n_\text{pairs} < d_h

), maintaining substantially higher retrieval performance than DeltaNet and standard linear attention under increasing memory load, and maintains 62% accuracy at the per-head capacity boundary. A Triton-fused kernel achieves

14\times

speedup over sequential Python and

\mathcal{O}(T)

scaling, crossing below softmax attention latency at approximately 43,000 tokens.

Vishal Pandey, Gopal Singh

Independent Researcher London, UK · Metriqual Athens, GR

Efficient Attention Linear Attention Associative Memory Memory-Efficient Attention Attention Approximation

May 10, 2026

Nectar: Neural Estimation of Cached-Token Attention via Regression

Evaluating softmax attention over a fixed long context requires reading every cached key-value pair for each new query token. For a given context (a book, a manual, a legal corpus) the attention output is a deterministic function of the query. We propose Nectar, which fits a compact neural network to this function for queries drawn from a task-relevant distribution. Nectar fits two networks per layer and KV-head: a target network that predicts the attention output and a score network that predicts the log-normalizer. The pair plugs into the standard masked self-attention at inference time, replacing the

O(n)

attention over the cache with a forward pass whose cost does not depend on

n

. Each module carries on the order of

|θ|

parameters per layer and KV-head, typically much smaller than the

2nd

KV-cache footprint at the same granularity. We report experiments on models from 1.7B to 8B parameters across five long-context datasets. The approximation error tracks the next-token accuracy gap to full attention, and allocating capacity non-uniformly across layers reduces that gap in our ablation. Beyond this analysis of metrics, we check that the text generations (following a question prompt) of a model equipped with a Nectar module match in semantic content those obtained by giving the same model access to the full cache.

João Monteiro, Michal Klein, Pierre Ablin +1

Apple

Attention Approximation Softmax Attention Self-Attention KV Caching Long-Context Language Model Inference Efficient Attention+2

May 10, 2026

Attention Itself Could Retrieve.RetrieveVGGT: Training-Free Long Context Streaming 3D Reconstruction via Query-Key Similarity Retrieval

Visual Geometry Grounded Transformer (VGGT) advances 3D reconstruction via scalable Transformer architecture, but the quadratic complexity of global attention prevents long context application. StreamVGGT enables streaming with causal attention, yet its KV cache grows linearly with frames, causing memory overflow and quality degradation. We present RetrieveVGGT, a training-free framework, which formulates context construction for VGGT as a retrieval problem. By retrieving a fixed number of relevant frames at each step, VGGT maintains a controllable memory budget, which is close to its training context length. Interestingly, we find that the similarity between current frame queries and cached history frame keys at the first global attention layer of VGGT is already a strong indicator of relevance, eliminating the need for additional learned scoring. To enhance information diversity similar to a recommender system, we propose Segment Sampling so that the retrieval spans distinct relevant segments rather than a single high-similarity region. We design a pose-aware spatial memory mechanism that organizes history frames according to their already estimated camera poses, enabling location-aware retrieval. Extensive experiments demonstrate that RetrieveVGGT achieves state-of-the-art performance, outperforming StreamVGGT, TTT3R, and InfiniteVGGT while maintaining constant memory usage regardless of sequence length. Code is available at https://github.com/zzctmd/RetrieveVGGT.

Zichen Zou, Xiaosong Jia, Zuxuan Wu +1

Institute of Trustworthy Embodied AI (TEAI), Fudan University · 2Shanghai Key Laboratory of Multimodal Embodied AI

Memory-Efficient Attention 3D ViTs Streaming 3D Reconstruction 3D Reconstruction Multi-View 3D Reconstruction Efficient Attention

May 7, 2026

Nearly Optimal Attention Coresets

We consider the problem of estimating the Attention mechanism in small space, and prove the existence of coresets for it of nearly optimal size. Specifically, we show that for any set of unit-norm keys and values

(K,V)

in

\mathbb{R}^d

, there exists a subset

(K',V')

of size at most

O({\sqrt{d} e^{ρ+o(ρ)}/\varepsilon})

such that

\left\| \operatorname{Attn}(q,K,V)- \operatorname{Attn}(q,K',V') \right\| \le \varepsilon

simultaneously for all queries whose norm is bounded by

ρ

. This outperforms the best known results for this problem. We also offer an improved lower bound showing that

\varepsilon

-coresets must have size

Ω({\sqrt{d} e^ρ/ε})

.

Edo Liberty, Alexandr Andoni, Eldar Kleiner

Softmax Attention Efficient Attention Coreset Selection Attention Approximation Memory-Efficient Attention

Mar 31, 2026

Tucker Attention: A generalization of approximate attention mechanisms

The pursuit of reducing the memory footprint of the self-attention mechanism in multi-headed self attention (MHA) spawned a rich portfolio of methods, e.g., group-query attention (GQA) and multi-head latent attention (MLA). The methods leverage specialized low-rank factorizations across embedding dimensions or attention heads. From the point of view of classical low-rank approximation, these methods are unconventional and raise questions of which objects they really approximate and how to interpret the low-rank behavior of the resulting representations. To answer these questions, this work proposes a generalized view on the weight objects in the self-attention layer and a factorization strategy, which allows us to construct a parameter efficient scheme, called Tucker Attention. Tucker Attention requires an order of magnitude fewer parameters for comparable validation metrics, compared to GQA and MLA, as evaluated in LLM and ViT test cases. Additionally, Tucker Attention~encompasses GQA, MLA, MHA as special cases and is fully compatible with flash-attention and rotary position embeddings (RoPE). This generalization strategy yields insights of the actual ranks achieved by MHA, GQA, and MLA, and further enables simplifications for MLA.

Timon Klein, Jonas Kusch, Sebastian Sager +2

Department of Mathematics; Otto von Guericke University Magdeburg; 39106 Magdeburg; Germany. · Scientific Computing, Norwegian University of Life Sciences, ˚As, Norway · Max Planck Institute for Dynamics of Complex Technical Systems; 39106 Magdeburg; Germany. +1

Tucker Decomposition Attention Mechanisms Attention Approximation Memory-Efficient Attention Multi-Head Attention Low-Rank Attention+3

Mar 21, 2026

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

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.

Dong Liu, Yanxuan Yu, Ben Lengerich +1

University of California, Los Angeles, Los Angeles, California, USA · Columbia University, New York, New York, USA · University of Wisconsin-Madison, Madison, Wisconsin, USA

Attention Mechanisms Memory-Efficient Attention Long-Context Language Modeling KV-Cache Management LLM Inference Acceleration Memory-Augmented Language Models

Mar 9, 2026

LinearARD: Linear-Memory Attention Distillation for RoPE Restoration

The extension of context windows in Large Language Models is typically facilitated by scaling positional encodings followed by lightweight Continual Pre-Training (CPT). While effective for processing long sequences, this paradigm often disrupts original model capabilities, leading to performance degradation on standard short-text benchmarks. We propose LinearARD, a self-distillation method that restores Rotary Position Embeddings (RoPE)-scaled students through attention-structure consistency with a frozen native-RoPE teacher. Rather than matching opaque hidden states, LinearARD aligns the row-wise distributions of dense

Q/Q

,

K/K

, and

V/V

self-relation matrices to directly supervise attention dynamics. To overcome the quadratic memory bottleneck of

n \times n

relation maps, we introduce a linear-memory kernel. This kernel leverages per-token log-sum-exp statistics and fuses logit recomputation into the backward pass to compute exact Kullback-Leibler divergence and gradients. On LLaMA2-7B extended from 4K to 32K, LinearARD recovers 98.3% of the short-text performance of state-of-the-art baselines while surpassing them on long-context benchmarks. Notably, our method achieves these results using only 4.25M training tokens compared to the 256M tokens required by LongReD and CPT. Our code is available at https://github.com/gracefulning/LinearARD.

Ning Yang, Hengyu Zhong, Wentao Wang +3

Institute of Automation, Chinese Academy of Sciences, Beijing, China · Southwest University, Chongqing, China · Dalian University of Technology, Dalian, Liaoning, China +3

Attention-Guided Distillation Memory-Efficient Attention Long-Context Language Modeling Language Model Distillation Knowledge Distillation Efficient Attention+1

Mar 8, 2026

SketchGraphNet: A Memory-Efficient Hybrid Graph Transformer for Large-Scale Sketch Corpora Recognition

This work investigates large-scale sketch recognition from a graph-native perspective, where free-hand sketches are directly modeled as structured graphs rather than raster images or stroke sequences. We propose SketchGraphNet, a hybrid graph neural architecture that integrates local message passing with a memory-efficient global attention mechanism, without relying on auxiliary positional or structural encodings. To support systematic evaluation, we construct SketchGraph, a large-scale benchmark comprising 3.44 million graph-structured sketches across 344 categories, with two variants (A and R) to reflect different noise conditions. Each sketch is represented as a spatiotemporal graph with normalized stroke-order attributes. On SketchGraph-A and SketchGraph-R, SketchGraphNet achieves Top-1 accuracies of 83.62% and 87.61%, respectively, under a unified training configuration. MemEffAttn further reduces peak GPU memory by over 40% and training time by more than 30% compared with Performer-based global attention, while maintaining comparable accuracy.

Shilong Chen, Mingyuan Li, Zhaoyang Wang +2

College of Computer, Qinghai Normal University, Xining, 810001, Qinghai, China · College of Computer, Nanjing University, Nanjing, 210023, Jiangsu, China · College of Computer, Qinghai Minzu University, Xining, 810007, Qinghai, China

Memory-Efficient Attention Graph Transformer Spatiotemporal GNNs Efficient Attention Graph Classification

Feb 24, 2026

Untied Ulysses: Memory-Efficient Context Parallelism via Headwise Chunking

Efficiently processing long sequences with Transformer models usually requires splitting the computations across accelerators via context parallelism. The dominant approaches in this family of methods, such as Ring Attention or DeepSpeed Ulysses, enable scaling over the context dimension but do not focus on memory efficiency, which limits the sequence lengths they can support. More advanced techniques, such as Fully Pipelined Distributed Transformer or activation offloading, can further extend the possible context length at the cost of training throughput. In this paper, we present UPipe, a simple yet effective context parallelism technique that performs fine-grained chunking at the attention head level. This technique significantly reduces the activation memory usage of self-attention, breaking the activation memory barrier and unlocking much longer context lengths. Our approach reduces intermediate tensor memory usage in the attention layer by as much as 87.5

\%

for 32B Transformers, while matching previous context parallelism techniques in terms of training speed. UPipe can support the context length of 5M tokens when training Llama3-8B on a single 8

\times

H100 node, improving upon prior methods by over 25

\%

.

Ravi Ghadia, Maksim Abraham, Sergei Vorobyov +1

Together AI

Memory-Efficient Attention Context Parallelism LLM Training Efficient Attention Self-Attention

Feb 19, 2026

2Mamba2Furious: Linear in Complexity, Competitive in Accuracy

Linear attention transformers have become a strong alternative to softmax attention due to their efficiency. However, linear attention tends to be less expressive and results in reduced accuracy compared to softmax attention. To bridge the accuracy gap between softmax attention and linear attention, we manipulate Mamba-2, a very strong linear attention variant. We first simplify Mamba-2 down to its most fundamental and important components, evaluating which specific choices make it most accurate. From this simplified Mamba variant (Mamba-2S), we improve the A-mask and increase the order of the hidden state, resulting in a method, which we call 2Mamba, that is nearly as accurate as softmax attention, yet much more memory efficient for long context lengths. We also investigate elements to Mamba-2 that help surpass softmax attention accuracy. Code is provided for all our experiments.

Gabriel Mongaras, Eric C. Larson

Lyle School of Engineering Southern Methodist University Dallas, TX 75205

Linear Attention Memory-Efficient Attention State Space Models

Feb 16, 2026

Thin Keys, Full Values: Reducing KV Cache via Low-Dimensional Attention Selection

Standard Transformer attention uses identical dimensionality for queries, keys, and values, yet these components serve different roles: queries and keys produce scalar attention weights (selection), while values carry rich representations (value transfer). We show that selection requires only

O(\log N)

dimensions to distinguish among

N

relevant token categories (e.g., syntactic roles, semantic clusters, positional patterns) -- far fewer than value transfer needs. We introduce factored keys, which exploit this asymmetry to physically shrink the KV cache of any pretrained model without retraining from scratch -- unlike Grouped-Query Attention (GQA) and Multi-Head Latent Attention (MLA), which must be designed into the architecture before pretraining. We factorize each key projection

W_K \approx A_{d \times r} B_{r \times d}

via truncated singular value decomposition (SVD) (where

r

is the chosen compression dimension), set

W_K' = A

as the new key projection producing compact

r

-dimensional keys for the cache, and absorb

B^\top

into the query projection (

W_Q' = W_Q B^\top

) at zero cost -- since queries are never cached. At the 7B scale, training from scratch with

r = d/4

(where

d

is the model dimension) matches full-attention perplexity (

9.24

vs

9.25

PPL after 20B tokens, mean over two seeds) while using 12% fewer parameters and training 8% faster. For existing models, SVD followed by QK fine-tuning (3 epochs, less than 1% of pretraining data) achieves 75% key cache savings at roughly 2% quality cost on both GPT-2 and Mistral-7B. The approach composes with GQA and quantization for up to

16\times

combined key cache compression. For a 7B model serving a 128K context, factored keys save 25 GB of KV cache per user, enabling roughly 60% more concurrent users on identical hardware.

Hengshuai Yao, Xing Chen, Ahmed Murtadha +1

Sapient Intelligence · Department of Computing Science, University of Alberta

Memory-Efficient Attention KV-Cache Compression LLM Compression Low-Rank Attention Dimensionality Reduction