Memory-Efficient Optimization

Follow this topic Ask about this topic

Momentum

1 paper in the last four weeks, down 83% on the four weeks before. 0.0% of all new papers.

Jul 13Week of Sep 28

Latest papers 64

Oct 1, 2026cs.LG

AF-Muon: An AdamW-Free Muon Optimizer for Tied-Embedding Models

Muon improves large-scale training by applying a spectral-norm steepest-descent update to matrix parameters, but practical models also contain parameter blocks that do not fit dense-matrix geometry. One important case is the tied vocabulary table, which appears in language models and other token generators and can receive multiple structurally different gradient sources, from sparse input lookups to dense output-classifier updates. In the reference recipe these blocks are handed to an auxiliary AdamW optimizer, which restores second-moment state and updates the aliased table as a generic tensor. We propose AF-Muon, an AdamW-free extension of Muon that keeps the Muon matrix update for hidden weight matrices while using a support-aware finite-cap linear minimization oracle for tied vocabulary tables and an RMS-normalized update for one-dimensional auxiliary parameters. AF-Muon therefore trains every parameter class with a single first-moment buffer and no second-moment state, saving around 20% optimizer-state memory relative to Hybrid Muon in our benchmark. Across nine tied-token settings - decoder-only language models from 124M to 1B parameters, a fully shared T5-style encoder-decoder, and ImageGPT-style image-token, protein, and sparse-MoE variants, spanning text, image, and protein-sequence data - AF-Muon improves mean validation loss and perplexity over both Hybrid Muon and a SCION-style Sign endpoint. Long-horizon runs and hyperparameter sensitivity studies confirm the gain is robust, and identical-momentum diagnostics attribute it to the finite cap, which preserves more within-row magnitude than Sign while bounding the coordinate concentration of row-RMS. These results identify tied vocabulary tables as a distinct optimizer geometry and yield a robust AdamW-free Muon variant across models, modalities, and architectures, with about 1% step-time overhead in matched training.

Arash Lagzian, Paniz Halvachi, Junming Zhang +2

National University of Singapore, Singapore · Independent Researcher · LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University, Shanghai, China

Memory-Efficient Optimization Muon Optimizer LLM Training

Sep 3, 2026cs.AI

GrowPage: On-Demand KV Budgeting for Efficient LLM Reasoning Serving

Long-output reasoning has made the key--value (KV) cache a critical memory bottleneck for efficient LLM serving. Existing KV compression methods usually rely on a predefined per-request budget and adjust only which KV states are retained, leaving the total capacity fixed throughout decoding. However, reasoning workloads exhibit substantial demand variation: different requests require different KV capacities, and the attention demand of an individual request evolves during generation. We introduce \textbf{GrowPage}, an on-demand KV budgeting framework that treats KV capacity as a runtime resource. GrowPage maintains lightweight dual-timescale query summaries to capture recent and long-term attention behaviors, and uses their relative attention working sets to estimate demand evolution. At each capacity boundary, GrowPage either compresses KV states within the current allocation or acquires an additional physical page when broader demand emerges. By integrating with PagedAttention's page-level memory abstraction, GrowPage preserves continuous batching and prefix caching. Experiments on reasoning benchmarks across multiple models show that GrowPage achieves a superior performance--throughput trade-off over existing approaches.

Qiankun Ma, Yanjiang Zhou, Zinan Xiong +5

Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences · Pengcheng Laboratory · University of Chinese Academy of Sciences

LLM Inference KV Caching Memory-Efficient Optimization KV-Cache Management LLM Inference Acceleration KV-Cache Compression+1

Sep 3, 2026cs.CL

Random Attention: Rethinking KV Cache Eviction for Efficient Reasoning

Large language models achieve superior performance on tasks that require extended reasoning, but long chains of thought make the KV cache a severe memory bottleneck. Existing KV cache compression methods share one paradigm: score each cached token by some estimate of how much it will matter later, and keep the top-scoring ones. We show that the selection signal contributes almost nothing. Random Attention keeps the prompt and evicts uniformly at random within each attention head, computing no score at all; across four models and six reasoning tasks, it matches the strongest baseline in task performance while delivering 32-43% higher throughput than that method when deployed with vLLM. Controlled experiments explain this by showing that 1) the prompt is the fragile part of the cache, and most of the gap between selectors is just whether their selection signal happened to keep it; 2) the reasoning trace protects itself against eviction with redundancy at two levels, in the text (the model restates what it still needs as it works) and across attention heads (each keeps its own copy of the trace), so once the prompt is safe, a random draw retains enough copies of what the model still needs, and no score is required to pick them. Our code is publicly available at https://github.com/SalesforceAIResearch/Random-Attention.

Heng Wang, Jielin Qiu, Wenting Zhao +8

Salesforce AI Research · University of Illinois Urbana-Champaign

KV Caching Memory-Efficient Optimization KV-Cache Eviction Efficient Language Model Reasoning LLM Inference Acceleration KV-Cache Compression

Sep 2, 2026cs.LG

LeanStream: A Speculate-and-Refine Streaming Framework for Efficient on-Device LLM Inference

