PaperlayerView exact revision on arXiv

Source-linked AI summary

UpliftMem: Learning Set-Level Uplift for Agent Memory Retrieval

Mengkun Liang, Haoran Qiang, Guannan Liu, Junjie Wu

arXiv:2609.36805v1cs.AI

TL;DR

LLM agents need to learn which memory sets improve execution, but ordinary retrieval provides feedback only for executed sets and evaluating alternatives costs additional rollouts. UPLIFTMEM learns set-level uplift against no-memory execution and selects probes by expected value for retrieval decisions. It achieves the best or joint-best success rates on the main evaluation sets across ALFWorld, WebShop, and BigCodeBench, with controlled evaluations also showing improved memory-use decisions and probe allocation.

  • Problem

    Memory retrieval needs outcome feedback to identify sets that improve execution, but ordinary retrieval labels only executed sets and evaluating alternatives requires costly rollouts.

    Source: §1 · pp. 1–2

  • Method

    UPLIFTMEM learns complete-set uplift relative to the same executor without memory and uses EVSI-guided probes to acquire feedback expected to improve retrieval decisions.

    Source: §1 · pp. 1–2

  • Results

    Across ALFWorld, WebShop, and BigCodeBench, UPLIFTMEM achieves the best or joint-best success rates on the main evaluation sets; matched-budget EVSI probing improves decision accuracy over random probing by 10.6 and 13.9 percentage points.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The results support treating the choice of which missing execution outcomes to observe as part of learning agent memory retrieval.

    Source: §6 · p. 10

    EVSI is derived under a correlated Gaussian surrogate for uncertainty about expected uplift, with execution-label variability modeled separately.

    Source: §4.3 · p. 5

Abstract

from arXiv · show

Large language model (LLM) agents reuse external memory to guide new tasks, but effective retrieval requires learning which memory sets improve execution. Such learning relies on costly outcome feedback: ordinary retrieval observes only executed sets, while evaluating alternatives requires additional rollouts. We introduce \textsc{UpliftMem}, which learns memory retrieval from set-level execution uplift relative to the same executor without memory. A theoretical analysis of how retrieval preferences restrict feedback coverage motivates targeted probing of alternative memory sets. Probe selection follows an expected value of sample information (EVSI) criterion, derived in closed form under a correlated Gaussian model, to allocate limited training rollouts according to their expected improvement in local retrieval decisions. The shared scorer is trained with a frozen executor and selects memory sets without test-time probes. Across ALFWorld, WebShop, and BigCodeBench, \textsc{UpliftMem} achieves the best success rates among evaluated baselines on the main evaluation sets. Controlled fixed-store and matched probe budget evaluations further demonstrate improved memory-use decisions and more effective use of execution feedback.

1 INTRODUCTION

Relevant memories can reduce an agent’s utility, and ordinary retrieval labels only the memory sets it executes, making alternative feedback costly. UPLIFTMEM learns set-level uplift and allocates additional rollouts to improve retrieval decisions, achieving strong benchmark and fixed-store results.

Motivation

  • Relevant retrieved memories may lower utility: in the backpack example, the agent succeeds without memory but follows a recommendation to buy an $89 item.Because only executed sets receive direct uplift labels, alternatives require additional rollouts, while the number of feasible combinations makes exhaustive labeling impractical under a limited budget.Source: §1 · pp. 1–2, Figure 1 · p. 1

Approach

  • UPLIFTMEM learns set-level execution uplift against no-memory behavior and uses targeted probes to acquire feedback for alternative sets expected to improve retrieval decisions.Its closed-form EVSI criterion uses a correlated Gaussian surrogate, and the shared value model evaluates candidate additions in the context already selected.Source: §1 · pp. 1–2

Results

  • Using Qwen3-4B and Qwen3-8B executors, UPLIFTMEM achieves the best success rates among evaluated systems on both ALFWorld splits, WebShop, and the overall BigCodeBench evaluation set.In fixed-store MemSyco-Bench evaluations, decision accuracy averaged over five memory stores and three scenarios assessing when to use memory rose from 41.7% to 66.5% with Qwen3-4B, and from 44.4% to 68.8% with Qwen3-8B.Source: §1 · p. 2

