PaperlayerView exact revision on arXiv

Source-linked AI summary

AutoMem: Automated Learning of Memory as a Cognitive Skill

Shengguang Wu, Hao Zhu, Yuhui Zhang, Xiaohan Wang, Serena Yeung-Levy

arXiv:2607.01224v1cs.AIcs.CLcs.MA

TL;DR

Long-horizon tasks expose limits of fixed external-memory mechanisms and make manual optimization of memory decisions impractical. AutoMem automates memory-structure revision and memory-proficiency training, and memory-only optimization improves performance roughly 2×–4× across three games.

  • Problem

    LLM memory is typically a fixed architectural mechanism, while long-horizon memory management remains difficult to optimize manually because its effects can surface thousands of steps later.

  • Method

    AutoMem makes file operations first-class actions and uses meta-LLM-driven loops to revise memory structure and train memory proficiency from trajectory experience.

  • Results

    ∼2×–4× gains across Crafter, MiniHack, and NetHack result from memory-only optimization without modifying the base model’s task-action weights.

  • Takeaways & Limitations

    Memory management is an independently learnable, high-leverage skill for LLM agents on long-horizon tasks.

  • Takeaways & Limitations

    The study uses episodic memory that resets each episode and evaluates the approach only in game environments.

Abstract

from arXiv · show

Memory expertise is a learned skill: knowing what to encode, when to retrieve, and how to organize knowledge--a capacity known in cognitive science as metamemory. We bring this perspective to LLMs by treating memory management as a trainable skill. We promote file-system operations to first-class memory actions alongside task actions, letting the model itself decide how to manage its memory. This memory skill improves along two axes: the structure that supports it (prompts, file schemas, action vocabulary), and the proficiency of the model exercising it. Both axes resist manual optimization: episodes in long-horizon tasks run for thousands of steps, and a single memory mistake can hide long before it surfaces, making human review of full trajectories impractical. We introduce AutoMem, a framework that automates both axes. In the first loop, a strong LLM reviews complete agent trajectories and iteratively revises the memory structure that shapes how the agent interacts with its memory files. In the second loop, the agent's own good memory decisions are identified from many episodes and used as training signal to sharpen the model's memory proficiency directly. Across three procedurally generated long-horizon games (Crafter, MiniHack, and NetHack), optimizing memory alone--without modifying the model's task-action behavior--improved the base agent's performance ~2x-4x, bringing a 32B open-weight model competitive with frontier systems such as Claude Opus 4.5 and Gemini 3.1 Pro Thinking. Our results show that memory management is an independently learnable skill, and a high-leverage objective yielding large gains on long-horizon tasks.

1 Introduction

AutoMem reframes memory management as a trainable metamemory skill for LLM agents, rather than a fixed architectural component. Its automated optimization closes the long-horizon memory bottleneck and produces large gains without changing task-action behavior.

  • Motivation: Metamemory is the learned ability to decide what to remember, when to retrieve it, and how to organize knowledge.External aids such as notes, indices, and files extend cognition beyond working memory.
  • Motivation: LLM context windows act as fixed-size working-memory buffers that long-horizon tasks routinely exceed.Existing retrieval databases, vector stores, scratchpads, and summary buffers typically treat memory as a fixed architectural mechanism.
  • Approach: File-system operations become first-class memory actions, giving the model observable control over what to store, retrieve, and organize.Read, write, search, append, and create operations share the action space with world actions, making every memory decision traceable.
  • Approach: Memory skill improves along two axes: the structure supporting it and the model’s proficiency in using available operations.Structure includes prompts, file schemas, validation logic, and action vocabulary; proficiency is the model’s ability to choose well among them.
  • AutoMem: Meta-LLMs automate both axes by reviewing complete trajectories, revising the scaffold, and selecting memory decisions for targeted training.This trajectory-level review addresses delayed consequences that make manual inspection of thousands-step episodes almost intractable.
  • Results: ∼2×–4× gains over the base agent show that memory-only optimization substantially improves a 32B open-weight model across three long-horizon games.The optimized agent reaches the performance level of frontier systems including Claude Opus 4.5 and Gemini 3.1 Pro Thinking.

