PaperlayerView exact revision on arXiv

Source-linked AI summary

Dynamic Mixture of Latent Memories for Self-Evolving Agents

Dianzhi Yu, Vireo Zhang, Hongru Wang, Yanyu Chen, Minda Hu, Wanghan Xu, Siki Chen, Philip Torr, Zhenfei Yin, Irwin King

arXiv:2605.21951v1cs.LG

TL;DR

Agents need to accumulate knowledge across changing task sequences without forgetting, while existing parameter updates and external memory have distinct limitations. MoLEM adds routed latent memory while keeping the reasoner frozen; after the full sequence, its average accuracy is 10.40% above the Vanilla pretrained baseline.

  • Problem

    Agents need to acquire knowledge across changing task sequences without forgetting, while parameter updates risk forgetting and external memory does not improve intrinsic capabilities.

    Source: §1 · pp. 1–2

  • Method

    MoLEM uses a dynamic mixture of experts to generate and route latent memories into reasoning while keeping the base reasoner frozen.

    Source: §1 · p. 2, §4.1 · p. 4

  • Results

    After the full continual-learning sequence, MoLEM’s average accuracy was 10.40% above the Vanilla pretrained baseline, and competitors did not consistently exceed it across training orders.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Experiments on cross-domain continual-learning sequences validate the framework’s effectiveness while the core reasoner remains frozen.

    Source: §6 · p. 11

    MoLEM adds expert, router, and autoencoder parameters, though these modules are compact relative to the backbone and provide capacity isolation across domains.

    Source: §6 · pp. 11–12

Abstract

from arXiv · show

Achieving self-evolution in intelligent agents requires the continual accumulation of new knowledge across changing task sequences without forgetting previously acquired abilities. Existing approaches either internalize knowledge by updating model parameters, which induces catastrophic forgetting, or rely on external memory, which fails to genuinely enhance the model's intrinsic capabilities. We propose MoLEM, a generative mixture of latent memory framework based on a dynamic mixture-of-experts (MoE). We treat multiple experts as independent carriers to generate memory. A router selects and weights experts through key-query matching, and the aggregated latent memory is injected into the reasoning process. The base model for reasoning remains entirely frozen, with all experiential knowledge internalized into the additional modules, avoiding catastrophic forgetting. For continual learning, each training stage is paired with a lightweight autoencoder that selects the appropriate routing group at inference, and inputs that match no stage fall back to the pretrained model. Experiments train the framework on continual-learning sequences spanning math, science, and code domains. After training, we evaluate the framework on the corresponding test sets to measure task learning and competence preservation across continual adaptation stages. After the full continual-learning sequence, our method improves the average accuracy by 10.40% over the Vanilla pretrained baseline, while none of the competing methods consistently exceed this baseline across different training orders.

1. Introduction

The paper asks how agents can acquire knowledge across changing tasks without forgetting, and proposes latent memory that can be used by a frozen reasoner. MoLEM uses a dynamic mixture of experts and stage-wise routing for this purpose.

Motivation

  • Parameter updates can cause catastrophic forgetting and task interference, while retrieval-based memory leaves the model’s intrinsic capabilities unchanged.The paper motivates latent memory as a way to internalize experience without modifying the original parameters.Source: §1 · pp. 1–2
  • The central challenge is to build agents that expand capacity across diverse tasks while preserving previously acquired abilities.The experiments evaluate new-task learning and competence preservation across continual-learning stages.Source: §1 · pp. 1–2

Method

  • MoLEM keeps the reasoning base frozen and updates MoE memory modules, using key-query routing to aggregate experts’ latent memories for reasoning.Stage-wise autoencoders select a routing group by reconstruction error; inputs matching no stage fall back to the pretrained model.Source: §1 · p. 2

Experimental Validation

  • 10.40% higher average accuracy than the Vanilla pretrained baseline after the full sequence, while no competing method consistently exceeded that baseline across training orders.Experiments cover continual-learning sequences spanning math, science, and code.Source: §1 · p. 2

2. Related Work

