GradMem: Gradient-Based Memory Techniques

• GradMem is a family of gradient-based memory techniques that iteratively update compact memory representations to preserve essential information.
• In transformer models, GradMem applies test-time gradient descent on meta-learned memory tokens, achieving up to 88% retrieval accuracy with significant context compression.
• GradMem methods extend to federated, multiobjective, and continual learning, effectively mitigating catastrophic forgetting and optimizing distributed update strategies.

GradMem encompasses a family of techniques and methods that utilize gradient-based mechanisms to maintain, write, or update memory in the context of optimization, learning systems, and memory-augmented neural architectures. Recent research demonstrates the application of gradient memory in diverse domains, including compressive context storage in LLMs, federated learning, continual learning, and multiobjective optimization. Across these, gradient-based memory serves as a tool to preserve essential information, alleviate forgetting, and maximize resource efficiency via loss-driven and iterative error-corrective updates.

1. Principle of Gradient-Based Memory

Gradient-based memory ("GradMem") refers to mechanisms where memory representations—usually compact, parameterized states—are iteratively modified using gradients derived from explicit loss functions. In contrast to forward-only memory writers, gradient-based approaches leverage per-sample or per-step optimization to encode, update, or edit memories, explicitly correcting errors and dynamically adapting to the current information state.

In context-rich tasks such as language modeling, the challenge is to compress long contexts efficiently into a writable memory so that downstream queries can be answered without access to the original sequence. GradMem mechanisms operate by treating memory slots as continuous parameters, optimizing them via gradient descent at test time to minimize a self-supervised reconstruction loss over the source context. This paradigm generalizes to optimization and federated settings, where historical gradient or update information is "memorized" and used to correct or regularize new updates via either direct inclusion in optimization steps or by solving constrained quadratic programs for memory-aligned corrections (Kuratov et al., 14 Mar 2026, Luo et al., 2023, Jin et al., 2020, Chen et al., 2022).

2. GradMem in Transformer Memory Compression

In "GradMem: Learning to Write Context into Memory with Test-Time Gradient Descent" (Kuratov et al., 14 Mar 2026), the GradMem method addresses the inefficiency of full-sequence KV-caches in long-context inference for pretrained LLMs. Standard transformer models require $\mathcal O(L c d)$ memory (with $L=$ layers, $c=$ context, $d=$ hidden-dim), which becomes prohibitive when repeatedly querying long documents.

GradMem introduces a set of $m$ prefix "memory tokens" prepended to the frozen model. The write operation consists of per-example test-time gradient descent: starting from a meta-learned initialization $\mathcal M_0$, memory tokens are updated by $K$ steps of gradient descent to minimize an explicit reconstruction loss,

$$\mathcal L_{\mathrm{write}}(\mathcal M;C) = -\sum_{i=1}^N \log f_\theta(t_i\,|\, [\mathcal M; t_{<i}])$$

yielding an adapted $\hat{\mathcal M}$. During reading, only $\hat{\mathcal M}$ and the query are provided to the model; the original context is not required. This approach enables iterative error correction, in contrast to forward-only memory writers, and scales memory capacity with $L=$0 (number of gradient steps). With $L=$1 vectors and $L=$2, GradMem stores up to $L=$390 key--value pairs with 88% retrieval accuracy, compared to $L=$4 keys for forward-only writers of the same size.

A break-even compute analysis reveals that, for large numbers of queries per context and substantial compression $L=$5, GradMem amortizes the WRITE cost and becomes more efficient than recomputing full-context attention. This mechanism is validated on both synthetic key--value retrieval and natural language tasks such as bAbI, SQuAD, and WikiText-103, where additional gradient steps during WRITE demonstrably increase the capacity and accuracy of the compressed memory (Kuratov et al., 14 Mar 2026).

3. Gradient Memory for Federated Learning: Catastrophic Forgetting and GradMA

In distributed federated learning, data heterogeneity (non-IID local data) and partial participation (subset selection at each round) cause the central model to "forget" contributions from less frequently sampled workers, hampering both convergence and final model utility.

The Gradient-Memory-Accelerated Federated Learning (GradMA) method introduces memory on both worker and server sides (Luo et al., 2023). Workers maintain a memory matrix consisting of prior and global gradients plus parameter differences, and use a quadratic program (QP) at each step to produce a locally corrected descent direction. The server stores per-worker memory buffers (accumulated update histories), employing a memory-reduction strategy to keep only the $L=$6 most recently active workers when $L=$7.

