AI Engram: In Search of Memory Traces in Artificial Intelligence

Summary

The paper introduces "AI engrams": a closed-form estimator for memory traces in trained neural networks obtained by translating four neuroscientific engram criteria (specificity, reactivation, sufficiency, necessity) into an algebraic constrained-optimisation problem. The KKT solution yields a strikingly simple per-layer operator, $W^+ = W\,\Sigma^+(\Sigma^+ + \Sigma^-)^\dagger$, computable in a single forward pass over target and reference statistics. A parallel derivation shows that, under a K-FAC block-diagonal Fisher approximation and an isotropic-gradient maximum-entropy prior, the same operator emerges as a natural-gradient step, reframing engram ablation as covariance-weighted forgetting rather than Mahalanobis normalisation. Empirically, the estimator supports class-level unlearning across MLP/ResNet/ViT/ConvAE, semantic attribute control on CelebA-WAE, and — with an adaptive per-layer scale from the engram-to-weight norm ratio — the highest Overall score among closed-form editing methods on TOFU forget10 with Llama-3.2-1B.

The conceptual contribution is genuine: a compact, single-pass, gradient-free procedure for concept-targeted weight editing that composes linearly and admits a clean geometric reading. Positioning it as a distinct point on the compute–accuracy frontier — dominant over other closed-form editors, still trailing dedicated iterative unlearners on Overall/Privacy — is honest and useful. The main conceptual limitation is a persistent slippage between the paper's hard-constraint framing ($W^+ X^- = 0$, "orthogonal task vectors", "zero interference with reference memories") and the soft, spectrally-weighted operator the estimator actually produces on rank-deficient covariances. That distinction is admitted once in Section 6.3 and again in Appendix H's "functional identification, not learning-trajectory reconstruction" caveat, but claims elsewhere — "linearizable subspace of the weight space", "spectral resolution of the entire network state", "exact reconstruction of causal traces" — go beyond what the estimator delivers. Methodologically, the tabula-rasa instantiation $\Delta W \approx W$ silently replaces the object of interest, evaluation is single-target on each benchmark, and the code release omits several artefacts (baseline unlearners, retrained references, CKA, ConvAE, naive-delta ablation) that would let a reader independently audit the head-to-head tables.

Strengths

• Conceptual contribution. The translation of the four neuroscientific engram criteria into KKT constraints is a clean and thought-provoking framing, and the resulting operator $P^+ = \Sigma^+(\Sigma^+ + \Sigma^-)^\dagger$ has both an intuitive interpretation (soft covariance-weighted projection) and useful compositional structure across concepts.
• Fisher / natural-gradient equivalence. Theorem 6.1 and Corollary 6.3 tie the biologically motivated estimator to a well-established geometric object under stated approximations. The reading of covariance-weighted (rather than Mahalanobis-inverse) forgetting is a genuinely useful lens.
• Scale of empirical scope. Discriminative (ResNet/ViT on CIFAR/ImageNet), generative (ConvAE/WAE on MNIST/CelebA), and 1B-parameter LLM (TOFU forget10) results with a single estimator template are impressive coverage, and the W-Norm localisation heatmap for Llama-3.2-1B is a genuinely informative interpretability artefact.
• Compute–accuracy positioning. Framing Engram as a distinct operating point on the compute–accuracy frontier — dominant over other closed-form editors, complementary to iterative unlearners — is honest and helpful to practitioners choosing between paradigms.
• Self-critique. Appendix G's structural analysis of tabula rasa vs. delta-weight instantiations and Appendix H's "functional identification, not learning-trajectory reconstruction" caveat show unusual awareness of the interpretive weight the framework can bear.

Weaknesses