The related work contrasts parameter-updating and retrieval-based memory with latent-memory methods. MoLEM addresses the stated need for latent memory that can scale under continual learning.

Parametric Methods

  • Parametric approaches internalize task knowledge through parameter updates but remain exposed to catastrophic forgetting in continual-learning settings.The cited examples include methods that integrate reasoning steps or memory mechanisms into training.Source: §2 · p. 3

Retrieval-based Memory

  • Retrieval-based systems avoid parameter modification but remain context engineering and do not improve the model’s intrinsic capabilities.Examples include external memory systems that retrieve experiences or organize structured notes.Source: §2 · p. 3

Latent Memory

  • Latent-memory methods use continuous representations to support reasoning without changing original parameters, but existing methods do not address architectural scalability under continual learning.Examples include MemGen, SoftCoT [10], and Coconut.Source: §2 · p. 3

3. Preliminary

The preliminary section describes latent-memory generation and notation, then formulates domain-incremental learning without task labels. Its objectives include learning on covered domains and preserving pretrained competence on uncovered ones.

3.1. Latent Memory Generation

  • Latent-memory generation produces a continuous sequence at delimiter positions and injects it as an auxiliary representation to support later reasoning and next-token prediction.MoLEM uses different experts to generate memories, then routes and aggregates them into an auxiliary signal.Source: §3.1 · p. 3

3.2. Notation

  • The reasoner π_θ stays frozen; experts generate candidate memories M_k,t, which are aggregated as M_t for injection into the reasoner.The experts and router are adapters attached to the pretrained reasoner.Source: §3.2 · p. 4

3.3. Problem Formulation

  • In the task-ID-agnostic setting, training accesses only the current domain’s data at each phase, and optimization changes expert and router parameters while θ remains fixed.The task sequence can span domains such as math, science, and code.Source: §3.3 · p. 4
  • Competence preservation is assessed on stage-wise test sets: learned domains should absorb new knowledge while preserving prior performance, and uncovered OOD domains should remain near the pretrained baseline.The frozen pretrained reasoner’s accuracy without a learned memory module defines that competence baseline.Source: §3.3 · p. 4

4. Methodology

MoLEM combines independent latent-memory experts with key-query routing while keeping the base reasoner frozen. Stage-specific autoencoders choose a routing group or reject inputs as out of distribution.

4.1. Overview: Dynamic MoE with Latent Memory

  • MoLEM uses independent experts, a key-query router, and per-stage autoencoders to select a routing group or reject out-of-distribution inputs.The base reasoner remains frozen, and inputs rejected by all stage thresholds fall back to the pretrained model.Source: §4.1 · pp. 4–5

4.2. Latent-Memory Experts

  • Each expert generates a distinct latent memory from its expert-modified hidden states, which is injected into reasoning.Parameter isolation allows different experts to produce different memories for the same input.Source: §4.2 · p. 5

4.3. Key-query routing

  • The router projects its last-position hidden state into a query, matches it by cosine similarity to learnable expert keys, and uses the scores to weight experts.The query represents the current reasoning state, and each key serves as a learnable routing prototype.Source: §4.3 · p. 6

4.4. Training Objective

  • Training uses supervised fine-tuning to maximize the frozen reasoner’s target-token probability given preceding tokens and aggregated memory, with an auxiliary load-balancing loss.Balancing is computed over routing instances, and previous-stage experts are excluded from the trainable expert set.Source: §4.4 · p. 6

4.5. Continual Expansion: Stage-wise Recruiting and Domain-Aware Routing

  • At each new domain, MoLEM freezes existing expert-router groups, recruits a randomly initialized group, and updates only the new group.Knowledge is isolated in separate parameter subspaces; stage autoencoders select among seen groups, while inputs exceeding all calibrated thresholds are rejected as OOD.Source: §4.5 · p. 7

4.6. Continual Learning Pipeline

  • The continual-learning sequence proceeds through Math, Science, and Code, training only on each current domain while evaluating seen and unseen domains.At inference, stage autoencoders route accepted inputs without a task identifier; otherwise the input is rejected as OOD.Source: §4.6 · pp. 7–8