At the server, QP-corrected momentum aggregates the updates such that new steps remain positively correlated with stored memories, preserving knowledge from inactive or infrequently participating clients. This mechanism ensures, via explicit constraints, that the optimization is aligned with both recent and historical gradients.

Empirically, GradMA achieves marked improvements over baselines (FedAvg, FedProx, Scaffold, MOON) in both non-IID, low-participation regimes and communication efficiency. For instance, on CIFAR-10 with extreme non-IID (Dirichlet $L=$8) and $L=$9 active workers, GradMA achieves 30.8% versus $c=$021% for the nearest baseline; on Tiny-ImageNet ($c=$1, $c=$2), 23.5% accuracy versus 17% for FedAvgM and 12% for MIFA. The convergence theory (Theorem 4.1, Corollary 4.2) establishes $c=$3 convergence and linear speedup with the number of active workers (Luo et al., 2023).

4. Memory Gradient Methods in Multiobjective Optimization

The memory gradient method (GradMem/MMG) for multiobjective optimization generalizes the principle of leveraging historical gradients to problems involving multiple objectives $c=$4 (Chen et al., 2022). Here, Pareto criticality requires that there is no common descent direction for all objectives.

MMG constructs each search direction as a combination of the current descent direction and $c=$5 past directions. Formally, at iteration $c=$6,

$c=$7

where $c=$8 is the steepest descent direction, and the memory coefficients $c=$9 are set to satisfy a sufficient descent condition. Two step-size strategies (Armijo-type backtracking and a Lipschitz-based "optimal" step) are analyzed. Under mild assumptions, global convergence to Pareto critical points with an $d=$0 rate is established; in the presence of a Polyak-Łojasiewicz-type condition, linear convergence can occur.

Numerical results over 32 benchmark problems show that MMG variants (with varying $d=$1 and $d=$2) outpace both steepest-descent and conjugate-gradient baselines in iteration efficiency and quality of Pareto front coverage (Chen et al., 2022).

5. Gradient-Based Editing of Memory for Continual Learning

The Gradient-based Memory Editing (GMED) method addresses catastrophic forgetting in task-free continual learning by actively editing stored memory examples in input space (Jin et al., 2020). In the online regime, only a replay buffer of limited size is maintained, and the utility of stored examples decays as the model adapts.

GMED modifies each stored example $d=$3 via a gradient ascent step to maximize its future loss under the next update, subject to regularization:

$d=$4

with $d=$5, where $d=$6 is the post-update parameter vector. This sharpens replay examples, making them more informative and resisting forgetting.

Experiments demonstrate consistent gains over unedited replay across Split/Permuted/Rotated MNIST, Split CIFAR-10/100, and mini-ImageNet benchmarks, with statistically significant improvements in average accuracy on five out of six datasets. GMED is modular and adds little computational cost (Jin et al., 2020).

6. Hardware Implementations: In-DRAM Gradient-Based Updates (GradPIM)

The concept of “GradMem” is further embodied in hardware by the GradPIM processing-in-memory architecture (Kim et al., 2021). GradPIM enables in-DRAM gradient computations for neural network training, particularly for parameter updates, using a set of fixed-function units collocated with DRAM bank-groups. This architecture supports scalar multiplication, addition/subtraction, and quantization/dequantization entirely within DRAM, leveraging bank-group parallelism.

In practice, GradPIM accelerates parameter update phases by $d=$7–$d=$8 and reduces off-chip bandwidth and energy consumption by $d=$9–$m$0, without significant area overhead ($m$1 die area). Although current functionality is limited to simple optimizers (SGD, momentum, weight decay), this shows the practical feasibility of memory-resident gradient logic as a pathway for hardware-accelerated “GradMem” operations (Kim et al., 2021).

7. Limitations and Outlook

GradMem techniques introduce additional memory or compute costs, especially in the form of extra backpropagation steps for test-time gradient descent (e.g., $m$2-step WRITE optimization or memory gradient QPs), and necessitate careful trade-off analysis between WRITE cost and read/query efficiency (Kuratov et al., 14 Mar 2026). First-order or implicit meta-learning may mitigate some training overhead. In federated and continual learning, constrained buffer size and memory reduction strategies become crucial for scalability.

Future directions include improved WRITE objectives tailored to downstream usage, integration of memory pruning or conditional writing, and hardware support for richer classes of update rules. Across all domains, gradient-based memory mechanisms provide a flexible and generalizable tool for efficient information retention, capacity scaling, and mitigation of forgetting in both algorithmic and practical deployments.