Gradient Projection Memory (GPM)

• Gradient Projection Memory (GPM) is a framework that uses geometric projections to protect critical neural activations and prevent catastrophic forgetting.
• It computes a Core Gradient Space via truncated SVD to isolate essential basis vectors and orthogonalizes new gradients against these directions.
• Extensions like Dynamic GPM and memory-efficient methods (GaLore, Grass) enhance adaptability and reduce state storage during large-scale model training.

Gradient Projection Memory (GPM) is a family of algorithms that leverages subspace geometry in the space of neural network gradients or activations to control weight updates, thereby reducing catastrophic forgetting in continual learning and enabling memory-efficient optimization in large model training. GPM identifies, stores, and protects subspaces critical to prior knowledge—for example, by orthogonalizing new-task gradients against directions essential for past tasks. It is implemented across multiple neural learning domains, such as catastrophic forgetting in continual learning and optimizer state compression in large-scale LLM training, by geometric projection within gradient or representation spaces (Saha et al., 2021, Deng et al., 2021, Zhao et al., 2024, Muhamed et al., 2024).

1. Catastrophic Forgetting and the GPM Principle

Continual learning with neural networks presents the well-known problem of catastrophic forgetting, wherein stochastic gradient descent on new tasks overwrites parameters important for previously learned tasks. Early GPM algorithms directly address this by constructing, for each layer, a subspace—termed the Core Gradient Space (CGS)—spanned by directions in activation space deemed critical for performance on prior tasks. This is operationalized by accumulating representative activations after learning each task, performing a truncated singular value decomposition (SVD), and selecting the minimal number of singular vectors to capture a threshold fraction (ε_th) of layerwise activation variance.

Gradient projection is then enforced during subsequent training: the raw gradient in each layer is replaced with its orthogonal complement respective to the stored CGS. Thus, all weight updates are zero-projected along directions deemed sensitive by past tasks, preventing destructive interference (Saha et al., 2021).

2. Mathematical Formulation and Projection Mechanisms

Given activations $X^{(\ell)} \in \mathbb{R}^{d_\ell \times n_s}$ sampled for layer $\ell$ after task $\tau$, SVD yields $X^{(\ell)} = U^{(\ell)}\Sigma^{(\ell)}V^{(\ell)T}$. The first $k$ left singular vectors $U^{(\ell)}_{1:k}$, chosen by a variance criterion, span the CGS. For a current-task gradient $g^{(\ell)}$, the orthogonal projection is

$$g^{(\ell)}_\perp = g^{(\ell)} - P^{(\ell)}g^{(\ell)},\quad P^{(\ell)} = U^{(\ell)}_{1:k} U^{(\ell)T}_{1:k}$$

giving weight updates $\Delta W^{(\ell)} \propto -\eta\,g^{(\ell)}_\perp$.

Across continual learning sequences, new task activations are projected out of the accumulated CGS before SVD, ensuring orthogonality and numerical stability. Memory cost is $O(d_\ell k_\ell)$ per layer and projection cost is $O(d_\ell k_\ell)$ per step. This is substantially lower than storing raw gradients or episodic exemplars, and subspace growth is controlled by the CGS approximation threshold (Saha et al., 2021).

3. Extensions: Dynamic and Sharpness-Regularized GPM

Dynamic GPM (DGPM) introduces a soft importance weighting $$\Lambda = \operatorname{diag}(\lambda_1,\ldots,\lambda_k)$$, allowing learned, adaptive continuous control over the extent to which each basis vector is “locked” (protected) or “released” (modifiable). Gradients and the soft-importance vector $\lambda$ are simultaneously updated. Bases with $\lambda_i \to 0$ are dropped, improving network plasticity without excessive forgetting (Deng et al., 2021).

Additionally, Flattening Sharpness (FS) regularization is integrated via inner maximization over the CGS, quantifying and minimizing the maximal loss increase caused by perturbations in the protected subspace. The training objective becomes

$$\min_{w,\Lambda} \max_{v\in V(\Lambda)} L_{D_t \cup \mathcal{M}}(w + v)$$

where $V(\Lambda)$ is the dynamic subspace and $\mathcal{M}$ denotes the replay buffer. The algorithm alternates between ascent in $v$, descent in $\Lambda$, and projected descent in $w$. This min–max approach further tightens expected risk bounds and yields higher empirical accuracy and stability on multi-task benchmarks (Deng et al., 2021).