2 RELATED WORK

Prior work has developed reusable agent memories and retrieval methods that estimate downstream usefulness, including uplift and context-dependent effects. UPLIFTMEM treats the allocation of costly feedback for untested memory sets as part of learning retrieval.

2.1 EXPERIENCE-BASED AGENT MEMORY

  • Agent-memory research has shifted from retaining trajectories toward reusable procedural and reasoning guidance (Shinn et al., 2023; Zhao et al., 2024).Later systems organize, revise, and selectively maintain stored experience using subsequent interactions or learned management policies.Source: §2.1 · p. 2

2.2 LEARNING TO RETRIEVE AGENT MEMORY

  • Retrieval work estimates usefulness or uplift and models context-dependent effects, while UPLIFTMEM also allocates costly rollouts to improve decisions about untested sets.Prior work addresses retrieval utility and contextual attribution (Li et al., 2026a; Zhang et al., 2026; Qu et al., 2025; Tu et al., 2026; Xu et al., 2025b; Li et al., 2026b); selective-feedback research studies costly or incomplete observations (Swaminathan & Joachims, 2015; Houlsby et al., 2011; Kirsch et al., 2019).Source: §2.2 · pp. 2–3

3 PRELIMINARIES

The agent has a controller that manages external experience and retrieves memory sets for an executor. The formulation values retrieval by expected execution utility relative to no-memory behavior, with execution cost included in utility.

3.1 INTERACTIVE AGENT WITH EXTERNAL MEMORY

  • The interactive agent pairs a memory controller, which manages reusable experience, with an executor that interacts with the environment to complete tasks.Within the controller, the retriever selects memory sets to support execution.Source: §3.1 · p. 3
  • At each episode, the retriever selects a subset of the current memory pool for the executor, and the controller uses the resulting trajectory and feedback to update future memory candidates and the pool.Source: §3.1 · p. 3

3.2 PROBLEM FORMULATION

  • Execution utility incorporates benchmark-assigned execution cost because task performance alone does not distinguish equally successful executions with different resource requirements.The formulation denotes trajectory performance by R(τ) and cost by C(τ), capped at Cmax.Source: §3.2 · p. 3
  • Once performance meets threshold ρ, lower execution cost increases utility through an efficiency bonus weighted by γ.Source: §3.2 · p. 3
  • The retriever predicts expected utility gain for a task and memory set relative to no-memory execution, with the empty memory set allowed.The predicted set must satisfy retrieval constraints for the current pool.Source: §3.2 · p. 3

4 METHODOLOGY

UPLIFTMEM scores complete memory sets by their expected execution uplift relative to no-memory execution by the same executor. It combines sequential retrieval with targeted probes and trains a shared scorer from execution feedback while keeping the executor frozen.

4.1 DECISION VALUE OF MEMORY SETS

  • UPLIFTMEM scores complete memory sets by predicted execution uplift relative to the same executor without memory.The empty set is the no-memory reference, and each observation corresponds to the complete set used during execution.Source: §4 · p. 3, §4.1 · p. 4

4.2 SET CONSTRUCTION AND EXECUTION FEEDBACK

  • Similarity recall supplies candidate memories; conditional gains build sets, while only selected terminal sets get feedback and locally disfavored paths suppress observations of the sets they construct.At test time, width-two beam search can retain intermediate sets with negative conditional gains so later additions can reveal complementarity.Source: §4.2 · pp. 4–5

4.3 DECISION-AWARE PROBING

  • A correlated-Gaussian surrogate gives closed-form EVSI for local set comparisons, combining resolvable uncertainty with the current utility margin to estimate decision improvement without simulating outcomes.The method aggregates EVSI across local comparisons to select a probe batch under a limited rollout budget.Source: §4.3 · pp. 5–6