On-device LLM inference is attractive for privacy and responsiveness, but remains challenging on mobile and embedded devices because model weights far exceed available DRAM. Prior systems exploit activation sparsity and offload weights to SSD or flash storage, but face a fundamental systems trade-off: accurate sparse execution decisions require the latest context, whereas efficient computation-I/O overlap requires early prediction. As a result, existing designs either serialize execution or incur redundant weight fetches, extra computation, and large cache overheads. We present LeanStream, a streaming speculate-and-refine framework for efficient on-device LLM inference. LeanStream progressively refines computation, loading, and cache-retention priorities using partial GPU results, enabling fine-grained overlap between GPU execution and storage I/O. We implement LeanStream on both mobile and embedded platforms. Compared with prior on-device LLM inference systems, LeanStream reduces memory usage by 4.8

\times

to 7.5

\times

at the best throughput achieved by prior work, while further improving token generation throughput by 1.6

\times

to 2.1

\times

.

Renyuan Liu, Yuyang Leng, Kaiyan Liu +7

Richard · George Mason University · Global Technology Applied Research, JPMorganChase +1

LLM Inference Efficiency On-Device Language Model Inference Memory-Efficient Optimization Memory-Efficient Inference Efficient Language Model Inference

Aug 23, 2026cs.LG

Beyond Dense Adam States: Adaptive Log-Space Quantization for Memory-Efficient Optimizers

Optimizer-state quantization is commonly designed for Adam's dense, parameter-aligned first- and second-moment arrays. This abstraction breaks for memory-efficient optimizers, whose states may be factored, confidence-modulated, or maintained in a projected space, so similar reconstruction error can produce different update error. We formulate optimizer-state quantization as a joint problem over representation, topology, and update semantics. We then introduce Adaptive Log-Space (AL) quantization for non-negative states. AL fits each block's observed nonzero logarithmic interval and reserves a separate code for exact zero, enforcing

q = 0 \Leftrightarrow x = 0

; signed momentum and state precision remain independently selectable. Controlled probes show that adaptive ranges reduce update error and temporal drift, exact-zero reservation preserves dormant states, and state topology constrains useful block granularity. End-to-end language-model training evaluates the resulting policy across dense, factored, confidence, and projected optimizer states. On TinyLlama-1.1B, AL8 with uniform 8-bit momentum reaches 72.90 perplexity versus 73.54 for bitsandbytes 8-bit AdamW, with comparable optimizer-state storage and higher throughput. CAME matches reference-level final perplexity across three seeds when its non-negative states use AL16, while a semantic grouping-and-protection policy closes most of quantized Adafactor's 100K-step late-loss gap. These results make state topology and update semantics first-class design constraints for optimizer quantization.

Yan Wang

Independent Researcher

Deep Learning Optimization Memory-Efficient Optimization Low-Bit Quantization Neural Network Quantization

Aug 16, 2026cs.LG

SubZero+: Memory-Efficient Adaptive Zeroth-Order LLM Fine-Tuning in Random Subspaces

Zeroth-order (ZO) optimization with SGD in random subspaces enables memory-efficient fine-tuning of large language models without backpropagation. However, high gradient estimation noise fundamentally undermines adaptive optimizers like Adam. We propose SubZero+, which achieves practical adaptive ZO optimization through a carefully designed dual low-dimensionality strategy: (i) multi-query forward-difference gradient estimation in periodically refreshed random subspaces to mitigate noise amplification in moment buffers, and (ii) Adam updates with periodic restarts performed directly in low-dimensional space rather than full-parameter space. In experiments, this dual design retains memory overhead comparable to momentum-free ZO methods while achieving stronger optimization performance than the evaluated ZO baselines. Theoretically, in the exact-directional limit,

K

-query averaging preserves conditional unbiasedness, while the coefficient estimator's covariance and mean-squared error, as well as query-induced second-moment inflation, scale exactly as

1/K

. Extensive experiments across SuperGLUE with models from 1.3B to 32B parameters under both full fine-tuning and LoRA schemes demonstrate consistent improvements over competing ZO methods. SubZero+ significantly narrows the performance gap with first-order optimization while preserving ZO's inference-time memory efficiency.

Ziming Yu, Shuyao Xiao, Xingyu Zhao +6

Beijing Normal University · Singapore Management University · Xiaomi Inc. +2

Memory-Efficient Fine-Tuning Zeroth-Order Optimization LLM Fine-Tuning Memory-Efficient Optimization

Aug 11, 2026cs.LG

Batch Size or Negatives? A Selection Rule for Memory-Constrained Recommender Training

Large-scale neural recommender systems are typically trained with a softmax cross-entropy objective over the full item vocabulary. For a typical large number of possible items

K

, the final classification layer dominates memory, requiring

O(nK)

logits and gradients to materialize for a batch of

n

examples. Sampled softmax reduces this cost by restricting the objective to only

k \ll K

candidate negative items, resulting in an

O(nk)

memory. However, for a fixed budget

B = n k

, it remains unclear whether one should prioritize larger batches or the inclusion of more negative items. We address this question by analyzing sampled-softmax training under a fixed memory constraint. Under standard smoothness and variance assumptions, our theoretical evidence suggests that the fastest convergence arises from an

n \sim B, k \sim 1

allocation. So, an actionable rule is to include as many objects as possible given computational constraints. Our theory is supported by controlled synthetic and synthetic and four real sequential recommendation benchmarks, including MovieLens-20M. The suggested configuration achieve faster convergence and better final recommendation quality than imbalanced alternatives within the same memory constraint. These findings provide a theoretical and empirical foundation for configuring memory during the training of recommender systems. Code, reproducibility materials, and all scripts for generating figures are available at https://anonymous.4open.science/r/LimitedMemoryRule-BBFB

