Technical Summary: ElasticMem

Problem Statement

Long-term memory is critical for enabling Large Language Model (LLM) agents to reason coherently across extended interactions, personalize responses, and reuse past experience. However, existing memory-augmented systems treat memory as a fixed resource, creating a mismatch between the query-dependent utility of memory and the static allocation of memory capacity.

Current approaches fall into two categories, both with significant limitations:

1. Text-Space Methods: These retrieve natural-language records and concatenate them into the context window. While interpretable, they incur substantial token overhead and are sensitive to noisy or redundant evidence, as irrelevant memories consume context tokens regardless of their utility.
2. Latent-Space Methods: These compress memories into continuous vectors or soft tokens to reduce textual costs. However, they typically rely on rigid retrieval mechanisms (e.g., fixed cosine similarity) and fixed-capacity interfaces where every retrieved chunk receives a uniform latent budget. This fails to distinguish between chunks containing crucial evidence versus those that are merely lexically similar but unhelpful.

The core challenge is that memory usefulness is highly query-dependent and task-dependent, yet retrieval and capacity allocation are often determined by fixed design choices rather than learned from task outcomes.

Methodology: ElasticMem

The authors propose ElasticMem, a framework that treats memory as a learnable, elastic latent resource. Instead of fixed retrieval and uniform allocation, ElasticMem learns to adaptively retrieve memories based on the reasoner's internal state and assign variable latent budgets to them based on their downstream utility.

The framework consists of four key stages:

1. Latent Memory Bank Construction

An offline latent memory bank is constructed once and kept frozen during training. Each memory chunk $m_i$ is encoded by a frozen base LLM encoder to produce:

• A retrieval key ($k_i$): The final hidden state of the chunk.
• A content cache ($C_i$): The last $N_c$ hidden states of the chunk.
  This allows for efficient retrieval without re-encoding the full corpus at inference time.

2. Query-Conditioned Retrieval

Unlike standard systems that use a separate, fixed retriever, ElasticMem derives the retrieval query from the reasoner's own hidden state.

• The LoRA-adapted reasoner samples a retrieval-control token $a$ given the query $q$.
• The reasoner processes $(q, a)$ to generate a query state $h_q$.
• Memories are retrieved based on the cosine similarity between $h_q$ and the cached retrieval keys $k_i$. This allows retrieval to adapt to the model's internal interpretation of the query rather than relying solely on surface semantic matching.

3. Elastic Latent Budget Allocation

After retrieving the top-$Z$ candidate chunks, a lightweight Transformer budget policy ($\pi_\phi$) assigns a variable number of latent tokens ($b_j$) to each chunk.

• The budget $b_j$ is sampled from a discrete distribution $\{0, 1, \dots, B_{max}\}$.
• A budget of $b_j=0$ allows the model to suppress retrieved chunks that are similar but unhelpful.
• The policy takes features including the query state, memory keys, retrieval scores, and learned rank embeddings, using self-attention to compare memories against one another before assigning capacity.
• Useful chunks receive larger budgets, while redundant or misleading ones are compressed or suppressed.

4. Latent Memory Injection and Optimization

Selected hidden states from the content caches are projected into soft memory tokens via a learned latent projector ($P_\psi$) and injected into the reasoner's input sequence.

• Training: The entire process (retrieval control, budget allocation, projection, and generation) is jointly optimized using Group-Relative Policy Optimization (GRPO).
• The reward signal (e.g., accuracy, task success) supervises not just the final answer, but also the retrieval decisions and budget allocations. This enables the model to learn which memories to retrieve and how much representational capacity each deserves based on actual task outcomes.

Key Contributions

1. Elastic Latent Memory: A novel paradigm shifting memory from a fixed textual or latent interface to a learnable, elastic resource where capacity is allocated dynamically based on utility.
2. Reasoner-State Adaptive Retrieval: A retrieval mechanism conditioned on the reasoner's internal state and a sampled control token, moving beyond rigid similarity matching.
3. Utility-Aware Budget Policy: A learned policy that assigns variable latent budgets to retrieved chunks, allowing the model to prioritize useful evidence and suppress noise.
4. End-to-End Optimization: A training framework that jointly optimizes retrieval, allocation, and generation via downstream task rewards, aligning memory management with task utility.

Experimental Results

ElasticMem was evaluated on MemorySuite, covering memory-intensive QA (PersonaMem, LoCoMo, LongMemEval) and embodied agentic decision-making (ALFWorld), using Qwen2.5-3B and Qwen2.5-7B backbones.

• Performance Gains:
  • QA Tasks: ElasticMem improved weighted average accuracy by 26.2% (Qwen2.5-3B) and 24.6% (Qwen2.5-7B) over the strongest baselines.
  • Agentic Tasks (ALFWorld): ElasticMem improved the weighted average success rate by 66.3% (Qwen2.5-3B) and 27.2% (Qwen2.5-7B).
• Efficiency: ElasticMem achieved the lowest token cost on ALFWorld among all compared methods, demonstrating that performance gains were not due to longer interaction trajectories or excessive memory use.
• Ablation Studies:
  • Budget Policy: The Transformer-based budget policy significantly outperformed random, uniform, and MLP-based allocation strategies, confirming the value of learning cross-memory interactions.
  • Retrieval: Query-state retrieval with a control token outperformed semantic and frozen-state retrieval, highlighting the importance of adapting retrieval to the reasoner's internal state.
  • Capacity Limits: A moderate per-chunk budget limit ($B_{max}=20$) yielded the best results, suggesting that overly large budgets can degrade performance by assigning capacity to less useful chunks.

Significance and Claims

The paper claims that ElasticMem addresses a fundamental mismatch in current memory systems: the static nature of memory allocation versus the dynamic nature of memory utility. By treating memory as a learnable resource, the framework enables agents to:

• Prioritize useful evidence and transferable plans even when they are not the most lexically similar to the query.
• Suppress noisy or redundant memories that would otherwise consume context or latent capacity.
• Achieve superior performance in both reasoning and decision-making tasks while maintaining high token efficiency.

The authors conclude that elastic latent memory is a promising direction for building more efficient, adaptive, and long-horizon LLM agents, moving beyond the limitations of fixed-context and rigid retrieval paradigms.