4.4 LEARNING FROM EXECUTION FEEDBACK

  • After value-only warm-up, uplift-based value and policy signals jointly train the shared scorer; signed, clipped uplift weights reinforce or suppress sampled retrieval actions.Validated and revised memories extracted from completed trajectories can contribute labels and enter the memory pool.Source: §4.4 · p. 6

5 EXPERIMENTS

The experiments compare complete agent workflows, isolate retrieval decisions with fixed memory stores, test learning and probe-acquisition components, and examine performance across retrieval budgets.

5.2 RQ1: END-TO-END AGENT PERFORMANCE

  • UpliftMem has the highest success rates on both ALFWorld splits and WebShop with both executors, and matches the highest overall BCB result with Qwen3-8B.The comparison evaluates complete agent workflows.Source: Table 1 · p. 7, §5.2 · p. 7

5.1 EXPERIMENTAL SETUP

  • Evaluation spans ALFWorld, WebShop, BigCodeBench, and dialogue-grouped MemSyco-Bench, using frozen Qwen3-4B and Qwen3-8B executors.MemSyco-Bench uses a 70:30 train/test split; the retriever is initialized from Qwen3-Reranker-0.6B and adapted separately for each executor.Source: §5.1 · p. 7

5.3 RQ2: CONTROLLED RETRIEVAL EVALUATION

  • With fixed stores and executors, retrieval improves accuracy across all five stores in Scope Control and Evidence Conflict, and in Valid Selection for all ten store–executor comparisons.Outdated-memory use decreases in nine Valid Selection comparisons.Source: §5.3 · pp. 7–8, Figure 3 · p. 8

5.4 RQ3: ABLATION STUDY

  • Across all three benchmarks and both executors, the full model has the highest overall success rates; on ALFWorld seen tasks, it exceeds warm-up-only and random-probing variants.For Qwen3-4B/Qwen3-8B, ALFWorld seen success is 76.86%/81.43% for the full model, 64.76%/66.19% for warm-up only, and 66.07%/72.85% for random probing.Source: Table 2 · p. 8, §5.4 · p. 8

5.5 RQ4: PROBE STRATEGY EFFICIENCY

  • On WebShop, probe strategies are compared across budgets of 20 to 200, tracking cumulative success rate and mean reward.The strategies are EVSI, random probing, Monte Carlo predictive standard deviation, and expected information gain.Source: §5.5 · p. 8, Figure 4 · p. 9

5.4 RQ3: ABLATION STUDY

  • EVSI-guided probing improves Decision accuracy over random probing by 10.6/13.9 percentage points and reduces sycophancy by 9.2/14.8 points for Qwen3-4B/Qwen3-8B.Removing either the policy or value learning objective also reduces Decision accuracy under both executors.Source: §5.4 · p. 8, Table 3 · p. 9

5.6 RQ5: SENSITIVITY TO RETRIEVAL BUDGET

  • Cross-model average gain over native retrieval remains positive at K = 3, 5, and 10: 14.2, 15.8, and 15.1 percentage points, respectively.Evidence Conflict gains increase as Scope Control gains decrease across retrieval budgets.Source: Figure 5 · p. 9, §5.6 · p. 9

6 CONCLUSION

UpliftMem learns set-level memory value against no-memory execution and allocates additional rollouts to improve retrieval decisions. The conclusion reports strong benchmark performance and better memory-use decisions under fixed-store and matched-budget evaluations.

6 CONCLUSION

  • UpliftMem addresses retrieval’s feedback coupling by learning set-level execution uplift against no-memory behavior and allocating rollouts by expected value for retrieval decisions.Retrieval determines both which experience is used and which execution outcomes become available as supervision.Source: §6 · p. 9
  • Across ALFWorld, WebShop, and BigCodeBench, UpliftMem achieves the best or joint-best success rates on the main evaluation sets.With fixed stores, Decision accuracy rises from 41.7% to 66.5% with Qwen3-4B and from 44.4% to 68.8% with Qwen3-8B.Source: §6 · p. 10

AI USE STATEMENT

  • Generative AI assisted manuscript organization, language editing, and literature searches during revision.Source: AI USE STATEMENT · p. 10

A THEORETICAL ANALYSIS