Artyom Sabitov, Daniil Volkov, Alexey Zaytsev

Applied AI Institute, Moscow, Russia Risk department

Memory-Efficient Optimization Recommender Systems Hyperparameter Optimization Memory-Efficient Training

Aug 9, 2026cs.LG

Memory-Efficient Activation Checkpointing with Sliding Window and Hirschberg's Algorithm for 0/1 Knapsack Solving in PyTorch

Activation checkpointing minimizes the runtime of neural networks under a given memory budget, by selecting which intermediate tensors to store and which to recompute. PyTorch solves this as a 0/1 knapsack problem, where operations from a joint forward-backward computation graph are items with a memory cost (weight) and a runtime saving (value). The default solver, dp_knapsack, allocates a full dynamic programming (DP) table of shape

(n+1) \times (W+1)

, where

n

is the number of operations and

W

is the quantized memory budget. This method is resource-hungry and crashes at

n = 100

items on a machine with 64 GB RAM. In this paper, we introduce dp_knapsack_sliding_hirschberg, which combines the sliding window trick and Hirschberg's algorithm to reduce peak memory from

O(nW)

to

O(W)

while preserving the exact optimal solution. Our experiments show successful knapsack execution at

n = 2000

, where dp_knapsack fails at

n = 100

, a 20

\times

increase in computable problem size. In addition, our benchmarks show a consistent 25-28% runtime speedup over dp_knapsack. The implementation is merged into PyTorch and released in version 2.10.

Jędrzej Maczan

Cohere Labs Community Poland

Dynamic Programming Memory-Efficient Optimization Combinatorial Optimization Memory-Efficient Training

Aug 8, 2026cs.LG

ZeroLock: Concurrent Memory-Efficient LLM Training via Modular Update Decoupling

Large language model (LLM) fine-tuning at the edge adapts the model to scenario-specific data while preserving privacy. Although existing studies proposed pipeline parallelism to address the limited memory and computing resources of edge devices, they commonly rely on backpropagation (BP) training, which has a fundamental limitation of update locking and could experience severe throughput and memory bottlenecks. In this work, we propose a BP-free algorithm, called ZeroLock, that decouples the model updates into independent chunk updates by local objective construction. It breaks the update locking of BP and hence can improve throughput at the algorithm level and lower memory usage by reducing activation storage. To the best of our knowledge, we provide the first theoretical framework for such local objective construction-based approaches under general model chunk division by mapping local objectives to the global objective. We prove that ZeroLock has a convergence rate of

\tilde{\mathcal{O}}(1/\sqrt{T})

, which differs from BP only by polylogarithmic factors. We design a system for ZeroLock and build real-world prototypes, incorporating techniques such as early forwarding and failure recovery for efficient and robust implementation. Experiments on the prototype show that compared to BP-based baselines, ZeroLock reduces the memory by 26.5% and improves throughput by 4.9%.

Wentao Dai, Xuanran Li, Yuxiang Zhang +2

Department of Computer Science and Engineering, Southern University of Science and Technology, China · Department of Mathematics, Southern University of Science and Technology, China · School of Computing, Montclair State University, New Jersey, USA

On-Device Training Memory-Efficient Fine-Tuning LLM Fine-Tuning Neural Network Optimization Memory-Efficient Optimization Memory-Efficient Training

Aug 5, 2026cs.LG

MESH: Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training

Memory-efficient matrix optimizers such as Sinkhorn gradient descent remove most AdamW optimizer state for dense Transformer matrices, but direct application to Mixture-of-Experts (MoE) training is unreliable. We study this failure in a controlled 110M-parameter nanowhale DeepSeek-style MoE pretraining setting. A SAGE/Sinkhorn hybrid reduces optimizer state from 0.883GB to 0.331GB but degrades evaluation loss to 3.8265, far above the AdamW baselines observed in the same setup (3.58--3.64 across the seeds we study). We show that routed MoE expert matrices are the dominant failure point: their gradients are conditional, temporally varying, and poorly served by stateless Sinkhorn normalization. We propose MESH, a hidden-momentum Sinkhorn update for MoE experts. MESH restores a temporal first-moment signal through the gradient-buffer lifecycle, without storing the expert first moment as optimizer state. MESH is an optional block-preconditioned variant that adds a coarse neuron/block inverse-RMS multiplier. Across ablations, temporal smoothing before matrix normalization is the primary causal ingredient; block/neuron preconditioning can improve the memory-quality frontier, but is not established as universally necessary. In two additional seeds, MESH and MESH-B reduce optimizer-state memory by 62.5% and peak PyTorch CUDA allocation by about 12.6% relative to AdamW, with a modest evaluation-loss gap. Full-state diagnostic variants recover AdamW-like performance in ablations, supporting the conclusion that MoE experts need temporal smoothing, but not necessarily full coordinate-wise AdamW state.

Masato Fujitake

FA Research, Fast Accounting Co., Ltd.

Neural Network Optimization Memory-Efficient Optimization Matrix Optimization Mixture-of-Experts Models Memory-Efficient Training

Aug 4, 2026cs.AI