4. Gradient Projection Memory in Large-Scale Model Training

In LLM training, the memory bottleneck arises from optimizer states and gradients, particularly for high-dimensional matrices. Gradient projection-based memory-efficient subspace optimization (MeSO) methods—such as GaLore and Grass—employ gradient subspace projection to reduce optimizer state storage from $O(mn)$ to $O(rn)$ for $W \in \mathbb{R}^{m\times n}$ and $r\ll m$. GaLore constructs dense low-rank projectors $(P, Q)$ by SVD of the full gradient, applies and tracks optimizer state in the projected space, and periodically updates the subspace (Zhao et al., 2024). Grass instead employs structured sparse projections—each projection matrix column selects only one row—enabling selective compression, high throughput, and reduced memory without forming dense full gradients (Muhamed et al., 2024).

Method	Projection Type	Memory Complexity
GPM, DGPM	Dense, orthonormal	$O(d_\ell k_\ell)$
GaLore	Dense, low-rank (via SVD)	$O(mr + rn)$ per matrix
Grass	Structured sparse	$O(rn)$ per matrix

In all cases, gradient updates are constructed either by projecting into the compact basis or by sparse selection, then applying the optimizer (Adam, Adafactor, 8-bit Adam) on the compressed states before reconstructing a sparse or low-rank incremental update to model weights (Zhao et al., 2024, Muhamed et al., 2024).

5. Empirical Performance and Application Domains

On continual learning benchmarks (Permuted MNIST, Split CIFAR-100, miniImageNet, multi-dataset sequences), GPM achieves superior or on-par accuracy with near-zero backward transfer (BWT ≈ 0), outperforming gradient episodic and regularization-based baselines, and matching or slightly surpassing methods such as HAT. For instance, in Split CIFAR-100, GPM achieves higher average accuracy and reduces memory footprint by 40–80% relative to replay or gradient-store baselines; online continual learning with DGPM and FS-DGPM yields additional improvements in both plasticity and stability (Saha et al., 2021, Deng et al., 2021).

In large model optimization, GaLore achieves full-rank performance—for instance, LLaMA 1B models with perplexity within 0.1 of the baseline—while reducing optimizer state memory by up to 65.5%. Grass enables 13B-parameter LLaMA pretraining on a single 40GB A100 GPU using only $O(rn)$ memory, with a 2× throughput improvement in distributed settings. Finetuning on GLUE and instruction-tuning tasks further demonstrates competitive accuracy with substantially lower memory and communication costs (Zhao et al., 2024, Muhamed et al., 2024).

6. Limitations, Tradeoffs, and Future Directions

GPM’s plasticity-stability trade-off is governed by the variance threshold (ε_th) or the soft-importance vector (Λ), determining the amount of past-task protection versus new-task learnability. For tasks with highly dissimilar or diverse representations, subspace dimensionality can grow to saturate network capacity, limiting the approach’s applicability to domains with shared representational structure (Saha et al., 2021, Deng et al., 2021).

In LLM training, dense projection variants incur SVD or sampling costs, while sparse methods require custom backward passes and careful coordination for distributed training (Zhao et al., 2024, Muhamed et al., 2024). Projector update frequency, subspace rank $r$, and structured selection strategies mediate the tradeoff between memory savings, computational cost, and convergence behavior.

Ongoing research targets online/incremental SVD, per-layer adaptive subspace control, integration with replay/generative memory, advanced sampling for projection selection, and extending theory to architectural variants such as self-attention. The geometric insight underlying GPM also motivates cross-domain hybrids blending subspace projection and gradient memory with other forms of episodic or generative rehearsal (Saha et al., 2021, Deng et al., 2021).

7. Connections and Impact Across Research Areas

The GPM framework unifies gradient-based continual learning and large-scale, memory-efficient optimization under a common geometric subspace perspective. It clarifies links among methods that project or regularize along task-relevant subspaces (e.g., GEM, A-GEM) and those using low-rank, sparse, or structured projection for efficiency or forgetting prevention. Its success in both catastrophic forgetting avoidance and LLM training efficiency highlights the broad relevance of subspace-based memory strategies for scaling and stabilizing future neural models (Saha et al., 2021, Deng et al., 2021, Zhao et al., 2024, Muhamed et al., 2024).