5. Experiments

Experiments evaluate MoLEM on continual-learning sequences across math, science, and code, testing both new-task learning and competence preservation. The results favor stage-isolated routing with multiple experts: MoLEM improves the final average over the Vanilla baseline, while analyses examine ablations, expert specialization, and task-free routing.

5.2. Main Results: Continual Learning Performance

  • 73.93% Stage 3 average accuracy exceeds Vanilla (63.53%) by 10.40%, MemGen (63.47%) by 10.46%, and SFT (34.40%) by 39.53%.The comparison is after the full training sequence.Source: §5.2 · p. 8, Table 1 · p. 9
  • Across the three examined training orders, only MoLEM delivers a positive net gain over Vanilla after the full sequence.Under the default order, competing methods finish at or below the 63.53% Vanilla reference; MoLEM’s Average gain is +10.40%.Source: §5.2 · pp. 8–9
  • Global parameter updates degrade earlier-domain performance, while retrieval-based ExpeL avoids parametric drift but remains limited on Math and falls below Vanilla on Code.SFT Math drops from 57.20% at Stage 1 to 12.80% at Stage 3; ExpeL Math stays in the 28.00%–32.00% range and its Stage 3 Code accuracy is 60.60%, below Vanilla’s 72.00%.Source: §5.2 · p. 9

5.3. Ablation Study

  • Stage isolation raises the multi-expert Stage 3 average from 51.80 to 73.93 and reduces Forget from 30.50 to 0.00.With isolated routing, multiple experts also outperform a single expert: Stage 3 average rises from 70.20 to 73.93, and Math accuracy from 53.80 to 59.20.Source: §5.3 · p. 10, Table 2 · p. 10
  • Removing the load-balance loss reduces Stage 1 Math Test accuracy from 59.80% to 58.00% under the same configuration.The loss encourages more balanced expert exposure during training.Source: §5.3 · p. 10

5.4. Analysis

  • Expert use varies by domain: Math favors E1 and E2, while Code distributes usage more broadly but assigns the largest share to E3.The analysis describes these input-dependent preferences as spontaneous expert diversity, without manually specified expert roles.Source: §5.4 · p. 10
  • Latent memories cluster by expert on a single domain and by domain across stages, indicating distinct, stage-specialized representations.Figures 4 and 5 show the reported t-SNE separations.Source: §5.4 · p. 11
  • Stage-wise autoencoders route inputs by reconstruction-error compatibility among seen stages and reject unmatched inputs to the pretrained reasoner.Before a dataset’s stage is learned, its inputs are mostly rejected; after its autoencoder is added, they are primarily routed to that stage. The routing decision uses the frozen reasoner’s hidden-state space rather than a provided task identity.Source: §5.4 · p. 11

6. Conclusion

MoLEM uses a dynamic mixture of latent-memory experts while keeping the core reasoner frozen. Its added expert, router, and autoencoder modules incur extra parameters, and the paper reports experiments on cross-domain continual-learning sequences.

6. Conclusion

  • MoLEM uses key-query routing to select and weight experts whose latent memory is injected into reasoning, while the core reasoner remains frozen.The paper describes the framework as averting catastrophic forgetting at the architectural level.Source: §6 · p. 11
  • Compared with a purely frozen backbone, MoLEM adds expert, router, and autoencoder modules that incur extra parameters.The authors describe the modules as parameter-efficient and compact relative to the backbone.Source: §6 · pp. 11–12
  • The paper reports validating the framework on cross-domain continual-learning sequences.It also states that the modules provide capacity isolation across domains so continual adaptation does not cause cross-domain interference.Source: §6 · p. 12

A. Implementation Details

Implementation pairs cached prompt-end features with independently trained domain autoencoders for routing, alongside a stage-specific expert router and fixed latent-memory injection settings.