TaskPress: Query-Agnostic KV Cache Compression via Task-Guided Pruning

Long-context inference with large language models is constrained by the linear growth of the key-value cache to sequence length. While pruning offers mitigation, prevailing methods determine query-specific token importance that cannot be reused across unseen queries. In contrast, we introduce TaskPress, a framework for task-guided, query-agnostic KV cache eviction. Instead of optimizing the cache for a single query, TaskPress constructs a reusable memory representation conditioned on a high-level task guide. The guide functions as a meta-query during prefill to filter irrelevant tokens before downstream queries are issued. In addition, TaskPress leverages quantization scale factors as a zero-cost signal for detecting influential representation outliers, providing an efficient proxy for token importance. Experiments on conducted on various tasks with long context input demonstrate that TaskPress efficiently creates a compact, reusable cache across diverse queries.

Wonpyo Park, Seung-won Hwang

Seoul National University · Google

LLM Pruning KV Caching Long-Context Language Model Inference Memory-Efficient Optimization KV-Cache Management KV-Cache Compression+1

Aug 3, 2026cs.LG

AnchorKV: Anchor-Residual KV Cache Compression

The key-value (KV) cache is the primary memory bottleneck in long-context LLM inference. Existing approaches attack it from opposite ends: eviction methods permanently discard tokens, degrading performance whenever a discarded token later proves essential, while quantization methods retain all tokens at low precision but offer limited compression. We propose AnchorKV, a compression scheme that shrinks the cache by

20\times

without discarding a single token. AnchorKV represents the cache using a small set of anchors stored exactly, expresses every other token through its most similar anchor, and refines only those whose approximation most affects the model's output. AnchorKV consistently preserves accuracy across models and datasets, retaining 99% of the full-cache score at the 70B scale, while keeping the entire context at a fraction of its cost.

Malik Khalaf, Yara Shamshoum, Nitzan Hodos +2

Department of Computer Science, Technion - Israel Institute of Technology

KV Caching Long-Context Language Model Inference Memory-Efficient Optimization Memory-Efficient Inference KV-Cache Compression

Aug 2, 2026cs.CL

RestoreKV: Recovering Full-Cache Behavior Under Aggressive Query-Agnostic KV Cache Eviction

Query-agnostic KV cache eviction compresses a context once and reuses the resulting cache for arbitrary future queries, but performance can collapse under tight budgets. Existing methods primarily improve which original KV pairs are retained. We introduce RestoreKV, which complements this selection-based formulation with learned restoration under the same total KV budget. Our key insight is that, although the information lost through eviction is context-specific, the mechanism for generating its compact complement can be shared across contexts. After context prefill, a few restore tokens attend to the full KV cache in a single LoRA-adapted pass, generating a compact, context-conditioned restore cache. The base importance scorer and eviction rule remain unchanged, and the adapters are disabled for all subsequent queries and decoding. RestoreKV is trained through parameter-efficient self-distillation from the frozen full-cache model, optimizing only

0.4\%

of the parameters and requiring no task-specific tuning. Across four backbones and four long-context benchmarks, RestoreKV substantially reduces compression-induced degradation. On Qwen3-4B, it improves 59 of 60 paired, budget-matched settings across five base eviction methods; at a

5\%

budget, it raises KVzip from

38.2

to

73.2

on RULER-4K. Applied to KVzip+, RestoreKV reaches

86.4

RULER accuracy at

16\times

compression on the KVPress Benchmark, while adding less than

0.5\%

one-time cache-construction overhead in a 32K-context evaluation. Our project page is available at https://paper.pnu-cvsp.com/RestoreKV/

Changwoo Baek, Seungjun Shin, Kyeongbo Kong

Pusan National University · Sookmyung Women’s University

KV Caching Memory-Efficient Optimization KV-Cache Eviction KV-Cache Compression

Jul 24, 2026cs.AR

HiKV: Hierarchical Importance-Aware KV Cache with Hardware Acceleration for LLM Decoding

With the rapid adoption of long-context large language models (LLMs), the continuously growing KV cache during decoding has become the critical memory bottleneck. To tackle this challenge, we propose HiKV, a novel algorithm-hardware co-design that exploits KV cache redundancy through hierarchical importance awareness. Algorithmically, HiKV compresses the KV cache at two granularities: Stage I evicts unimportant tokens within a fixed budget, and Stage II further loads only the significant elements of each retained token, reaching compression ratios unattainable at a single granularity. Architecturally, we develop a dedicated accelerator centered on a reconfigurable importance sorter that switches between the distinct sorting datapaths each stage requires, unifying the two-stage acceleration in one circuit with minimal overhead. Evaluated on representative LLMs, HiKV achieves up to 7.95x speedup and 90% energy reduction in the attention computation over the vanilla KV cache baseline within negligible 1% accuracy loss. Under iso-accuracy constraints, HiKV outperforms state-of-the-art importance-based methods by achieving an additional 1.82~4.87x reduction in external memory accesses. These benefits are enabled by specialized hardware components that add only 8% to the system area.

Chao Fang, Jun Yin, Man Shi +1

ESAT-MICAS, KU Leuven, Leuven, Belgium · UC Berkley, CA, USA

KV Caching Memory-Efficient Optimization KV-Cache Eviction LLM Inference Acceleration KV-Cache Compression KV-Cache Pruning