• Hard-constraint framing vs. soft-projection reality. Section 4.3, Appendix A.2, and Section 5.3 repeatedly speak of "strict specificity", "zero interference with reference memories", and "orthogonal task vectors...without interference". The proof sketch of Proposition 4.1 and Section 6.3 admit that $P^+$ is a soft filter, not a hard projector, and Appendix C.4 shows measurable within-superclass interference. The paper would be substantially clearer if the soft-projection caveat were promoted from an aside to the definition itself.
• "Linearizable subspace of the weight space" overclaims what has been shown. What the derivation establishes is that the extraction operator is linear per layer; the claim that memory itself resides in a linearizable subspace is a hypothesis consistent with, but not established by, the qualitative arithmetic experiments. Similarly, "spectral resolution of the entire network state" (Appendix D) and "exact reconstruction of causal traces" are stronger than the actual retrospective identification of a functional component.
• Tabula rasa is a definitional shift, not a numerical convenience. The KKT solution is derived for $\Delta W = W_1 - W_0$, and Eq. (4) then substitutes $\Delta W \approx W$. Appendix G shows this is a structural, not numerical, choice; Appendix H acknowledges that the resulting object is a functional component of the converged manifold, not the concept-induced weight perturbation. That reframing deserves to be signposted at first use of Eq. (4).
• Maximum-entropy prior is load-bearing. Theorem 6.1's Fisher–Engram equivalence requires $G_l \approx \sigma^2 I$. The paper defends this as "least informative given no gradient information", but the entire input-covariance-only estimator matches the Fisher metric only because the output-side geometry has been assumed away. Please either show empirically that $G_l$ is close to isotropic on the tested networks or reframe the theorem as an equivalence to an input-only Fisher surrogate.
• Statistical rigour of head-to-head tables. Table 2 reports means over five seeds but no standard deviations; on CIFAR-100 the Engram advantage over SalUn on ToW is 0.003, comparable to the retrain-to-retrain noise floor. Hyperparameter selection maximises the ToW metric that is then reported, which biases every method's ToW upward and makes near-ties hard to interpret.
• Evaluation scope is narrow at the head-to-head level. All CIFAR/TOFU head-to-head tables target a single class or the single forget10 split. The class-wise heatmaps sweep all classes, but the quantitative comparisons that carry the main claims never aggregate across targets, forget-set sizes, or LLM scales beyond Llama-3.2-1B.
• Adaptive $\alpha_{\text{W-Norm}}$ carries most of the LLM gain, unattributed. The Overall score on TOFU rises from 0.698 (uniform $\alpha=0.6$) to 0.818 with W-Norm scaling — a large fraction of the paper's LLM headline. Since W-Norm is derived from the same weights being edited, an alternative-schedule ablation (e.g. validation-tuned per-layer $\alpha$, rank-scaled $\alpha$) is needed to attribute the gain to "localisation" vs. "a well-tuned knob".
• Normalisation / bias handling glossed. ResNet BatchNorm, ViT LayerNorm, and Llama RMSNorm interact with the input covariance in ways that Eq. (4) does not obviously respect (the footnote about bias absorption via input augmentation does not cover normalisers). It would help to describe explicitly how BN running statistics and normaliser gains are handled during covariance collection and application.
• Rank of $\Sigma^+$ at LLM scale. The $N \gg d$ justification for the $O(d^2)$ memory footprint holds for the retain set, but for the forget subset $\Sigma^+$ can be highly rank-deficient, so the "soft filter" behaviour is dominated by its null-space. A rank analysis on TOFU would help characterise when the estimator degrades gracefully vs. when it silently truncates.

Reproducibility & code

The release at https://github.com/jeakwon/ai-engram covers the Engram estimator itself and enough plumbing to reproduce several qualitative figures, but multiple head-to-head comparisons and figures cannot be regenerated without independent implementation effort. All findings below come from static inspection of the code (paper–code mode); nothing was executed.

• TOFU headline pipeline. tests/test_tofu_evaluate.py + tests/_tofu_evaluate.py reproduce the Engram row of Table 3/4 end-to-end on Llama-3.2-1B, and src/engram/scaling.py:weight_norm implements $\alpha_{\text{W-Norm}}$. Both Engram variants (uniform and adaptive) are reachable, and the OpenUnlearning 14-metric harness is bundled. This part of the release is solid.
• Baselines absent from Table 2. Fine-tune, NegGrad, NegGrad+, $l_1$-Sparse, RandomLabel, and SalUn are not implemented in src/engram, and no retrained-from-scratch reference training script is shipped. Reproducing Table 2 requires re-implementing every baseline and re-training references (including a ResNet-50 CIFAR-100 checkpoint not present in the repository — see next point).
• CIFAR-100 architecture mismatch. Table 2's CIFAR-100 row is labelled ResNet-50, but Appendix B.2 cites ResNet-18 checkpoints and the shipped fig_resnet18_cifar100.ipynb uses edadaltocg/resnet18_cifar100. No ResNet-50 loader or checkpoint is included; this row is currently irreproducible from the release.
• Closed-form baselines (Task Arithmetic, UCE). Neither is implemented in the repository, though Appendix E.5 describes them algebraically and lists sweeps. Table 4's positioning claim rests on rows a reader must code from scratch.
• ConvAE MNIST (Fig. 5). Only fig_mlp_mnist.ipynb (a fully-connected MLP) is shipped. No ConvAE architecture, training script, or reconstruction-MSE evaluation code exists in the release, so the "extends to unsupervised generative regime" evidence is currently unauditable.
• CKA plots (Figs. 6 and 10). No CKA implementation, no baseline unlearned checkpoints, and no retrained references are shipped. The "closest to the retrained model" claim cannot be regenerated as released.
• Naive-delta ablation (Table 9). compute_engram_weights implements only the tabula-rasa instantiation; the naive $\Delta W = W_{\text{ft}} - W_{\text{pt}}$ variant used to justify Appendix G's structural argument has no reproducer script.
• Layer-type ablation (Table 8). target_modules supports the ablation mechanically, but no shipped script parameterises the TOFU pipeline over the four layer-subset conditions.
• Graceful-degradation reduction (Fig. 14). The raw 100×100 accuracy matrix is producible from fig_resnet18_cifar100.ipynb, but the superclass grouping, cosine-similarity binning, and two-panel plotting logic are not in the release. The 0.80 pp headline number requires the reader to guess an averaging protocol.
• Compute–accuracy claims (Table 7, "$\leq$ 2 minutes on A100"). Table 7's FLOP and memory figures are analytical, not measured. No profiling harness or wall-clock instrumentation is shipped, so the 30x and 2.3x multipliers and the "2 minutes" claim are unverifiable without adding instrumentation.
• ImageNet-1k Fig. 12. The fig_vit_imagenet1k.ipynb notebook exists but data access is HF-gated (~150 GB) and covariance collection alone is ~6 h; the 20-class grouping strategy is not fully scripted per the veritas assessment.

