cs.LGFeb 26, 2026
SaveTell Me What To Learn: Generalizing Neural Memory to be Controllable in Natural Language
Organizations: Columbia University
Abstract
Modern machine learning models are deployed in diverse, non-stationary environments where they must continually adapt to new tasks and evolving knowledge. Continual fine-tuning and in-context learning are costly and brittle, whereas neural memory methods promise lightweight updates with minimal forgetting. However, existing neural memory models typically assume a single fixed objective and homogeneous information streams, leaving users with no control over what the model remembers or ignores over time. To address this challenge, we propose a generalized neural memory system that performs flexible updates based on learning instructions specified in natural language. Our approach enables adaptive agents to learn selectively from heterogeneous information sources, supporting settings, such as healthcare and customer service, where fixed-objective memory updates are insufficient.
Figures & tables
Figure 1: An AI system with memory can adapt to its environment continuously by integrating diverse information sources. We propose a generalized neural memory system that performs flexible long-term updates based on learning instructions specified in natural language. Our approach enables important use cases in critical domains such as healthcare, where an adaptive agent must learn from heterogeneous documents that preclude using a neural memory system with a fixed objective.
Figure 2: Examples from our benchmark, including a document and a sample of three possible learning instructions, each with a sample of possible queries and correct responses.
Figure 3: Continual Learning of Targeted Facts (Heatmaps). Shows average performance on each desiderata across an episode. The x-axis ‘Documents Seen’ represents the sequence index that the memory state is in (e.g., x-axis value of represents the memory state after seeing the first documents in an episode). The y-axis ‘Query Timestep’ represents the sequence element that a set of queries derives from (e.g., y-axis value of represents the queries that are sampled from the document-instruction pair that was the th element in the episode sequence). The bottom diagonal is ignored because these represent queries for document-instruction pairs that have not yet been seen.
Figure 4: Continual Learning of Targeted Facts (Overall Performance). The left chart (‘Total Performance’) shows the average performance across all time steps. ‘Score’ is the harmonic mean of these averages across accuracy, specificity, and selectivity. ‘Base’ is the performance of Llama-3 probed on only queries, without any document or learning instruction. We show base model performance for reference. The base model naturally does well on specificity and selectivity, as it represents the performance of the model without any interference from new facts. The right chart (‘Performance By Recency’) reports the ‘score’ of the queries associated with document learned steps ago, where can be 0, 4, or 9. Error bars show 95% confidence intervals (CI).
Figure 5: Continual Learning of Knowledge, Styles, and Behaviors. Here we report results of our ‘Continual Learning of Knowledge, Styles, and Behaviors’ experiment. For the left three charts, we report results by recency to evaluate retention performance over the course of an episode. For the plot on the far right, we show the inference cost in FLOPs per token based on how many document-instruction pairs have been learned. Error bars show 95% CI. Full details are reported in Appendix C.6 .
Figure 6: Format Generalization. Format accuracy performance on formats seen during training (‘Train Formats’) vs. those used only in validation (‘Val-ID Formats’) vs those used only in our test-ood data (‘Test-OOD Formats’) in our ‘Continual Learning of Knowledge, Styles, and Behaviors’ experiment. Error bars show 95% CI. See Table 8 for details on format types.
Figure 7: Compositional Generalization. Results of ‘Compositional Generalization’ experiment. See Section 5.3 and Appendix D for details. Error bars show 95% CI.
Figure 8: Ablation . Shows how performance degrades when training an ablated version of GNM (‘GNM (Ablation)’) that only passes gradients over the inference step, not over the learning step. Error bars show 95% CI.
Figure 9: Memory Analysis . ‘Target Alignment’ is computed as the dot product , where is the mean memory update across all new memory tokens in a given layer, and is the normalized direction from distractor to target hidden states at each layer. Positive values indicate the memory update encodes more target information. Error bars show 95% CI. Left plot shows alignment across all layers. Right plot shows averaged alignments on layers 15-30. See Appendix E for details.
Figure 10: Layer-Wise Ablation : for each layer , we replace GNM’s layer with the corresponding ablation model layer and measure the resulting accuracy and selectivity on held-out examples. We report on the degradation in accuracy and selectivity performance across each of these swaps. y-axis is the drop in performance (0.25 means a drop of 25%, so selectivity might go from 75% to 50%, for example). x-axis is the layer being swapped. Results show clearly that GNM is doing something unique in layers 5-14 that causally improves selectivity, while not being particularly important for accuracy.
Appendix figures & tables16 assets
Supplementary material from the paper’s appendix.
Appendix
| Grouping | Category | # total | # train | # val-id | # test-ood | Example |
| Language | Major western european languages | 2,210 | 1,997 | 213 | 213 | The mother tongue of Michel Poniatowski is Spanish |
| eastern european mediterranean languages | 625 | 559 | 66 | 66 | The language of Allegro Non Troppo was Serbian. | |
| asian middle eastern languages and pacific | 469 | 426 | 43 | 43 | Alain de Benoist speaks Sanskrit | |
| northern central european languages | 410 | 0 | 0 | 410 | In Kuusamo, the language spoken is Swedish | |
| Location | non us cities or states | 4,285 | 3,833 | 426 | 452 | Institut Polair originated in Budapest. |
| country | 3,138 | 2,826 | 312 | 312 | Delta Goodrem, that originated in India. |
Table 1: Distribution of factual categories and subcategories, ordered by frequency within each category. Each cell represents the number of unique facts in each dataset. Note that the ”# total” is not a sum of the rows to the right, because val-id and train-id share non-held out facts, the difference between them is that test-ood contains 4 special held-out fact categories while train and val-id do not.
| What Facts Are Targeted | Example Learning Instruction |
|---|---|
| All | Please memorize all facts. Ignore the document’s formatting; only learn the factual information. |
| Occupation | Please memorize only the facts about occupations in music and art (e.g., painter, novelist), sports (e.g., quarterback, midfielder), science and academia (e.g., physicist, historian), law (e.g., lawyer, judge), medicine (e.g., physician, nurse), politics (e.g., mayor, governor), entertainment (e.g., actor, comedian), and religion (e.g., priest, rabbi). Exclude locations, languages, or companies. Ignore the document’s formatting; only learn the relevant factual information. |
| Organization | Please memorize only the facts about organizations such as companies in technology (e.g., Apple), entertainment (e.g., Netflix), automotive (e.g., Ford), industrial (e.g., General Electric), oil (e.g., ExxonMobil), gaming (e.g., Nintendo), and religions (e.g., Christianity). Exclude locations, languages, or jobs. Ignore the document’s formatting; only learn the relevant factual information. |
| Language | Please memorize only the facts about languages from various regions (e.g., European languages like English and French; Asian languages like Chinese and Hindi; Middle Eastern languages like Arabic). Exclude locations, companies, religions, or jobs. Ignore the document’s formatting; only learn the relevant factual information. |
| Location | Please memorize only the facts about locations, including cities, states, countries, and continents. Exclude languages, companies, religions, or jobs. Ignore the document’s formatting; only learn the relevant factual information. |
| US cities or states | Please memorize only the facts about cities, states, boroughs, and neighborhoods that are located within the United States. Exclude any cities outside the US, continents, countries, languages, companies, or jobs. Ignore the document’s formatting; do not learn it—only learn the relevant factual information. |
Table 2: Examples of learning instructions in experiment 1 (part one) . The top 17 are those types of learning instructions used during training, the bottom four are the held out learning instructions used in testing.
| What Facts Are Targeted | Example Learning Instruction |
|---|---|
| Tech Industrial or Gaming Company | Please memorize only the facts about technology companies, industrial manufacturers, oil companies, and gaming companies. Exclude TV networks, car companies, religions, places, languages, or jobs. Ignore the document’s formatting; do not learn it—only learn the relevant factual information. |
| TV Entertainment or News Organization | Please memorize only the facts about entertainment studios, record labels, TV channels, news outlets, and media companies. Exclude tech companies, car companies, religions, places, languages, or jobs. Ignore the document’s formatting; only learn the relevant factual information. |
| Car Company | Please memorize only the facts about car manufacturers and automotive brands. Exclude other kinds of companies, places, religions, languages, or jobs. Ignore the document’s formatting; only learn the factual information. |
| Music or Art Related Occupation | Please memorize only the facts about music, art, literature, and entertainment genres (e.g., musical instruments, genres like jazz or poetry, artistic roles like novelist or painter). Exclude sports, science, politics, places, languages, or companies. Ignore the document’s formatting; learn only the factual information. |
| Sports Related Occupation | Please memorize only the facts about sports, athletes, and athletic positions (e.g., quarterback, midfielder). Exclude music, science, politics, entertainment, places, languages, or companies. Ignore the document’s formatting; learn only the factual information. |
| Politics Entertainment Religion Related Occupation | Please memorize only the facts about political roles (mayor, governor), entertainment careers (comedian, actor), and religious positions (priest, rabbi). Exclude music, sports, science, places, languages, or companies. Ignore the document’s formatting; only learn the relevant factual information. |
Table 3: Examples of learning instructions in experiment 1 (part two) . The top 17 are those types of learning instructions used during training, the bottom four are the held out learning instructions used in testing.
| Prompt | Accuracy | Specificity | Selectivity | Harmonic Mean |
|---|---|---|---|---|
| Prompt 1 | 84.1 | 66.2 | 16.6 | 34.4 |
| Prompt 2 | 82.8 | 69.5 | 16.1 | 33.9 |
| Prompt 3 | 82.1 | 67.5 | 17.9 | 36.21 |
| Prompt 4 | 88.1 | 58.1 | 12.6 | 27.8 |
| Prompt 5 | 81.7 | 69.0 | 17.8 | 36.18 |
Table 4: Performance metrics for different prompt formulations in warmup experiment.
Figure 11: Heatmap for all models on experiment 1 when rolling out to sequence length of 20 instead of 10.
| What Is Targeted For Learning | Example Learning Instruction |
|---|---|
| Formats | Learn only the document’s formatting style (e.g., JSON, YAML, XML, CSV, TOML, etc.), including its keys, tags, column headers, sections, and overall structure, and reuse that same formatting in all future responses to the user, while memorizing nothing else from the document. |
| Refuse US cities or states | Learn the contents of the document, but if any of the contents of the document contain information about US cities or states, refuse to answer any questions from the user about this specific fact, if a user asks about it, respond ’Sorry’. Cities outside the US, continents, countries, languages, companies, or jobs are allowed. |
| Refuse Non-US cities or states | Learn the contents of the document, but if any of the contents of the document contain information about cities, regions, provinces, or subnational areas located outside the United States, refuse to answer any questions from the user about this specific fact, if a user asks about it, respond ’Sorry’. US places, continents, countries, languages, companies, or jobs are allowed. |
| Refuse Countries | Learn the contents of the document, but if any of the contents of the document contain information about about sovereign countries or widely recognized nations, refuse to answer any questions from the user about this specific fact, if a user asks about it, respond ’Sorry’. Facts about cities, subnational regions, continents, languages, companies, or jobs are allowed. |
| Refuse Major Western European Languages | Learn the contents of the document, but if any of the contents of the document contain information about major Western European languages (e.g., English, French, Spanish, Italian), refuse to answer any questions from the user about this specific fact, if a user asks about it, respond ’Sorry’. Other languages, places, companies, or jobs are allowed. |
| Refuse Eastern European Mediterranean Languages | Learn the contents of the document, but if any of the contents of the document contain information about Eastern European and Mediterranean languages (e.g., Russian, Polish, Greek), refuse to answer any questions from the user about this specific fact, if a user asks about it, respond ’Sorry’. Other languages, places, companies, or jobs are allowed. |
Table 5: Examples of additional learning instructions in experiment 2 (part one) . The top 18 are those types of learning instructions used during training, the bottom four are the held out learning instructions used in testing.
| What Is Targeted For Learning | Example Learning Instruction |
|---|---|
| Refuse Car Company | Learn the contents of the document, but if any of the contents of the document contain information about car manufacturers and automotive brands, refuse to answer any questions from the user about this specific fact, if a user asks about it, respond ’Sorry’. Facts about other kinds of companies, places, religions, languages, or jobs are allowed. |
| Refuse Music or Art Related Occupation | Learn the contents of the document, but if any of the contents of the document contain information about music, art, literature, and entertainment genres (e.g., musical instruments, genres like jazz or poetry, artistic roles like novelist or painter), refuse to answer any questions from the user about this specific fact, if a user asks about it, respond ’Sorry’. Facts about sports, science, politics, places, languages, or companies are allowed. |
| Refuse Sports Related Occupation | Learn the contents of the document, but if any of the contents of the document contain information about sports, athletes, and athletic positions (e.g., quarterback, midfielder), refuse to answer any questions from the user about this specific fact, if a user asks about it, respond ’Sorry’. Facts about music, science, politics, entertainment, places, languages, or companies are allowed. |
| Refuse Politics Entertainment Religion Related Occupation | Learn the contents of the document, but if any of the contents of the document contain information about political roles (mayor, governor), entertainment careers (comedian, actor), and religious positions (priest, rabbi), refuse to answer any questions from the user about this specific fact, if a user asks about it, respond ’Sorry’. Facts about music, sports, science, places, languages, or companies are allowed. |
| Refuse Northern Central European Languages | Learn the contents of the document, but if any of the contents of the document contain information about Northern and Central European languages (e.g., German, Swedish, Finnish), refuse to answer any questions from the user about this specific fact, if a user asks about it, respond ’Sorry’. Other languages, places, companies, or jobs are allowed. |
| Refuse Science academia law medicine related occupation | Learn the contents of the document, but if any of the contents of the document contain information about science, academia, law, journalism, and medicine (e.g., physicist, historian, lawyer, physician), refuse to answer any questions from the user about this specific fact, if a user asks about it, respond ’Sorry’. Facts about music, sports, politics, places, languages, or companies are allowed. |
Table 6: Examples of additional learning instructions in experiment 2 (part two) . The top 18 are those types of learning instructions used during training, the bottom four are the held out learning instructions used in testing.
| Prompt | Fact Accuracy | Fact Specificity | Fact Selectivity | Format Accuracy | Refusal Precision | Refusal Recall | Refusal Specificity | Harmonic Mean |
|---|---|---|---|---|---|---|---|---|
| Prompt 1 | 79.5 | 70.8 | 20.3 | 0.0 | 71.4 | 12.9 | 85.6 | 33.5 |
| Prompt 2 | 81.2 | 70.2 | 19.1 | 0.0 | 55.6 | 15.8 | 87.8 | 35.0 |
| Prompt 3 | 85.7 | 60.7 | 18.6 | 0.0 | 41.5 | 69.8 | 26.5 | 38.0 |
| Prompt 4 | 82.2 | 70.0 | 22.0 | 0.0 | 56.1 | 25.0 | 71.8 | 41.8 |
| Prompt 5 | 82.4 | 68.3 | 21.2 | 0.0 | 53.8 | 24.1 | 72.9 | 40.6 |
Table 7: Fact, formatting, and refusal performance across prompt variants. Format accuracy ignored from harmonic mean calculation because all was 0.
| Format type | Key used | Exact format | Datasets used within (train/val-id/test-ood) |
|---|---|---|---|
| JSON | response | {’response’:’…’} | train |
| message | {’message’:’…’} | train | |
| string | {’string’:’…’} | train | |
| chat | {’chat’:’…’} | train | |
| reply | {’reply’:’…’} | train | |
| answer | {’answer’:’…’} | train |
Table 8: Types of Formats used.
Figure 12: Heatmap for all models on experiment 2. X-axis is how many document-instruction pairs has been provided thus far in the sequence, and the y-axis is which sequence index the queries derive from. The bottom diagonal is ignored because these represent queries for document-instruction pairs that have not yet been memorized by the model.
| Facts | Formats | Refusal | |||||||
|---|---|---|---|---|---|---|---|---|---|
| Model | Score | Acc. | Sel. | Spec. | Acc. | Sel. | F1 | Precision | Recall |
| RAG-FT | 70.5 | 74.5 | 61.8 | 77.1 | 0.0 | 99.9 | 69.4 | 65.6 | 73.7 |
| ICL-FT | 67.3 | 79.1 | 51.5 | 80.1 | 4.4 | 99.7 | 89.8 | 92.7 | 87.1 |
| MemLLM | 38.7 | 81.8 | 19.2 | 75.2 | 0.0 | 100.0 | 100.0 | ||
| GNM | 84.1 | 95.8 | 77.4 | 81.1 | 67.4 | 100.0 | 93.2 | 88.7 | 98.1 |
| GNM (Ablation) | 44.5 | 75.4 | 24.8 | 72.2 | 0.0 | 75.4 | 51.0 | 47.7 | 54.7 |
Table 9: Experiment #2 detailed results for most recent time step. Note that Format Selectivity is an additional metric we report on here, defined as: when the model is instructed to ignore the formatting, what percentage of the time does the model correctly ignore the document’s formatting. Values below 1 are reported as .
Figure 13: Performance on experiment 2 when rolling out to sequence length of 20 instead of 10.
| Targeted Category For Fact Learning | Targeted Category For Refusal Learning | Learning Instruction |
|---|---|---|
| Non-US cities or States | Major Western European Languages | Please memorize only the facts about cities, regions, provinces, or other subnational areas outside the United States. If a fact in the document concerns major Western European languages such as English, French, Spanish, or Italian, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| Non-US cities or States | Eastern European Mediterranean Languages | Please memorize only the facts about cities, regions, provinces, or other subnational areas outside the United States. If a fact in the document concerns Eastern European and Mediterranean languages such as Russian, Polish, or Greek, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| Non-US cities or States | Tech Industrial or Gaming Company | Please memorize only the facts about cities, regions, provinces, or other subnational areas outside the United States. If a fact in the document concerns technology, industrial, oil, or gaming companies, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| Non-US cities or States | TV Entertainment or News Organization | Please memorize only the facts about cities, regions, provinces, or other subnational areas outside the United States. If a fact in the document concerns entertainment studios, record labels, TV channels, news outlets, or other media organizations, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| Non-US cities or States | Music or Art Related Occupations | Please memorize only the facts about cities, regions, provinces, or other subnational areas outside the United States. If a fact in the document concerns music, art, literature, or entertainment-related roles (for example, genres, instruments, or artistic occupations), refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| Non-US cities or States | Sports Related Occupation | Please memorize only the facts about cities, regions, provinces, or other subnational areas outside the United States. If a fact in the document concerns sports, athletes, or athletic positions such as quarterback or midfielder, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
Table 10: Compositional learning instructions used in experiment 3 (Part One).
| Targeted Category For Fact Learning | Targeted Category For Refusal Learning | Learning Instruction |
|---|---|---|
| Major Western European Languages | Music or Art Related Occupations | Please memorize only the facts about major Western European languages such as English, French, Spanish, or Italian. If a fact in the document concerns music, art, literature, or entertainment-related roles (for example, genres, instruments, or artistic occupations), refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| Major Western European Languages | Sports Related Occupation | Please memorize only the facts about major Western European languages such as English, French, Spanish, or Italian. If a fact in the document concerns sports, athletes, or athletic positions such as quarterback or midfielder, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| Eastern European Mediterranean Languages | Non-US Cities or States | Please memorize only the facts about Eastern European and Mediterranean languages such as Russian, Polish, or Greek. If a fact in the document concerns cities, regions, provinces, or other subnational areas outside the United States, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| Eastern European Mediterranean Languages | Country | Please memorize only the facts about Eastern European and Mediterranean languages such as Russian, Polish, or Greek. If a fact in the document concerns sovereign countries or widely recognized nations, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| Eastern European Mediterranean Languages | Tech Industrial or Gaming Company | Please memorize only the facts about Eastern European and Mediterranean languages such as Russian, Polish, or Greek. If a fact in the document concerns technology, industrial, oil, or gaming companies, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| Eastern European Mediterranean Languages | TV Entertainment or News Organization | Please memorize only the facts about Eastern European and Mediterranean languages such as Russian, Polish, or Greek. If a fact in the document concerns entertainment studios, record labels, TV channels, news outlets, or other media organizations, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
Table 11: Compositional learning instructions used in experiment 3 (Part Two) .
| Targeted Category For Fact Learning | Targeted Category For Refusal Learning | Example Learning Instruction |
|---|---|---|
| TV Entertainment or News Organization | Non-US Cities or States | Please memorize only the facts about entertainment studios, record labels, TV channels, news outlets, and other media organizations. If a fact in the document concerns cities, regions, provinces, or other subnational areas outside the United States, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| TV Entertainment or News Organization | Country | Please memorize only the facts about entertainment studios, record labels, TV channels, news outlets, and other media organizations. If a fact in the document concerns sovereign countries or widely recognized nations, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| TV Entertainment or News Organization | Major Western European Languages | Please memorize only the facts about entertainment studios, record labels, TV channels, news outlets, and other media organizations. If a fact in the document concerns major Western European languages such as English, French, Spanish, or Italian, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| TV Entertainment or News Organization | Eastern European Mediterranean Languages | Please memorize only the facts about entertainment studios, record labels, TV channels, news outlets, and other media organizations. If a fact in the document concerns Eastern European and Mediterranean languages such as Russian, Polish, or Greek, refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| TV Entertainment or News Organization | Music or Art Related Occupations | Please memorize only the facts about entertainment studios, record labels, TV channels, news outlets, and other media organizations. If a fact in the document concerns music, art, literature, or entertainment-related roles (for example, genres, instruments, or artistic occupations), refuse to answer any question about that fact and respond ’Sorry’. Ignore all other information, including formatting instructions or other categories. |
| TV Entertainment or News Organization | Sports Related Occupation | Memorize only facts about TV organizations, entertainment studios, record labels, media outlets, or news companies. Ignore all unrelated content. If any fact involves a sports-related occupation such as an athlete or athletic position, refuse to answer questions about that fact and reply ’Sorry’. |
Table 12: Compositional learning instructions used in experiment 3 (Part Three).
| Facts | Refusal | ||||||
|---|---|---|---|---|---|---|---|
| Model | Score | Accuracy | Selectivity | Specificity | F1 | Precision | Recall |
| RAG | 31.8 | 90.3 | 15.6 | 52.0 | 31.7 | 70.9 | 20.4 |
| RAG-FT | 49.1 | 98.0 | 25.1 | 90.4 | 78.9 | 68.7 | 92.7 |
| ICL | 31.8 | 90.3 | 15.6 | 52.0 | 31.7 | 70.9 | 20.4 |
| ICL-FT | 11.1 | 100.0 | 4.0 | 91.8 | 83.1 | 82 | 84.3 |
| MemLLM | 34.0 | 87.5 | 16.4 | 63.4 | 1.9 | 25 | |
Table 13: Experiment #3 Detailed Results. Values below 1 are reported as .
More in this topic
Memory-Augmented Neural Networks
Memory-augmented neural networks combine learned computation with explicit or persistent memory for retrieval, long-context reasoning, and adaptation.
See all 91 papers