Technical Summary: AI Engrams: In Search of Memory Traces in Artificial Intelligence

1. Problem Statement

Deep neural networks store knowledge in a deeply entangled manner, where individual parameters participate in multiple behaviors simultaneously. This distributed and entangled nature of learned representations makes it difficult to perform "surgical" knowledge manipulation, such as removing specific concepts (machine unlearning) or isolating specific memories without retraining the entire model. While neuroscience posits that memories are stored in identifiable physical substrates called "engrams," it remains an open question whether analogous, identifiable memory traces exist within artificial neural networks. Current methods for knowledge editing or unlearning often rely on iterative gradient-based optimization or heuristic attribution, which can be computationally expensive, lack theoretical guarantees of specificity, or fail to scale to large models.

2. Methodology

The paper introduces a geometric framework to identify "AI engrams"—causal memory traces in artificial intelligence—by formalizing neuroscientific criteria into a constrained inverse problem.

2.1. Neuroscientific Criteria as Constraints

The authors define the identification of an engram $\mu^+_c$ for a target concept $c$ based on four canonical criteria (Josselyn & Tonegawa, 2020; Tomé et al., 2024):

1. Specificity: The engram must selectively encode the target experience ($X^+$) while remaining inert to reference experiences ($X^-$).
2. Reactivation: The engram must consistently reproduce the internal representations ($Z^+$) associated with the target.
3. Sufficiency: Injecting the engram into a naive (null) state must induce the learned behavior.
4. Necessity: Ablating the engram from a learned state must eliminate the target behavior while preserving unrelated knowledge.

2.2. Algebraic Formulation

The authors formulate the search for the synaptic engram $W^+$ as an optimization problem. They contrast observed states (natural learning) with intervened states (algebraic manipulation) to derive linear constraints.

• Objective: Minimize the discrepancy between intervened and observed states across target and reference sets.
• Constraints: Enforce orthogonality to the reference subspace ($W^+ X^- = 0$) to ensure specificity.
• Solution: The authors derive a closed-form estimator for the engram. By utilizing the Moore-Penrose pseudoinverse and uncentered covariance matrices ($\Sigma^+ = X^+ X^{+\top}$ and $\Sigma^- = X^- X^{-\top}$), they obtain:
  $$W^+ = \Delta W \Sigma^+ (\Sigma^+ + \Sigma^-)^\dagger$$
  In practice, for a trained model with weights $W$, the update is approximated retrospectively as $W^+ = W \Sigma^+ (\Sigma^+ + \Sigma^-)^\dagger$. This allows for the extraction of engrams in a single forward pass without backpropagation.

2.3. Geometric Interpretation

The paper establishes a link between these biological constraints and information geometry. Under the Kronecker-Factored Approximate Curvature (K-FAC) assumption and the Principle of Maximum Entropy, the derived closed-form solution corresponds to the minimum-norm projection onto a constraint manifold under the Fisher Information Metric. This reveals that the engram direction aligns with the natural gradient direction on the parameter manifold.

2.4. Combinatorial Unlearning

Because the identified engrams are linearly additive, the framework enables combinatorial unlearning. To remove a subset of concepts, one can simply subtract the corresponding engrams from the model weights using linear arithmetic:
$$\text{Engram}(\alpha) := W - \alpha \sum_{k \in U} W^+_k$$
This bypasses the exponential complexity of iterative unlearning methods.

3. Key Contributions

1. A Framework for AI Engrams: The authors translate biological memory criteria into algebraic constraints, yielding a scalable, closed-form estimator that isolates memory traces from globally entangled parameters.
2. Amortized Memory Isolation: The method decouples memory traces, enabling the instantaneous composition or erasure of knowledge via simple linear arithmetic. This resolves the combinatorial complexity of unlearning and demonstrates robustness across architectures from simple MLPs to billion-parameter LLMs.
3. Link to Information Geometry: The work demonstrates that the biologically-inspired closed-form estimator recovers the Fisher curvature, coinciding with the natural gradient direction, thereby bridging biological memory theories and artificial representation learning.

4. Experimental Results

The authors validate the framework across diverse architectures and datasets:

• Scalability and Versatility: The method successfully isolates and erases concepts in MLPs, CNNs (ResNet-18/50), Vision Transformers (ViT), and generative models (ConvAE, WAE). In CIFAR-10 and CIFAR-100 experiments, ablating a specific engram selectively removes the target class accuracy while preserving performance on reference classes (surgical precision).
• Quantitative Performance: On class-wise unlearning benchmarks (CIFAR-10/100), the Engram method achieves superior Tug-of-War (ToW) scores compared to state-of-the-art baselines (e.g., NegGrad, SalUn, Fine-tuning), indicating a closer behavioral alignment to a model retrained from scratch. It also shows effective dissipation of structural traces in latent space (DA and NMI metrics).
• Large Language Models (LLMs): Applied to Llama-3.2-1B on the TOFU benchmark, the method demonstrates high efficiency.
  • Efficiency: It requires roughly 30x fewer FLOPs and 2.3x less peak memory than gradient-based iterative methods, completing in ~2 minutes versus tens of minutes.
  • Performance: While it trails some dedicated iterative methods (e.g., AltPO, NPO) on overall metrics, it significantly outperforms other closed-form editing baselines (UCE, Task Arithmetic) and achieves the lowest Exact Memorization (EM) scores among gradient-free methods.
  • Localization: Analysis of the engram weights reveals that memory traces are not randomly distributed but concentrated in specific layers: the Query (Q) and Key (K) projections of self-attention and the Gate projections of MLPs. Ablation studies confirm that targeting these specific layers is sufficient for effective unlearning.

5. Significance and Claims

The paper claims to bridge the gap between biological memory theories and artificial intelligence by demonstrating that causal substrates of knowledge can be isolated and manipulated through deterministic spectral resolutions, even within highly non-linear systems.

• Paradigm Shift: The work proposes shifting the paradigm of machine unlearning from stochastic gradient search to structural identification.
• Orthogonal Regime: The authors position Engram not as a replacement for iterative unlearning but as an orthogonal regime (O(1)-pass, gradient-free) that is particularly valuable when:
  1. The number of concepts to unlearn is large (leveraging combinatorial compositionality).
  2. Compute resources or weight access are restricted.
  3. The goal is structural identification rather than distributional indistinguishability.
• Interpretability: The discovery of condensed localization patterns in LLMs provides a new perspective on neural interpretability, suggesting that specific functional topologies of memory reside in linearizable subspaces of the weight space.

The authors conclude that their framework offers unprecedented efficiency for surgical memory isolation, enabling safer model behavior and clearer interpretability, while opening avenues for future research into the temporal dynamics of engrams in continual learning.