Recommended Changes

Essential

• Reconcile hard-constraint text with the soft-projection operator. Rewrite Proposition 4.1, Section 4.3, Section 5.3, and Appendix A.2 so the "strict specificity" and "orthogonal task vectors...without interference" language is replaced by the soft, spectrally-weighted characterisation. Cross-reference Appendix C.4's measured 0.80 pp within-superclass interference so readers see the operator's soft behaviour up-front, not only in Section 6.3.
• Signpost the tabula-rasa reframing at first use. Rewrite Section 4.5 to state explicitly that Eq. (4) identifies a functional component of the converged manifold rather than a reconstruction of the training-induced perturbation, promoting the Appendix H caveat to a Section 4 remark.
• Soften "linearizable subspace" / "exact reconstruction" language. Restrict "linearity" language in Section 4.5, Appendix D, and the abstract to properties of the estimator; frame the memory-topology hypothesis as tested (partially) by the arithmetic experiments rather than as an established theoretical result.
• Restate Theorem 6.1 and Corollary 6.3 with visible knobs. Reframe Theorem 6.1 as an equivalence to an input-only Fisher surrogate (or add empirical evidence that $G_l$ is close to isotropic on the tested networks). Rewrite Corollary 6.3 so the implicit $1/(2\sigma^2)$ step size is visible instead of "unit-step".
• Report per-seed variance and validate hyperparameter selection. Add standard deviations across the five seeds to Table 2 (and to Tables 3/4/8/9 for TOFU, adding at least a small seed set for the LLM runs). Move $\alpha$ selection to a held-out validation split so ToW is not both the selection and reporting criterion.
• Release the missing artefacts needed to audit head-to-head tables. Ship (i) baseline unlearners for Table 2 (or a documented wrapper over OpenUnlearning), (ii) the ResNet-50 CIFAR-100 checkpoint and loader, (iii) a Task Arithmetic and UCE implementation for Table 4, (iv) a ConvAE MNIST training + evaluation notebook for Fig. 5, (v) a CKA harness plus the underlying similarity matrices for Figs. 6 and 10, and (vi) a naive-$\Delta W$ code path for Table 9.

Suggested

• Ablate $\alpha_{\text{W-Norm}}$ against alternative schedules. Add at least one alternative per-layer $\alpha$ (e.g. rank-scaled or validation-tuned) so the LLM gain from W-Norm is attributable to localisation rather than to any well-tuned per-layer knob.
• Broaden LLM evaluation. Report TOFU forget01 and forget05, and if possible one additional target model (e.g. Llama-3-8B via Q/K + Gate-only extraction as suggested by Appendix E.6). At CIFAR scale, aggregate the head-to-head tables over multiple targets rather than only Class 0.
• Add a multi-concept compositional test. Demonstrate the $2^n - 1$ compositionality claim with $n \geq 3$ concepts on CIFAR-10 by comparing composite ablations to retrained-without-subset references, not only two-attribute WAE arithmetic.
• Add a rank / sample-size analysis for TOFU. Report effective ranks of $\Sigma^+$ and $\Sigma^-$ at each Llama layer and characterise how the estimator degrades when $|X^+| \ll d$.
• Add profiling harness and timing hooks. Add a lightweight profiler for Table 7 and time.time() instrumentation for the "$\leq$ 2 minutes" claim; alternatively, clearly relabel Table 7 as analytical.
• Add reduction scripts for Fig. 14 and Table 8. Ship the superclass grouping / cosine-similarity binning pipeline for Fig. 14 and a driver script that reproduces Table 8's four rows from one command.
• Discuss normalisation handling. Add a short paragraph on how BatchNorm/LayerNorm/RMSNorm running statistics and gains are handled during covariance collection and application.
• Fix the MNIST caption typo. "The remaining 99 classes" should read "the remaining 9 classes" in the Fig. 11 caption.
• Temper the CIFAR-100 headline. Rewrite "outperforming all baselines" to reflect that on CIFAR-100 the ToW lead over SalUn is 0.003 and within plausible seed noise.