2 AUTOMEM

AutoMem combines file-based memory with two sequential outer loops: scaffold revision first, then dedicated memory-proficiency training. Together, these loops optimize what memory operations are available and how effectively the model uses them.

  • 2.1 Inner-loop agent: memory as file system: The inner-loop agent uses a directory of files as external memory while executing one long-horizon episode.The unified action space makes file choices observable and therefore suitable for trajectory-level optimization.
  • 2.1 Inner-loop agent: memory as file system: LOG records useful consequences of prior actions, while PLAN retrieves relevant file information before committing the next world action.The routines divide memory management into encoding and retrieval decisions.
  • 2.2 Outer-loop 1: structure optimization: Outer-loop #1 revises the scaffold, including code, prompts, file schema, and action vocabulary, using complete episode traces.The meta-LLM diagnoses delayed memory failures and retains revisions only when measured progression improves on fixed seeds.
  • 2.3 Outer-loop 2: training memory proficiency: Outer-loop #2 curates the agent’s good memory decisions and trains a dedicated memory specialist with a matched LoRA configuration.The gameplay model remains unmodified, preserving its world-action competence while sharpening memory operations.
  • 2.3 Outer-loop 2: training memory proficiency: All values are progression rate (%), reported as mean ± standard error, with scaffold optimization and memory training applied sequentially.Table 1 compares Qwen2.5-32B-Instruct baselines, scaffold-optized agents, memory-trained agents, and frontier systems.
  • 2.3 Outer-loop 2: training memory proficiency: The two loops optimize different parameters and update signals: scaffold revisions improve structure, while supervised training updates memory-model weights.The scaffold sets the structural ceiling, and proficiency training pushes the model toward that ceiling.

3 Experiments

Across three long-horizon games, memory-only optimization substantially improves performance and reshapes both task behavior and memory use. Scaffold revision delivers the largest gains, while specialist training adds complementary lift through more disciplined retrieval.

  • Performance gains: Crafter 25.0%→47.27% (×1.89), MiniHack 7.5%→27.5% (×3.67), and NetHack 0.42%→1.57% (×3.74) after scaffold optimization alone.Model weights remain untouched; only agent code, prompts, and file schema change.
  • Performance gains: The scaffolded 32B agent outperforms Qwen2.5-72B-Instruct and basic sliding-window context management by wide margins.The comparison indicates structured external memory can matter more than model scale or short-context strategies on these tasks.
  • Performance gains: Training the memory specialist raises progression to 51.36% on Crafter (+4.09), 30.0% on MiniHack (+2.5), and 1.85% on NetHack (+0.28).This lift is obtained on top of the optimized scaffold and is comparable to one or two scaffold iterations.
  • Performance gains: Together, the two optimization axes bring the 32B model to Claude Opus 4.5-level performance and within a few points of Gemini 3.1 Pro Thinking.Claude Opus 4.5 scores 49.5/27.5/2.0 and Gemini 3.1 Pro Thinking scores 55.0/27.5/2.6 across Crafter/MiniHack/NetHack.
  • Behavioral effects: Scaffold optimization reduces unproductive gameplay actions by 32–65% across all three environments.The rate counts stuck steps and oscillation, recovering opportunities for exploration, gathering, and crafting.
  • Behavioral effects: Redundant writes fall −68 to −83%, empty-search rates fall −13 to −50%, and per-step input context shrinks −3 to −30%.These changes indicate leaner memory, more targeted retrieval, and less information for the model to attend to.
  • Memory structures: Training shifts LOG-phase memory writes per SEARCH from 0.84→0.39 in Crafter, 2.89→0.82 in MiniHack, and 4.66→1.31 in NetHack.The specialist increasingly searches existing files before appending new content.
  • Memory structures: The optimized NetHack schema replaces duplicate append-only map entries with coordinate-keyed <|UPSERT_MAP|> updates and synchronized status files.It also pre-populates a strategy reference, making the memory structure more compact and useful.

