Selective Memorization in ML (SeMem)

• Selective Memorization (SeMem) is an approach that enforces explicit boundaries between memorized and generalized content in machine learning models.
• Techniques like MemSinks, model fusion, and gated subspace steering enable selective retention, reducing memory overhead and safeguarding sensitive data.
• Practical implementations in transformers, diffusion models, and state-space systems demonstrate substantial memory savings and controlled utility loss, enhancing privacy and scalability.

Selective Memorization (SeMem) is an umbrella term referring to neural or algorithmic mechanisms that control, constrain, or leverage the specific subset of information retained, isolated, or reproduced by a machine learning model, in contrast to indiscriminate or uniform memorization. SeMem arises in diverse architectures—including transformers, state space models, diffusion models, and semiparametric retrieval-augmented LMs—whenever one aims to enforce explicit boundaries between memorized and generalized content, allocate memory selectively according to salience or privacy needs, or surgically forget dangerous information while retaining utility. Distinct instantiations of SeMem now underpin competitive approaches for privacy-preserving training, scalable learning under resource budgets, and targeted unlearning.

1. Formal Characterizations of Selective Memorization

The concept of Selective Memorization appears in multiple formalizations depending on the model and problem domain:

1. Neuron-level isolation (MemSinks): In transformer models, SeMem can mean explicit partitioning of neuron populations for memorization and generalization. The MemSinks approach splits each MLP into always-on generalization neurons (fraction $g$) and sequence-specific "sink" neurons (fraction $1-g$), activated via a deterministic but sequence-dependent mask (Ghosal et al., 14 Jul 2025). Each repeated training sequence is routed to a disjoint subset of sink neurons, allowing direct control over where and how memorization occurs.
2. Shared vs. unshared memory in model fusion: The "Fuse to Forget" paradigm studies SeMem via weight-averaged fusion of models fine-tuned on non-identical data. The fusion preserves knowledge (measured by utilization score $\Psi_{D, T}(\theta)$) present across models but "forgets" content unique to only one constituent model, a direct manifestation of SeMem (Zaman et al., 2023).
3. Token- and subspace-selectivity: In both LLMs and diffusion models, not all tokens, features, or subspaces are equally likely to be memorized. High-variance directions in generative diffusion collapse to memorization earlier than low-variance ones, leading to a staged, selective loss of data manifold dimensions (Achilli et al., 2024); auto-encoder probes reveal that SSM LMs like Mamba differentially forget numerical and infrequent tokens as a function of sequence length and pretraining statistics (Hossain et al., 17 Dec 2025).
4. Budgeted or policy-driven selective memory: Memory-augmented LMs may use gating or salience-scoring to select only a subset of data for external storage, ensuring that memory costs grow sublinearly with data or model size (Alla et al., 7 Nov 2025, Peng et al., 2023).
5. Selective unlearning and mitigation: Methods such as Gated Subspace Steering (GSS) implement token-level, context-aware interventions that probe for memorization-relevant modes and remove them at inference only where needed, ensuring minimal collateral degradation (Zhang et al., 9 Feb 2026). Gradient projection methods in diffusion enforce SeMem by zeroing gradients aligned with prohibited features (Kothandaraman et al., 12 Dec 2025).

2. Mechanistic Implementations and Algorithms

A core ingredient in SeMem frameworks is architectural or algorithmic separation—across neurons, parameter subspaces, tokens, or policy gates—between what is retained, what is forgotten, and what remains for generalization:

• MemSinks (Ghosal et al., 14 Jul 2025): Each document/sequence is hashed to generate a unique mask $m_i \in \{0,1\}^h$ for the sink neurons in each MLP block. During forward propagation, only the designated subset for that sequence is active; all others are dropped. No explicit penalties are placed on parameter separation; all isolation is effected via the mask at both forward and backward passes. At inference, setting all sink neuron activations to zero erases memorized content with negligible loss to generalization.
• Model Fusion (Zaman et al., 2023): For $M$ models with parameter vectors $\theta_1,\ldots,\theta_M$, the convex weighted sum $$\theta_{\mathrm{fused}} = \sum_{i=1}^M \alpha_i \theta_i$$ produces SeMem, preserving knowledge present in all (shared memorization) while attenuating that which is idiosyncratic. The measure $\Psi_{D,T}(\theta)$ formalizes knowledge utilization; after fusion, $\Psi$ on "private" data (unique sequences) is suppressed.
• Gated Subspace Steering (GSS) (Zhang et al., 9 Feb 2026): Here, memorization is empirically sparse and token-conditioned. GSS learns a "memorization subspace" $S_m$ (via SVD of a memorization-relevance matrix), defines a probe $1-g$0, intervenes with a steer $1-g$1 only when $1-g$2 exceeds a per-coordinate threshold, and guarantees the lowest possible correction rank consistent with safety constraints.
• BudgetMem (Alla et al., 7 Nov 2025): Chunks of context are scored by a shallow net using entity density, TF-IDF, discourse marker score, position bias, and numeric density. Only the highest-scoring $1-g$3 chunks (under a hard memory budget) are stored, enabling practical long-context LLM deployment with minimal F1 loss.
• Gradient Projection in Diffusion (Kothandaraman et al., 12 Dec 2025): All updates are projected onto the orthogonal complement of forbidden concept gradients, blocking the first-order learning signal for prohibited features while allowing free learning in unrestricted subspaces.

3. Empirical Findings and Measured Trade-offs

SeMem approaches are empirically validated using detailed ablations and quantitative metrics, illustrating the precision-versus-memory (or precision-versus-forgetting) trade-offs:

Framework	Domain	Main SeMem Mechanism	Memorization Attenuation	Utility Cost
MemSinks	Transformer LMs	Sequence-tied sink neuron mask	>50% gap closed	<0.02 nats loss
Model Fusion	BERT fine-tuning	Convex parameter averaging	Private sets forgotten	Generalization ↑
GSS	LLMs (Pythia, GSM8K)	Token-gated rank-K projection	Greedy memorize 52.9→6.9%	PPL +2–7; math ~id.
BudgetMem	Long-context QA	Feature-gated memory writes	72.4% mem. saved (long)	F1 −1% (long docs)
Gradient Projection	Diffusion models	Orthogonalization of gradients	SSCD −47–64% (adversary)	CLIP ∆<0.2%

Empirical analyses highlight that selective interventions can sharply reduce memorization risks (e.g., privacy leakage, copyright violations) while minimizing degradation of core utility metrics (PPL, F1, generalization accuracy). In fusion, shared knowledge is actually amplified, demonstrating a counterintuitive "SeMem boosting" effect for intersectional content (Zaman et al., 2023).

4. Theoretical Analyses and Guarantees

Key theoretical advances under SeMem regimes include:

• Gradient orthogonality guarantee: In gradient-projected diffusion models, all update steps are strictly orthogonal to prohibited-concept directions, making it impossible for the forbidden subspace to acquire weight by first-order learning (Kothandaraman et al., 12 Dec 2025).
• MemSinks minimum-norm bias: Standard gradient descent tends to entangle memorization and generalization, but MemSinks’ per-sequence sink neuron isolation shifts gradient flow, minimizing cross-interference and supporting efficient post-hoc erasure (Ghosal et al., 14 Jul 2025).
• Selective collapse in diffusion models: Theoretical results using random energy models and Jacobian spectral gap analyses demonstrate that, given a data manifold with subspaces of varying variance, the high-variance directions overfit and "collapse" first; memorization thus proceeds by staged freezing of geometric features (Achilli et al., 2024).
• Capacity and scaling bounds: In semiparametric continual learning, memory size grows sublinearly with continued training and model scale; larger, better parametric models require fewer retained memories for the same accuracy (Peng et al., 2023).