Jul 21, 2026cs.LG

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

Efficient Transformer Inference Grouped-Query Attention Transformer Inference GPU Kernel Optimization KV Caching Memory-Efficient Optimization+2

Jul 21, 2026cs.LG

Where Should Optimizer State Live? Tiered State Allocation for Memory-Efficient Mixture-of-Experts Training

Optimizer state is the largest single line item in the memory budget of mixture-of-experts (MoE) training. On a 6.78B-parameter MoE language model AdamW keeps 50.6 GB of first and second moments to update 12.6 GB of bfloat16 weights. We study SkewAdam, an optimizer built on the observation that the three parameter populations of an MoE differ enough in size and gradient statistics that they should not receive the same state. Those populations are the dense backbone, the experts and the router. SkewAdam keeps float32 momentum plus a factored second moment for the backbone (5% of parameters), a factored second moment alone for the experts (95%) and an exact second moment for the router (<0.01%). The resulting state occupies 1.29 GB or 2.6% of AdamW's and peak training memory falls from 81.4 GB to 31.3 GB, within the budget of a 40 GB accelerator. In a controlled comparison from identical initializations over 82M tokens, SkewAdam reaches validation perplexity 108.4, ahead of AdamW (126.8), Muon (120.2) and Lion (393.7), and settles router load balance to within 1% of its uniform floor. The allocation is not what earns that perplexity. A tier ablation reaches the same value while carrying twenty times the state, so the tiers buy memory rather than accuracy. Same-platform runs separate what does earn it. Removing momentum costs 31 perplexity points (tuned Adafactor, 139.7) and replacing the factored second moment and its update clipping with a full second moment costs 10 (tuned AdamW, 118.5), so neither tuned baseline reaches the untuned tiered policy. Where optimizer state lives, these results suggest, matters at least as much as how much of it there is.

Nuemaan Malik

Independent Researcher

Deep Learning Optimization Mixture-of-Experts Language Models Memory-Efficient Optimization Efficient Language Model Training Memory-Efficient Training

Jul 21, 2026cs.DS

Stronger Memory-Query Tradeoffs for Convex Optimization: The Limitations of Subquadratic Memory

We prove two lower bounds for the first order oracle complexity of minimizing a

d

-dimensional

1

-Lipschitz convex function over the unit ball with

m

bits of memory. We first show that any such (possibly randomized) algorithm must make

\tildeΩ(\frac{d^2}{\sqrt{m}})

oracle queries. For deterministic optimization algorithms, we show that

\tildeΩ(\min\{d^{1.6},\frac{d^{8/3}}{m^{2/3}}\})

queries are required. For all memory regimes of interest, these improves upon the previous best known lower bounds of

\tildeΩ(\max\{\frac{d^{8/3}}{m^{4/3}},\frac{d^{4/3}}{m^{1/6}}\})

and

\tildeΩ(\frac{d^{5/3}}{m^{1/3}})

for randomized and deterministic algorithms respectively. Notably, due to existing upper bounds, our lower bound for deterministic algorithms is the first to show a sharp oracle complexity phase transition around

m\approx d^2

, where a polylogarithmic change in memory leads to a

\mathsf{poly}(d)

change in the number of required oracle calls. Further, when the suboptimality is polynomially small in

d

, our lower bound randomized algorithms is the first to show that

\tildeΩ(d^2)

memory is necessary to nearly match the optimal query complexity among algorithms without memory constraints. Previously, such a result was only known for the regime where the suboptimality is quasipolynomially small in

d

.

Michael Menart, Aleksandar Nikolov, Ohad Shamir

Department of Computer Science, University of Toronto, Vector Institute · Department of Computer Science, University of Toronto · Weizmann Institute of Science and University of Toronto +1

Query Complexity Oracle Complexity Memory-Efficient Optimization Constrained Optimization

Jul 17, 2026cs.LG

PagedWeight: Efficient MoE LLM Serving with Dynamic Quality-Aware Weight Quantization

Mixture-of-Experts (MoE) is a popular class of large language models (LLMs), offering high efficiency and accuracy. However, in KV-cache-intensive serving scenarios, MoEs often exhibit a tension between the GPU memory requirements of the model weights and the growing KV cache. We propose PagedWeight, a novel management method for MoE LLM serving that dynamically quantizes MoE model's weights at runtime and balances expert-weight precision with the KV cache sizes. PagedWeight exposes and effectively navigates the complex tradeoff between the model's task accuracy, memory consumption, and throughput/latency. Across several memory-sensitive MoE serving scenarios, PagedWeight improves the quality-memory tradeoff over several existing quantization baselines. PagedWeight achieves FP16-equivalent accuracy with up to 72.0% GPU memory savings and 1.94

\times

throughput improvement, and improves quality over quantization methods by up to 39.3% at a similar memory budget with at most 4.1% throughput loss.

Yuchen Yang, Yifan Zhao, Anisha Dasgupta +1

University of Illinois Urbana-Champaign

LLM Quantization Mixture-of-Experts Language Models Memory-Efficient Optimization LLM Inference Acceleration Mixture-of-Experts Quantization Weight-Only Quantization

Jul 14, 2026cs.CL

Adaptive Filtering of the KV Cache: Diagnosing and Correcting Structural-Role Bias in LLM Inference