This section analyzes how selective retrieval limits feedback on unexecuted memory sets and derives criteria for choosing probes, then describes how execution and comparison feedback train the scorer and how beam retrieval works at test time.

A.1 SELECTIVE RETRIEVAL AND OBSERVATION COVERAGE

  • Direct coverage of a memory set sums the probabilities of all feasible construction paths, and early deficits can make even complementary sets unlikely to be observed.For a set B, the analysis bounds its execution probability by the sum over paths of exp(−d(a)/τtr); in the example, p{m1,m2} ≤ min{1, 2e^(−a/τtr)}.Source: §A.1 · pp. 13–14

A.2 IDENTIFICATION UNDER INCOMPLETE COVERAGE

  • When an alternative set is unobserved, its uplift sign can remain unidentifiable: ordinary retrieval has worst-case sign error at least 1/2(1−p̄B)^n, even with no-memory references.If pB = 0, the sign is unidentifiable over the stated outcome class.Source: §A.2 · p. 14

A.3 CHANGING POLICIES AND ADDITIONAL PROBE COVERAGE

  • With changing retrieval policies, a set executed with conditional probability at most εt per round has worst-case sign error at least 1/2 max{0, 1−∑t εt}; independent probes raise direct-coverage probability to 1−(1−pB)^n(1−rB).A set excluded throughout training has zero ordinary coverage.Source: §A.3 · p. 14

A.4 GAUSSIAN BELIEFS FOR PROSPECTIVE EXECUTION EVIDENCE

  • Prospective probe outcomes follow a Gaussian model combining scorer-feature uncertainty with execution noise, whose covariance captures shared no-memory references and coupled runs.The local surrogate is conditional on available data; its posterior covariance defines resolved uncertainty before the probe outcomes are observed.Source: §A.4 · p. 15

A.5 DECISION VALUE OF GAUSSIAN OBSERVATIONS

  • For a Gaussian margin with mean μj and resolved standard deviation sj, binary EVSI is sjϕ(|μj|/sj) − |μj|Φ(−|μj|/sj), the expected reduction in decision opportunity loss.EVSI is zero when sj = 0; the analysis derives the expression for comparisons between uncertain utilities and for fixed thresholds tied to a reference action.Source: §A.5 · p. 16

A.6 FINITE-FAMILY DECISIONS AND INFORMATION QUANTITY

  • For a finite family of executable sets, summing binary comparisons against one predicted-best incumbent upper-bounds selection EVSI, while mutual information ranks probes like EVSI only for a fixed comparison.Across comparisons, EVSI also reflects decision margin and utility scale; duplicate comparisons can count the same improvement more than once.Source: §A.6 · p. 17

A.7 BATCH ACQUISITION AND EVIDENCE TIMING

  • Batch acquisition accounts for overlapping evidence before execution, but marginal decision value need not decrease, and the probe-selection score is distinct from the future uplift change caused by training updates.The implemented H(Q) is computed before execution from current local comparisons; acquired outcomes then train the shared scorer.Source: §A.7 · pp. 17–18

B ALGORITHMIC DETAILS

  • The training algorithm maintains value observations, sampled paths, and an acquisition window, while observations retain reference identities for pairwise supervision and writing validation records enter the value buffer.OBSERVE returns an assisted trajectory, its no-memory reference, and their uplift.Source: B ALGORITHMIC DETAILS · p. 18
  • Training bootstraps memory and uplift observations, samples retrieval paths, records execution outcomes, and writes new memory from the resulting trajectories.The loop stores value and policy observations separately and accumulates retrieval windows for periodic comparison and probing.Source: B ALGORITHMIC DETAILS · pp. 18–19
  • At each acquisition interval, the algorithm builds local comparisons, selects probes within the probe budget, and adds their outcomes to the value dataset.Probe selection uses the current scorer and evidence at batch entry, with that belief held fixed during selection.Source: B ALGORITHMIC DETAILS · pp. 18–19
  • Scorer updates optimize the combined value and policy losses, then clear policy observations so paths sampled under earlier scorer states are not reused.The value loss uses eligible observations, and the policy loss assigns zero loss to the empty set.Source: B ALGORITHMIC DETAILS · pp. 18–19
  • At inference, UpliftMem recalls candidate memories and uses width-two beam retrieval, including the empty set, before executing the selected set.The algorithm returns the executor’s result for that selected set.Source: B ALGORITHMIC DETAILS · pp. 18–19