A. Implementation Details

  • Each domain’s autoencoder reconstructs the cached final-token hidden state from the frozen reasoner at the latent-memory insertion position.The prompt-end feature is cached once per prompt; one small autoencoder is trained per domain in an unsupervised manner.Source: A. Implementation Details · p. 15
  • Each domain uses an independent symmetric MLP autoencoder with 2560→512→128→512→2560 layers and ReLU, trained on reconstruction loss with Adam at 1 × 10−3.The objective minimizes the squared reconstruction error between the autoencoder output and its input feature.Source: A. Implementation Details · p. 15
  • Each stage’s rejection threshold is calibrated independently on its validation set to cover 95% of validation samples.Source: A. Implementation Details · p. 15
  • Each stage’s MoE has four LoRA experts (rank 16, alpha 32); the router selects the top two using normalized cosine similarity and softmax weights.Source: A. Implementation Details · p. 15
  • The load-balance loss weight is 0.01, and inference inserts one latent-memory segment after the prompt plus at most five additional segments.Each injected segment contains eight latent tokens.Source: A. Implementation Details · p. 15
  • Figure 2 visualizes frozen-reasoner hidden states at the first latent-memory insertion point for math, science, and code training prompts.Source: Figure 2 · p. 16
  • Using a fixed random seed, each domain contributes 3,000 training, 100 validation, and 500 test instances; the learning rate is 1 × 10−5 and warmup ratio 0.1.Maximum generation length is 2048 tokens. Training and evaluation use NVIDIA H200 GPUs.Source: A. Implementation Details · p. 16

B. Analysis

The analysis examines hidden-state signals for autoencoder routing, stage-wise routing statistics, expert usage, and the separation of learned latent memories across experts and domains.

B. Analysis

  • Separated prompt-end hidden-state clusters for math, science, and code provide distributional signals for router-group selection and out-of-distribution fallback.Figure 2 visualizes the frozen reasoner’s hidden states at the prompt end.Source: B. Analysis · p. 16
  • Figure 3 reports overall expert usage on the Math, Science, and Code test sets, supporting the analysis of spontaneous expert specialization.Source: B. Analysis · p. 16
  • On Math prompts, different experts’ latent memories occupy separated t-SNE regions, suggesting they learn distinct memory representations.Source: B. Analysis · p. 16
  • Corresponding experts’ latent memories form well-separated clusters across math, science, and code prompts, indicating that stage-isolated pools specialize by domain.Source: B. Analysis · p. 16
  • Table 3 summarizes autoencoder routing statistics across continual-learning stages, comparing prompt-end features with autoencoders for domains introduced so far and routing distant samples to the pretrained model.Source: B. Analysis · p. 16, Table 3 · p. 17

C. Additional Results on Training Order

Across the reported training orders, MoLEM has the highest final Average Performance after the full continual-learning sequence. The appendix reports results for two alternative orders alongside the paper’s default order.

C. Additional Results on Training Order

  • Across all three training orders, MoLEM consistently achieves the highest final Average Performance after the full continual-learning sequence.The reported advantage is stable under task-order variation.Source: C. Additional Results on Training Order · p. 17, Table 4 · p. 19, Table 5 · p. 19
  • The default order is Math → Science → Code; the additional orders are Code → Science → Math and Science → Code → Math.Source: C. Additional Results on Training Order · p. 17
  • Figure 3 shows overall expert usage on Math, Science, and Code test sets.Source: Figure 3 · p. 17
  • Figure 4 visualizes t-SNE latent memories from different experts on Math prompts, while Figure 5 shows corresponding within-pool index 0 experts on math, science, and code prompts.Source: Figure 4 · p. 18, Figure 5 · p. 18
  • Table 4 reports continual-learning results for Code → Science → Math, including per-test accuracy and three auxiliary metrics after each stage.Its caption notes that auxiliary metrics are undefined at the first stage.Source: Table 4 · p. 19
  • Table 5 reports continual-learning results for Science → Code → Math, including per-test accuracy and three auxiliary metrics after each stage.Its caption notes that auxiliary metrics are undefined at the first stage.Source: Table 5 · p. 19
Loading 2605.21951v1…