Attention-based KV cache eviction (H2O and its descendants) compresses the memory-constrained state of a long-context model by ranking tokens on accumulated attention mass, treated here as signal energy, and keeping the heaviest. On schema-dense input streams such as nested JSON, this score acts as a non-stationary filter that disproportionately retains noise: a non-content sink role (delimiters or whitespace) carries an order of magnitude more energy than any content role, and structural KEY tokens are over-retained at roughly 1.8x the rate of the answer-carrying VALUE tokens, collapsing exact-match accuracy from 88% to 0% at a 5% budget as the signal-to-noise ratio of the retained state degrades. A counterfactual experiment establishes that suppressing KEY tokens is the best deployable filter. Our retraining-free, role-conditional allocation over SnapKV's windowed score, governed by a single tuned hyperparameter, closes 63-98% of the H2O gap at sub-20% budgets and, at higher budgets, modestly matches or exceeds full-cache accuracy -- a small, seed-sensitive denoising effect (borderline significant at B=0.50; not distinguishable from zero at B=0.30 over four seeds). A 15 MB linear role probe supplies these labels at negligible inference cost, though matching parser-level downstream accuracy remains open.

Soumil Mandal

ServiceNow

KV Caching Long-Context Language Model Inference Memory-Efficient Optimization KV-Cache Eviction KV-Cache Pruning

Jul 14, 2026cs.AR

Full-Pipeline Inference Optimization for MiMo-V2.5 Series: Pushing Hybrid SWA Efficiency to the Limit

We present a full-pipeline inference optimization for the MiMo-V2.5 model family, which combines Hybrid Sliding Window Attention (Hybrid SWA), sparse Mixture-of-Experts (MoE), and multimodal encoders. While Hybrid SWA can ideally reduce both attention compute and KVCache storage significantly compared to Full Attention, realizing these gains in production requires substantial engineering effort. We systematically optimize the KVCache system with layerwise prefetch, SWA-aware prefix cache trees, and specialized placement strategies, achieving strict

O(W)

SWA storage and high cache hit rates. We further build GCache, a high-performance distributed cache infrastructure with RDMA-optimized networking, and develop a KVCache-affinity router to reduce computation while preserving load balancing. We also optimize for multimodal inputs, including GPU image preprocessing, parallel video decoding, and multimodal cache sharing. Together, these optimizations constitute the first large-scale LLM serving system in production that efficiently covers the Hybrid SWA + MoE + multimodal composite architecture.

Xiaomi MiMo Team, Anqi Liu, Aoxin Ma +28

MiMo Team, Xiaomi

Efficient Multimodal Inference KV Caching Memory-Efficient Optimization KV-Cache Management LLM Inference Serving LLM Inference Acceleration

Jul 7, 2026cs.AI

DepthWeave-KV: Token-Adaptive Cross-Layer Residual Factorization for Long-Context KV Cache Compression

Long-context language model inference is increasingly limited by the memory bandwidth and capacity required to store key-value caches, yet existing compression methods often apply uniform budgets across layers or tokens and degrade retrieval when lexical cues and semantic states require different preservation. We introduce DepthWeave-KV, a token-adaptive cache compression method that factorizes key and value states across neighboring transformer layers using shared low-rank channel bases while retaining lightweight token-specific residuals where attention behavior is sensitive. DepthWeave-KV combines cross-depth residual factorization with a token-conditional depth router that allocates higher reconstruction rank to instruction-bearing and retrieval-critical tokens, and uses calibration-free online error tracking from attention-output probes to adapt compression during generation without retraining the base model. A fused CUDA implementation jointly performs basis lookup, residual dequantization, and attention projection to reduce decode-time memory traffic. Across LongBench, Needle-in-a-Haystack, L-Eval, and long-form QA and summarization benchmarks, DepthWeave-KV achieves near-full-cache task quality with substantially lower memory use, improving average score and retrieval accuracy over prior compressed caches while reaching 8.3x KV memory reduction and 72.8 tokens per second at 64K context.

Anna Cordoba, Adam Puente Tercero, Nerea Angulo Hijo +4

Instituto de Investigación en Visión Artificial

KV Caching Long-Context Language Model Inference Memory-Efficient Optimization LLM Inference Acceleration Low-Rank Compression KV-Cache Compression

Jul 7, 2026cs.AI

FreqDepthKV: Frequency-Guided Depth Sharing for Robust KV Cache Compression in Long-Context LLM Inference

Long-context LLM inference is increasingly limited by the memory and bandwidth cost of KV caches, yet aggressive compression can remove the layer-specific evidence needed for retrieval and multi-step reasoning. We introduce FreqDepthKV, an inference-time cache compression method that factorizes adjacent-layer KV states into shared low-frequency depth components and sparse high-frequency residuals. A lightweight online probe assigns attention heads to shared-depth, residual-depth, or exact cache modes according to their contribution to reconstruction-sensitive attention logits, allowing the compression policy to adapt to prompt structure without retraining. Across long-context question answering, needle retrieval, summarization, and code generation benchmarks, FreqDepthKV preserves task accuracy under substantially smaller cache budgets. With a 32k-token prefill window, FreqDepthKV reaches 58.3 Exact Match, 63.0 F1, 32.5 ROUGE-L, and 48.1 pass@1, closely matching full KV while outperforming prior compressed-cache methods. It also improves decoding throughput to 70.4 tokens/s, reduces TTFT to 2.06 seconds, and lowers peak KV memory to 6.2 GB, achieving a 3.9x effective compression ratio.