4 Related Work

Related work spans external-memory architectures, metamemory-inspired agent design, automated agent optimization, and long-horizon game benchmarks. AutoMem differs by revising memory structure from complete trajectories and training memory proficiency directly.

  • External memory: External-memory systems address fixed context windows through retrieval, paging, timestamped memories, and related mechanisms.These approaches generally treat memory as an architectural module designed into the system.
  • Cognitive science: Metamemory frames memory as monitoring and regulating one’s own memory processes, while the Extended-Mind thesis treats notes and files as cognitive aids.AutoMem applies these cognitive-science perspectives to LLM memory management.
  • Automated optimization: Automated agent-optimization methods search architectures, workflows, prompts, or code, whereas AutoMem rewrites code, prompts, and memory-file schemas.Its update signal comes from complete long-horizon trajectory analysis rather than per-question logs.
  • Automated optimization: AutoMem additionally trains memory proficiency, unlike inference-only optimization methods.This is its second outer loop, complementing scaffold optimization.
  • Game agents: Long-horizon games provide stochastic, extended testbeds, but AutoMem studies memory management as the primary lever rather than reasoning, planning, or retrieval architectures.The benchmark context includes Minecraft and NetHack agents alongside broader embodied-agent methods.

5 Conclusion

AutoMem shows that memory management is an independently learnable, high-leverage capability for LLM agents. By optimizing structure and proficiency through meta-LLM-driven loops, memory-only improvements raise long-horizon performance without changing gameplay-model weights.

  • Core finding: ∼2×–4× performance improvement across three long-horizon game environments came from targeting memory alone.The gameplay model’s weights were not modified.
  • Core finding: AutoMem separates memory skill into structure and proficiency, then automates both through meta-LLM-driven outer loops.The framework combines scaffold revision with targeted training of a memory specialist.
  • Core finding: Meta-LLMs can autonomously decompose long-horizon improvement into trajectory-level review and targeted revision.This addresses the impracticality of human review for full episode traces.

6 Limitations and Future Work

The study is limited to episodic memory in game environments and uses environment-specific optimizations. Future work should test persistence, real-world tasks, and shared memory systems.

  • Scope limits: Memory is episodic because the file system starts fresh at the beginning of each episode.Persistent memory carrying knowledge across episodes remains unexplored.
  • Scope limits: Experiments use game environments, whose long horizons, procedural generation, and information-management demands may not represent real-world tasks.The authors identify memory-intensive real-world applications as a future direction.
  • Future work: Each game receives a separate scaffold and memory specialist because the environments differ in structure and objectives.Whether one scaffold or specialist can transfer across environments remains open.

7 Broader Impacts

AutoMem improves long-horizon performance through automated scaffold optimization and targeted memory training, potentially making capable agents more accessible on open-weight models. Deployment in high-stakes settings still requires safety review.

  • Benefits: Automated memory optimization lowers the model-scale threshold at which long-horizon agents become practical.This creates an accessibility gain for open-weight deployment.
  • Broader applicability: The techniques could extend beyond games to other long-horizon, memory-intensive tasks.The passage presents this as an adaptation opportunity rather than a demonstrated deployment result.
  • Safety: Released AutoMem artifacts are not directly applicable to high-stakes deployment without further safety review.The paper explicitly limits immediate operational use in such settings.

A Implementation Details