5. Practical Implications: Privacy, Scalability, and Evaluation

Key practical outcomes of SeMem frameworks include:

• Privacy mitigation: SeMem enables controlled forgetting without catastrophic loss. For transformer LMs, MemSinks removes the ability to reconstruct sensitive repeated documents by zeroing a neuron subset, protecting privacy even at scale (Ghosal et al., 14 Jul 2025). In model fusion, privacy-sensitive knowledge unique to a single dataset is erased (Zaman et al., 2023).
• Long-context resource efficiency: BudgetMem and related gating approaches reduce the footprint of memory-augmented LMs by 60–72% for long documents at only 1–3% loss in F1, democratizing access to powerful models on modest hardware (Alla et al., 7 Nov 2025).
• Cautions in benchmarking and inference: The presence of SeMem (or "selective perfect memory") in LLMs challenges the validity of any claimed forecasting or backtesting performance for in-sample (pre-training) periods; even explicit instruction or masking fails to fully block memorized recall (Lopez-Lira et al., 20 Apr 2025). Only genuine out-of-sample periods should be used for evaluation.
• Selective retention biases in fixed-state LMs: SSMs exhibit content-dependent forgetting, disproportionately losing rare, numerical, or domain-specific tokens as context length grows or under-represented in pretraining (Hossain et al., 17 Dec 2025).

6. Limitations and Open Problems

While SeMem frameworks offer tangible advances, open questions remain:

• Scaling limits: Current SeMem methods are tested at O(10^9–10^10) parameter and token scales; extension to larger, multi-lingual, or highly entangled corpora remains to be proven (Ghosal et al., 14 Jul 2025, Alla et al., 7 Nov 2025).
• Robustness to adaptive attacks: Most demonstrated SeMem defenses operate under benign or known threat models. Resistance to adversarial extraction, paraphrase, and attribute entanglement is an area of future work (Kothandaraman et al., 12 Dec 2025).
• Mechanistic co-adaptation: In sink-based architectures, over-activation or under-activation of memory neurons (extreme p/g choices) can lead to capacity underuse or excessive cross-interference, demanding more principled policy optimization (Ghosal et al., 14 Jul 2025).
• Dynamic/adaptive policies: Hard-coded thresholds, static gating, and linear salience models are prevalent; more expressive, adaptive, and context-aware SeMem policies may be needed for real-world scaling (Alla et al., 7 Nov 2025, Peng et al., 2023).
• Fine-grained attribute isolation: While per-sequence, per-chunk, or per-feature isolation is possible in principle, optimal partitioning and mask design is an open area, particularly for semantically entangled concepts (Ghosal et al., 14 Jul 2025, Kothandaraman et al., 12 Dec 2025).

7. Cross-Domain Extensions and Future Directions

Recent advances demonstrate that SeMem is a unifying paradigm across language, vision, and multi-modal domains:

• VLN and multi-modal models: State-space and cross-modal SeMem modules can replace attention machinery in navigation systems, yielding substantial computational savings with no loss in accuracy (Zhang et al., 31 Mar 2025).
• Diffusion models and geometry: By linking SeMem to manifold geometry and spectral theory, new approaches for monitoring, regularizing, and intervening in the memorization trajectory of generative models become evident (Achilli et al., 2024).
• Semiparametric continual learning: Selective memory policies achieve sublinear memory growth and stable adaptation, creating scalable, low-forgetting learners suitable for continual and streaming applications (Peng et al., 2023).

Overall, the SeMem paradigm—realized through neuron isolation, model fusion, feature gating, selective subspace projection, or policy-driven external memory allocation—marks a shift toward more precise, parsimonious, and privacy-aware machine learning. It offers both practical tools and theoretical frameworks for rigorously controlling the boundary between knowledge retention, memorization, and forgetting across heterogeneous learning systems.