Tell Me What To Learn: Generalizing Neural Memory to be Controllable in Natural Language

📝 Paper Summary

Neural Memory Continual Learning

Generalized Neural Memory (GNM) enables users to explicitly control what a model learns, ignores, or forgets from new documents by conditioning memory updates on natural language instructions.

Core Problem

Existing neural memory systems optimize a fixed objective (usually next-token prediction) on all data, preventing users from specifying which parts of a document should be remembered, ignored, or treated as behavioral updates.

Why it matters:

Real-world data is heterogeneous; a single document may contain useful behavioral heuristics (e.g., escalation protocols) alongside outdated facts or sensitive data that must be ignored
Current methods like RAG or ICL are either imprecise or computationally expensive, while standard fine-tuning suffers from catastrophic forgetting and cannot selectively filter information within a training document
Users lack a mechanism to align the model's long-term memory updates with specific downstream intents, such as adopting a tone but rejecting the associated factual content

Concrete Example: A medical agent reading nurse-patient transcripts should learn the 'escalation heuristics' (when to call a doctor) but explicitly ignore the 'outdated dosing protocols' mentioned in the same text. Current models would learn both indiscriminately.

Key Novelty

Language-Controlled Neural Memory Updates

Modifies the memory update mechanism to accept a natural language instruction (e.g., 'Learn the format but ignore the facts') alongside the input document
Transforms the memory writing process from an automatic, fixed-objective operation into a controllable action where the instruction dictates how the document is compressed into memory

Breakthrough Assessment

8/10

Introduce a novel paradigm of 'instruction-conditioned memory,' addressing a critical gap in continual learning: the ability to selectively learn from mixed-quality data streams without retraining.

⚙️ Technical Details

Problem Definition

Setting: Continual learning from a stream of document-instruction pairs, where the model must update its memory state to answer future queries based on the instruction constraints

Inputs: A sequence of pairs (Instruction I_t, Document D_t) and intermittent user queries q

Outputs: An updated memory state M_t and predicted answers y conditioned on M_t

Pipeline Flow

Instruction-Document Pair Input
Memory Update Step (Instruction-Conditioned)
Inference Step (Memory-Conditioned Generation)

System Modules

Update Mechanism (U_psi)

Compress the new document into memory tokens based on the learning instruction

Model or implementation: Modified MemoryLLM Encoder (Llama-3 based)

Memory Bank

Store compressed knowledge as continuous vector slots or memory tokens

Model or implementation: Layer-wise memory embeddings (from MemoryLLM architecture)

Response Generator (f_theta)

Generate answers to user queries by attending to the current memory state

Model or implementation: Llama-3-8B (MemoryLLM backbone)

Novel Architectural Elements

Instruction-conditioned memory update rule: U_psi(M_{t-1}, I_t, D_t), explicitly incorporating the instruction I_t into the memory writing process

Modeling

Base Model: MemoryLLM (built on Llama-3-8B)

Training Method: Supervised Fine-Tuning (SFT) on synthetic instruction-following episodes

Objective Functions:

Purpose: Ensure the model learns/ignores specific information according to instructions.

Formally: Minimize expected sequence loss L(y, p_theta(y|q, M_t)) over probes q drawn from current and past timesteps.

Adaptation: Full fine-tuning of memory update parameters (and potentially base model parameters during training)

Training Data:

11,849 documents in train
351 documents in val-id (in-distribution)
2,276 documents in test-ood (out-of-distribution categories and instructions)
Documents generated by sampling 3-8 facts from CounterFACT and rendering in bullet points

Key Hyperparameters:

effective_batch_size: 6 episodes (24 documents)
episode_length: 4 steps
memory_bank_size: 7,098 tokens per layer
+ 1 more
new_memory_size: 256 embeddings per update

Compute: Not reported in the paper

Comparison to Prior Work

vs. MemoryLLM: GNM adds the 'instruction' input to the update step, allowing selective rather than indiscriminate memory updates
vs. ROME/MEMIT: GNM handles streams of heterogeneous documents and behavioral updates, not just isolated fact edits
vs. ICL-FT: GNM compresses history into fixed-size memory, avoiding quadratic attention costs and context window limits
+ 1 more
vs. RAG-FT: GNM integrates information into the model's state rather than relying on retrieval heuristics which can be imprecise

Limitations

Relies on a synthetic benchmark constructed from CounterFACT rather than real-world organic data streams
Performance depends on the base capability of the MemoryLLM architecture
The approach requires training on explicit document-instruction pairs, necessitating a specialized dataset construction pipeline

Reproducibility

Code: https://github.com/maxbennett/Generalized-Neural-Memory

Code and model are open-sourced at https://github.com/maxbennett/Generalized-Neural-Memory. The benchmark is synthetic, constructed from CounterFACT using GPT-5.1 for categorization and document generation.

📊 Experiments & Results

Evaluation Setup

Episodic evaluation where the model processes a sequence of (Document, Instruction) pairs and answers probes

Benchmarks:

Synthetic CounterFACT Benchmark (Continual Instruction Following / Memory Editing) [New]

Metrics:

Fact Accuracy (did it learn the target fact?)
Fact Specificity (did it avoid altering neighborhood facts?)
Fact Selectivity (did it ignore what it was told to ignore?)
Format Accuracy (did it adopt the style?)
Refusal Precision/Recall (did it refuse restricted topics?)
Statistical methodology: Not explicitly reported in the paper

Main Takeaways

Qualitative findings indicate the model successfully generalizes to 'test-ood' instructions (never seen during training), validating the effectiveness of natural language control.
The method reportedly outperforms strong baselines (ICL-FT and RAG-FT) on selectivity and efficiency metrics, particularly in preventing the learning of disallowed information.
The synthetic benchmark demonstrates that GNM can handle compositional instructions (e.g., 'learn facts but refuse category X'), a capability lacking in fixed-objective memory systems.
Note: Specific numeric results were not extractable from the provided text, which ended prior to the Results section, but the abstract and setup describe these outcomes.

📚 Prerequisite Knowledge

Prerequisites

Understanding of Neural Memory (external differentiable memory)
Familiarity with Continual Learning concepts (catastrophic forgetting)
Basic knowledge of Transformer architectures and embeddings

Key Terms

GNM: Generalized Neural Memory—the proposed framework where memory updates are conditioned on natural language instructions

MemoryLLM: The specific base architecture used in this paper, which adds writable memory embeddings to a Llama-3 backbone

ICL: In-Context Learning—providing examples or history directly in the prompt context window

RAG: Retrieval-Augmented Generation—fetching relevant documents from a database to augment the prompt

catastrophic forgetting: The tendency of neural networks to lose previously learned information upon learning new information

CounterFACT: A dataset originally designed for fact editing, adapted here to create a synthetic benchmark for instruction-following memory

test-ood: Out-of-Distribution Test Set—a data split containing fact categories and learning instructions never seen during training

harmonic mean: A type of average used here to aggregate different metrics (accuracy, specificity, selectivity) into a single score