B.1 MEMORY SET SEARCH AND STOPPING

The appendix specifies constrained memory-set search, local probe-batch construction, and iterative writing with source-task validation. Search scores complete sets by predicted value, while probes target local decisions and memory candidates are admitted using validated uplift.

B.1 MEMORY SET SEARCH AND STOPPING

  • Search masks used or infeasible memories, permits STOP, deduplicates sets, and ranks complete sets by predicted value rather than path likelihood.Width-two beam search archives sets across depths and can retain a negative-gain intermediate set when later additions may expose complementarity.Source: §B.1 · p. 19

B.2 PROSPECTIVE PROBE BATCH CONSTRUCTION

  • Local probe candidates are complete retrieved sets and feasible one-memory additions, removals, and replacements; each executed task–set pair supplies one complete-set observation.The comparison family covers retaining or extending a set, choosing between replacements, and stopping versus extending.Source: §B.2 · p. 20

B.3 ITERATIVE MEMORY WRITING AND VALIDATION

  • Each candidate memory is validated alone on its source task; the first version with positive uplift is selected, or the highest-uplift version after at most three revisions.The selected version enters value supervision whether or not it is admitted to the pool; admission requires positive singleton validation uplift.Source: §B.3 · p. 21

C TRAINING DETAILS

Training details define a calibrated set scorer and value, sign, pairwise, and retrieval-path supervision. The experimental protocols describe benchmark partitions, including a dialogue-grouped MemSyco-Bench split.

C.1 SCORER PARAMETERIZATION

  • The scorer maps a task and complete memory set to a calibrated uplift estimate, with the empty set fixed at zero.Its trainable parameters include the affine calibration coefficients and LoRA adapter parameters.Source: §C.1 · p. 21

C.2 VALUE SUPERVISION

  • Value, sign, and pairwise supervision use uplift labels, retaining sign labels above a threshold and within-task pairs whose labels differ sufficiently.Pairwise preferences follow the larger observed uplift even when both labels have the same sign.Source: §C.2 · p. 21

C.3 OUTCOME-WEIGHTED RETRIEVAL-PATH LEARNING

  • Retrieval-path updates weight signed terminal uplift, excluding near-zero outcomes, clipping magnitude, and scaling negative feedback.Updates use only the latest batch of paths sampled under the current scorer.Source: §C.3 · p. 21

D.1 BENCHMARKS AND DATA PARTITIONS

  • MemSyco-Bench is split into 1,085 training and 465 test instances across five scenarios, with canonical-dialogue groups kept within one partition.The task-stratified 70:30 split uses random seed 42.Source: §D.1 · p. 22, Table 5 · p. 22

D.2 MEMORY SOURCES AND EVALUATION PROTOCOLS

The protocols distinguish complete-workflow evaluation from retrieval-only evaluation and specify utility, rollout accounting, and the models used. Retrieval-only comparisons hold the post-extraction store, executor, and evaluator fixed.

D.2 MEMORY SOURCES AND EVALUATION PROTOCOLS

  • Complete-workflow results include both memory-content and retrieval effects, whereas the controlled evaluation compares retrieval methods within the same store, executor, and evaluator.The controlled evaluation uses five MemSyco-Bench memory sources: NaiveRAG, Mem0, A-MEM, MemGPT, and LightMem.Source: §D.2 · p. 22

D.3 UTILITY CONFIGURATION AND ROLLOUT ACCOUNTING

  • Training cost counts reference, ordinary, probe, and writing-validation executions, counting a reused execution once; evaluation feedback does not trigger additional probes.New probe count measures sample allocation, while total training cost also includes required reference and validation executions.Source: §D.3 · p. 23