The implementation materials are released with the paper, including complete prompts and code for both optimization loops.

  • Implementation details: The codebase releases the complete prompt templates for both outer loops.These templates cover the scaffold-optimization and memory-proficiency-training pipeline.
  • Implementation details: The implementation configuration is organized by system components.This organization provides a component-level reference for reproducing the framework.
  • Implementation details: The AutoMem codebase is publicly available at its listed GitHub repository.The passage provides the repository link for access to the released implementation.

A.1 Game environment configuration

The evaluation spans Crafter, eight MiniHack tasks, and NetHack under specified episode limits and prompt-handling configurations. Fixed seeds make scaffold and training comparisons directly comparable.

  • Harness configuration: MiniHack enables autopickup=True, while fallback actions use search where supported and Noop for Crafter.These changes align MiniHack with NetHack and preserve environment-specific fallback behavior.
  • Crafter: Crafter uses a 64×64 world, 9×9 agent view, dense reward, unique_items=True, and 2000 maximum episode steps.The environment also enables skip_more=True to suppress the –More– prompt.
  • MiniHack: MiniHack covers eight tasks with 100 maximum episode steps and a −0.01 step penalty.The task set includes Boxoban, MazeWalk, Corridor, and Quest variants.
  • NetHack: NetHack uses 100,000 maximum episode steps and a no-progress timeout of 150.The shared configuration suppresses the –More– prompt with skip_more=True.
  • Evaluation seeds: Evaluation uses seeds 42 through 51 with 10 Crafter episodes, 40 MiniHack episodes, and 5 NetHack episodes.Episode i uses seed 42 + i, enabling direct per-episode comparisons across versions.

A.2 Outer-loops and training pipeline

AutoMem uses one outer loop to improve the memory scaffold and another to train memory proficiency. Both stages rely on evaluation-driven selection and deployment into the optimized scaffold.

  • Outer-loop #1: Scaffold revisions are accepted only when average progression on the same fixed evaluation seeds strictly improves.A failed revision receives up to one retry before the meta-loop restarts from a clean session.
  • Outer-loop #2: The second outer loop collects data, selects training examples and configurations, then performs LoRA finetuning with two-model deployment.It operates on the final scaffold: V5 for Crafter, V4 for MiniHack, and V2 for NetHack.
  • Data collection: Training collects 100 Crafter episodes, 400 MiniHack episodes, and 50 NetHack episodes under the final scaffold.Training seeds are explicitly disjoint from evaluation seeds [42, . . . , 51].
  • Data-engine selection: Claude Opus 4.7 selects dataset composition, trace examples, and LoRA configurations based on evaluation trajectories.The training engine iteratively refines these choices across passes.
  • Data filtering: Postprocessing removes action-only examples and retains memory-operation reasoning when traces contain both memory and gameplay actions.It also cleans code-block wrappers and trims gameplay action commitments.
  • LoRA training: LoRA training uses cutoff_len=16384, bf16 precision, AdamW, cosine scheduling, and two GPUs with DeepSpeed ZeRO-3.The reported LoRA rank, batch size, and epoch count vary by environment.

B Memory-scaffold evolution by iteration

The scaffold evolves by adding structured knowledge, tracking progress and exploration, and introducing specialized memory operations. Each revision targets a concrete failure mode while preserving the core interface.

  • Evolution strategy: Every scaffold change edits prompts, memory-file layout, or automatic observation-derived hints.The two-phase LOG/PLAN loop and memory-operation interface remain unchanged.
  • Crafter: Crafter revisions add game knowledge, maintained goals, resource warnings, inventory deltas, achievement tracking, and movement-loop detection.These changes progressively make planning state and failure signals explicit.
  • MiniHack: MiniHack revisions track recent actions, visited cells, unexplored directions, task-specific rules, and constrained replies to directional questions.The scaffold also auto-loads recent action-log entries.
  • NetHack: NetHack adds UPSERT_MAP to overwrite stale coordinate entries rather than accumulating duplicates.Later revisions also trim the action log, pre-fill strategy guidance, and issue staircase descent directives.
Loading 2607.01224v1…