Anna Córdoba, Adam Puente Tercero, Nerea Angulo Hijo +4

Instituto de Investigación en Visión Artificial

KV Caching Long-Context Language Model Inference Memory-Efficient Optimization LLM Inference Acceleration KV-Cache Compression

Jul 1, 2026cs.LG

Token Geometry

Language models learn continuous programs over discrete symbols, with the embedding table and LM-head acting as the read/write interface between them. We show that this interface has gradient geometry distinct from dense hidden weights which can be exploited to improve the Pareto frontier across supervised finetuning, RL, and pretraining, while only utilizing kilobytes of optimizer state. We introduce Ember, a lightweight optimizer for embedding and LM-head matrices that utilizes O(V + D) VRAM, instead of Adam's O(2VD), and forgoes the need to shard both token table optimizer states. We provide empirical evidence that Ember scales effectively across batch size and parameter count. We show that the optimization trajectory of tokens can be well described by a simple 1D ray, counter to the popular belief that neural net parameters navigate a heavily nonconvex landscape. We provide a principled view on the surprisingly narrow space of optimizers that suffice for Transformer training. Finally, we open-source our distributed Ember implementation that merges cleanly with existing ZeRO/FSDP setups to support further research at https://github.com/katop1234/ember

Kathan Shah

Preprint.

Language Model Pretraining LLM Fine-Tuning Memory-Efficient Optimization Efficient Language Model Training Representation Geometry in Language Models

Jun 22, 2026cs.LG

FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training

Reverse-mode differentiation computes every weight gradient, writes it to memory, and only then lets the optimizer read it back. This two-phase schedule sets the memory ceiling of modern training: at the seam between the phases, every layer's gradient is live at once. We argue that this materialized gradient is an artifact of how differentiation is staged, not a quantity that learning requires -- and we eliminate it. FORGE folds the optimizer step into the backward pass and applies it one tile at a time, entirely in registers, so each gradient tile is consumed the instant it is produced and never becomes a tensor. The fusion changes only when the update happens, not what it computes: in full precision the fused step is provably exact -- the identical optimizer update, for every element-wise rule -- and that exactness survives tensor- and sequence-parallel sharding; in the bf16 and 8-bit regimes used in practice it is faithful rather than bit-identical, its deviation bounded and, for the weight store, rendered unbiased by stochastic rounding. Because each gradient tile is born and consumed in the same registers, it is never converted down to bf16 to be stored and read back; FORGE thus preserves the full-precision fidelity that both bf16 and 8-bit optimizers lose to that conversion. Nor is the method tied to one architecture or one optimizer: linear layers are ubiquitous, and FORGE reclaims the gradient memory of any of them under any element-wise rule. Empirically FORGE more than halves the memory of an optimizer step and, at the small batch sizes typical of fine-tuning and continued pretraining, runs about 1.5x faster; integrated into tensor-parallel Megatron-LM it fits 8B training at four times the micro-batch a standard optimizer allows on the same GPUs.

Dikshant Kukreja, Kritarth Prasad, Avinash Anand +6

Puch AI · Singapore Institute of Technology · Nanyang Technological University +2

Automatic Differentiation Backpropagation Memory-Efficient Optimization Efficient Language Model Training LLM Training Memory-Efficient Training

Jun 20, 2026cs.CL

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

Jun 14, 2026cs.LG

SILAGE: Memory-Efficient, Full-Gradient-Free Nonconvex Optimization for Nested Finite Sums

Empirical risk minimization on massive datasets naturally exhibits a nested double finite-sum structure, where

N=nm

total samples are logically or physically partitioned into

n

blocks of size

m

(e.g., in pooled data silos, out-of-core learning, or deliberate stratification). While variance-reduced methods achieve optimal oracle complexities for nonconvex objectives, they suffer from severe scaling bottlenecks in this centralized regime. Recursive estimators, such as PAGE, require periodic global full-gradient refreshes over all

nm

samples, which are computationally expensive. Conversely, single-loop methods, such as SILVER, avoid such refreshes but require an impractical

\mathcal{O}(nm)

memory footprint to store a control variate for every sample. In this paper, we propose SILAGE, a variance-reduced algorithm that addresses this trade-off. By actively exploiting the double-sum structure, SILAGE eliminates periodic global full-gradient refreshes over all

nm

components (evaluating at most one local group gradient per iteration) while requiring only

\mathcal{O}(n)

memory. Furthermore, we provide a tight convergence analysis that avoids pessimistic worst-case Lipschitz constants. Instead, SILAGE's complexity natively adapts to the underlying data geometry via nested functional similarities: across-group (

δ_1

) and within-group (

δ_2

) heterogeneity. Our results improve existing state-of-the-art bounds in several practically relevant regimes.

Igor Sokolov, Laurent Condat, Peter Richtárik

*Center of Excellence for Generative AI, King Abdullah University of Science and Technology (KAUST), Thuwal, Saudi Arabia. · KAUST

Memory-Efficient Optimization Nonconvex Stochastic Optimization Stochastic Optimization Convergence Variance Reduction

Jun 11, 2026cs.LG

Gefen: Optimized Stochastic Optimizer