D.4 MODELS AND IMPLEMENTATION

  • The frozen task executors are Qwen3-4B and Qwen3-8B, each paired with an independently adapted scorer initialized from Qwen3-Reranker-0.6B.The scorers are trained with LoRA.Source: §D.4 · p. 23, Table 6 · p. 23

D.5 EVALUATION METRICS AND STATISTICAL REPORTING

The evaluation section defines task-specific metrics and reporting conventions, while the appendices specify prompt interfaces, information boundaries, and memory construction and revision workflows.

D.5 EVALUATION METRICS AND STATISTICAL REPORTING

  • ALFWorld, WebShop, and BigCodeBench use separate task metrics, while MemSyco-Bench reports scenario-specific accuracy, memory-use, sycophancy, and outdated-memory measures.Direction-aligned gains compare UpliftMem with native retrieval, equally weight the five memory sources, and average task scores for the overall summary.Source: §D.5 · p. 23
  • Configurations are trained once; benchmark evaluations use three runs, while main MemSyco-Bench results use five and ablations use three.Reported standard deviations reflect inference-run variability conditional on the trained checkpoint; reused Native and Full ablation results retain their original five-run statistics.Source: §D.5 · pp. 23–24

E PROMPTS FOR EXECUTION, VALUATION, AND MEMORY UPDATE

E.1 INFORMATION BOUNDARIES

  • Executor prompts exclude hidden goals, test code, expected answers, and evaluator internals, while writers ground guidance in visible evidence.Quantitative execution outcomes supervise the value scorer.Source: §E.1 · p. 24

E.2 EXECUTOR PROMPTS

  • The three executors have distinct action and generation contracts, but paired memory-assisted and no-memory runs use the same executor instructions and examples.The prompts specify the message fields through which retrieved memory enters the task.Source: §E.2 · p. 24

E.3 SET-VALUE SCORING PROMPT

  • The set-value scoring template presents the complete candidate set and produces yes/no logits calibrated in Appendix C.1.Source: §E.3 · p. 27

E.4 MEMORY CONSTRUCTION PROMPTS

  • Memory writers use task-specific representations, and a successful trajectory can still have negative retrieval uplift relative to the no-memory reference.The appendix gives writer templates, builder fields, and output contracts for interactive control, shopping, and code generation.Source: §E.4 · p. 27

E.5 REVISION AND ADMISSION

  • The revision loop sends the latest candidate, its realized trajectory, and validation feedback to the writer before environment-specific output contracts are applied.Source: §E.5 · p. 30

F ADDITIONAL PROBE STRATEGY EFFICIENCY RESULTS

EVSI achieves the highest cumulative success rate at every nonzero probe budget on BigCodeBench and leads or ties most ALFWorld cumulative-coverage comparisons. The section reports both cumulative coverage and checkpoint success across probe budgets.

F ADDITIONAL PROBE STRATEGY EFFICIENCY RESULTS

F.1 ALFWORLD

  • At 200 probes, EVSI cumulative success is 84.3% on seen and 81.3% on unseen ALFWorld tasks, versus 81.4% and 79.9% for random probing.At 50 probes, EVSI reaches 72.1% on seen and 65.7% on unseen tasks, versus 66.4% and 62.7% for random probing; checkpoint rates fluctuate more, with EIG slightly exceeding EVSI on seen tasks at 200 probes.Source: Figure 7 · p. 31, §F.1 · p. 31

F.2 BIGCODEBENCH

  • 51.8% cumulative success at 200 probes on BigCodeBench, versus 47.4% for random probing and 50.3% for MC-STD; EVSI leads at every nonzero budget.At the final budget, EVSI checkpoint success is 37.7%, versus 35.1% for MC-STD and 33.3% for random probing.Source: Figure 8 · p. 31, §F.2 · p. 31

G ADDITIONAL MEMSYCO RESULTS

This section points to absolute per-store MemSyco-Bench scores and per-system ablations at K=5 for both Qwen backbones.

G ADDITIONAL MEMSYCO RESULTS

Loading 2609.36805v1…