AdamW is a default optimizer for deep learning, but its moment states add two parameter-sized buffers to training memory, increasing the cost of large-scale pretraining. We propose Gefen, a memory-efficient optimizer that automatically shares second-moment estimates across parameter blocks and quantizes the first moment using a learned codebook. Gefen reduces AdamW's optimizer memory footprint by up to 8x while maintaining performance, saving 6.5 GiB per billion parameters. Prior work shares second moments across parameters grouped along the Hessian's block-diagonal structure, but relies on hand-specified architectural rules and leaves unexplained why such grouping works. We prove that large mixed Hessian entries constrain the ratio of squared gradients toward one, explaining why shared second moments are accurate when the squared gradients they pool are similar. The Hessian need not be computed: its block structure is inherited by squared gradients, allowing blocks to be found directly. Gefen therefore infers block structure from initial squared gradients, requiring no architecture-specific metadata or user-tuned hyperparameters beyond AdamW defaults. Gefen learns an exact histogram-based dynamic-programming quantization codebook and reuses the blocks for first-moment scaling. Across diverse pretraining experiments, Gefen achieves the lowest peak optimizer memory among compared methods that maintain AdamW-level performance. In single-machine and distributed training, the reduced footprint enables larger microbatches and substantially improves throughput over AdamW, making Gefen a drop-in replacement that can train larger models or use larger global batch sizes. We provide the complete Python implementation, including fused CUDA kernels at https://github.com/ndvbd/Gefen

Nadav Benedek, Tomer Koren, Ohad Fried

Reichman University · Tel Aviv University, Google Research

Deep Learning Optimization Memory-Efficient Optimization Adaptive Gradient Methods Memory-Efficient Training

Jun 9, 2026cs.AI

ReasonAlloc: Hierarchical Decoding-Time KV Cache Budget Allocation for Reasoning Models

Long chain-of-thought (CoT) trajectories in large language model (LLM) reasoning cause severe inference bottlenecks due to rapid key-value (KV) cache growth. Current decoding-time compression methods mitigate this issue via token eviction, but typically assume a uniform budget distribution across all layers and heads. In contrast, existing non-uniform budget allocation methods are predominantly designed for the static prompt prefill phase, and they do not capture the stepwise context demands of autoregressive reasoning. To bridge this gap, we propose ReasonAlloc, a training-free framework that recasts decoding-time KV compression as a hierarchical budget allocation problem. ReasonAlloc operates at two complementary levels: an offline layer-wise preallocation strategy captures an architecture-driven demand pattern which we call ``\textit{Reasoning Wave}'', while an online head-wise strategy reallocates resources during decoding to information-rich heads based on real-time utility. Evaluations on mathematical reasoning benchmarks (MATH-500, AIME~2024) using DeepSeek-R1-Distill-Llama-8B, DeepSeek-R1-Distill-Qwen-14B, and AceReason-14B show that ReasonAlloc outperforms uniform-budget R-KV, SnapKV, and Pyramid-RKV (a baseline enforcing a static, monotonically decreasing layer budget), with the largest gains at small budgets (128-512 tokens). ReasonAlloc is plug-and-play with existing token-eviction policies and introduces negligible inference-time overhead.

Wenhao Liu, Hao Shi, Yunhe Li +7

1Tsinghua University · 2City University of Hong Kong · 3Peking University +2

KV Caching Memory-Efficient Optimization KV-Cache Management LLM Inference Acceleration KV-Cache Compression Online Resource Allocation+1

Jun 5, 2026cs.LG

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

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

Jun 4, 2026cs.LG

Tangram: Unlocking Non-Uniform KV Cache Compression for Efficient Multi-turn LLM Serving

Multi-turn LLM serving accumulates dialogue history whose Key-Value (KV) cache grows with every turn and every user, quickly exceeding the model weights themselves and making memory -- not compute -- the binding constraint on throughput. Non-uniform KV compression, which allocates heterogeneous budgets across attention heads, preserves accuracy far better than uniform schemes, yet remains impractical: modern serving stacks assume identical KV lengths across heads, so heterogeneity traps freed memory as page fragmentation, spends up to 25% of prefill time reclaiming scattered pages, and skews GPU workloads that inflate decode latency by up to

1.7\times

or burn 15--20% of each decode step on re-planning. We observe that this heterogeneity need not be discovered at runtime: head-wise retention follows a two-level structural regularity -- an input-invariant head ranking with narrowly bounded per-head ratios -- that can be calibrated offline from as few as 50 samples. Building on this insight, we present Tangram, a serving framework that statically resolves what prior systems handle dynamically: Budget Reservation fixes each head's post-compression footprint at scheduling time, eliminating page reclamation; Ragged Paging clusters similar-budget heads into independent page tables, turning fragmentation into reclaimable memory; and Ahead-of-Time Load Balancing precomputes balanced GPU partitions with zero runtime planning. Implemented on vLLM, Tangram serves as a drop-in substrate for existing non-uniform compression methods, matching their accuracy while improving end-to-end throughput by up to

2.6\times

over the full-KV baseline. Our implementation is publicly available at https://github.com/aiha-lab/TANGRAM.

Hyungmin Kim, Minsoo Kim, Hongseok Kim +1

LLM Serving KV Caching Memory-Efficient Optimization LLM Inference Acceleration KV-Cache Compression