PaperlayerView exact revision on arXiv

Source-linked AI summary

AttriMem: Attribution-Guided Process Feedback for Agent Memory Construction

Qinfeng Li, Yuntai Bao, Xinyan Yu, Hongze Chen, Yanming Liu, Huifeng Zhu, Yier Jin, Jintao Chen, Wenqi Zhang, Xuhong Zhang

arXiv:2607.21106v3cs.AI

TL;DR

Long-horizon dialogue memory learning must identify which intermediate memory contents support an answer, yet existing feedback is coarse and process ground truth is difficult to specify. AttriMem derives local memory-token rewards from final-answer attribution and combines them with outcome reward. The authors report improved performance, memory quality, and RL stability on long-horizon dialogue question answering.

  • Problem

    Heuristic memory policies can be subjective and task-specific, while outcome- or action-level RL feedback does not pinpoint which intermediate memory contents support an answer.

    Source: §1 · p. 1

  • Method

    AttriMem estimates memory-token contributions by masking token subsets and measuring changes in final-answer score, then maps those scores to local process rewards alongside outcome reward.

    Source: §1 · pp. 2–3

  • Results

    AttriMem improves overall performance, intermediate memory quality, and RL training stability over outcome-only or action-level rewards on long-horizon dialogue question answering.

    Source: §7 · p. 11

  • Takeaways & Limitations

    Final-answer support can provide process feedback for memory construction without requiring a universal static ground-truth target for intermediate decisions.

    Source: §1 · p. 2

    Attribution quality can be challenged by exceptionally long contexts or weaker answer models, though the authors report limited impact on LongMemEval and say the framework can use more robust attribution methods.

    Source: §6 · p. 10

Abstract

from arXiv · show

Effective memory is crucial for LLM agents, yet constructing it effectively remains challenging. A memory-construction policy decides what information to extract, store, update, compress, or discard as interactions accumulate. Heuristic memory methods rely on subjective, task-specific rules, which can misalign with downstream objectives and limit cross-task adaptability. RL-based methods, by contrast, learn from task feedback but mainly use outcome- or module-level rewards. These coarse signals indicate task success but cannot identify which intermediate memory contents support the final answer, creating a fine-grained credit-assignment bottleneck. However, constructing such process feedback is prohibitively difficult because intermediate memory decisions lack unique ground-truth targets, while the appropriate credit varies with the agent's uncertain reasoning trajectory and therefore cannot be specified in advance. We propose AttriMem, an attribution-guided process-feedback framework for learning memory-construction policies with RL. AttriMem augments the global outcome reward with local rewards derived from token-level contributions to the final answer. Experiments on long-horizon dialogue question answering show that AttriMem outperforms retrieval-based, heuristic, and RL-based baselines, generalizes across benchmarks and answer models, stabilizes RL optimization.

1 Introduction

Long-horizon dialogue memory must preserve useful information scattered across sessions, but existing heuristic and RL approaches leave memory construction subjective or its intermediate decisions weakly supervised. AttriMem uses final-answer attribution to reward useful memory tokens and improve memory-policy learning.

1 Introduction

  • Heuristic memory methods use manually designed, task-specific rules that can misalign with downstream objectives and limit generalization.Table 1 identifies AttriMem as the only compared paradigm combining task feedback, feedback beyond outcomes, and action- and token-level guidance.Source: §1 · p. 1, Table 1 · p. 2
  • Existing RL methods provide sparse outcome- or action-level feedback, which does not identify how specific memory contents contributed to success or failure.Action-level rewards can miss the critical tokens or spans within a memory action.Source: §1 · p. 2
  • Intermediate memory decisions lack unique ground-truth targets because useful evidence depends on the agent’s uncertain reasoning trajectory.The paper notes that constructing process ground truth is prohibitively difficult.Source: §1 · p. 2
  • AttriMem masks subsets of intermediate memory tokens, uses changes in final-answer score to estimate token contributions, and maps those scores to local process rewards alongside the global outcome reward.This provides feedback on specific memory content within an action.Source: §1 · pp. 2–3
  • AttriMem’s token-level attribution enables more effective memory-policy optimization than outcome-only or action-level rewards, with stronger task performance and more reliable memory construction.The authors report extensive experiments supporting this comparison.Source: §1 · p. 3

2 Preliminaries

The preliminaries define long-term dialogue question answering as a persistent memory-construction problem separated from query-specific answering. They describe the memory modules and the SFT-plus-RL training setup used for memory policies.

2.1 Problem Setting and Scope

  • The task is to answer questions using information distributed across a conversation history, while a persistent memory policy incrementally updates memory without observing the future question.A fixed retriever and answer model handle query-specific reasoning, separating memory construction from answer-model adaptation.Source: §2.1 · p. 3

2.2 Memory Architecture

  • The memory architecture has core, episodic, semantic, and procedural modules, each with its own action space and action-specific prompts.Actions include operations such as ADD, UPDATE, and MERGE.Source: §2.2 · p. 4

2.3 Training a Memory Agent

  • SFT initializes the memory policy before RL optimizes memory decisions for downstream utility beyond imitation.The passage notes that smaller models particularly benefit from this initialization.Source: §2.3 · p. 4
  • Outcome-only GRPO normalizes final-answer rewards across sampled trajectories but gives a single feedback signal for each whole trajectory, which can make training inefficient and leave incorrect intermediate actions unpenalized when the outcome is correct.The retriever and answer model are separate from the trained memory policy.Source: §2.3 · pp. 4–5
  • AttriMem’s training overview pairs token-level attribution-derived local process rewards with the global outcome reward after memory construction and retrieval-based answering.The memory bank is built through extraction, compression, and update.Source: Figure 1 · p. 5
  • Context attribution estimates a source’s contribution to an answer by measuring the change in response score when that source is removed.The response-score measurement can use model output logits.Source: §2.3 · p. 5

3 Method

AttriMem isolates memory construction as the learned component in long-term dialogue question answering while keeping retrieval and answer generation fixed. It trains that policy with token-level attribution rewards alongside outcome feedback.

3.1 Problem Formulation

  • AttriMem constructs a memory store from dialogue history before a fixed retrieval-and-answering interface answers a later question.The method targets memory construction rather than the retriever or answer generator.Source: §3.1 · p. 5
  • The memory-construction policy decides what information to extract, update, merge, compress, and retain in the store.The paper frames learning this policy as the central question of the formulation.Source: §3.1 · pp. 5–6

3.2 Attribution-based Process Reward

  • GRPO combines the outcome advantage with token-level process advantages to update the memory-construction policy.The outcome advantage broadcasts final-answer quality across action tokens, while process advantages assign distinct attribution-based credit.Source: §3.2 · p. 6
  • Token contributions are estimated by removing memory-action tokens and measuring the resulting change in the generated answer’s score.AttriMem uses ContextCite, which fits a lightweight linear attribution model from randomly ablated contexts and evaluates up to 32 variants in parallel.Source: §3.2 · p. 6
  • Each token receives a distinct signed attribution reward, rather than sharing an aggregated reward for its action.The token scores are used as process rewards.Source: §3.2 · p. 6

4 Experiments

Experiments train the memory policy on LongMemEval and test zero-shot transfer on LoCoMo and PerLTQA, using matched MemBuilder and AttriMem setups. AttriMem leads the reported baselines, with token-level attribution also improving memory quality, training dynamics, and token efficiency.

4.2 Main Results

  • Token-level credit distinguishes exact times, durations, and quantities, while MemBuilder assigns one reward to the whole output.After RL, AttriMem retains critical evidence and removes distractions, whereas the MemBuilder-trained policy omits or distorts critical details.Source: Figure 2 · p. 8, §4.2 · p. 8
  • Claude, GPT, and Qwen3 can all use AttriMem's learned memory, indicating its utility is not tied to one answer model.The memory policy is optimized using feedback from a specific answer model.Source: Table 3 · p. 8, §4.2 · p. 8
  • AttriMem surpasses the strongest prior methods by 1.95 points on LoCoMo, 1.50 on LongMemEval, and 0.75 on PerLTQA.Trained only on LongMemEval, it also exceeds the strongest non-RL baselines on LoCoMo and PerLTQA.Source: §4.2 · p. 8

4.3 Ablation Study

  • Token-level attribution yields the strongest reward-granularity results, outperforming action-level attribution with and without SFT.Without SFT, it improves over the base model by 4.17, 5.50, and 2.26 points on LoCoMo, LongMemEval, and PerLTQA, respectively.Source: §4.3 · p. 8, Table 4 · p. 9

4.4 Further Analysis

  • AttriMem's reward rises steadily and stays stable, while MemBuilder plateaus after early gains and surges around step 320.Source: Figure 3 · p. 9, §4.4 · p. 9

4.3 Ablation Study

  • Without SFT, AttriMem gains 2.26–5.50 points over the base model, while MemBuilder gains only 0.10–2.00 points.The authors interpret this as fine-grained reward reducing reliance on SFT.Source: §4.3 · p. 9

4.4 Further Analysis

  • A GPT-based judge prefers AttriMem's intermediate memory to MemBuilder's on LoCoMo, LongMemEval, and PerLTQA, with win rates of 73.91%, 72.94%, and 77.78%.Source: §4.4 · p. 9
  • AttriMem lies on the LoCoMo performance–efficiency Pareto frontier, attaining the highest accuracy with token use comparable to MemBuilder.Its token consumption is substantially lower than Mem-T and GAM, while A-Mem uses fewer tokens at considerably lower accuracy.Source: §4.4 · p. 9

5 Related Work

Related memory approaches either use predefined retrieval or heuristic policies, or derive RL rewards from downstream outcomes. AttriMem's distinction from these RL approaches is token-level feedback tied to specific memory content.

  • Retrieval and heuristic memory policies are predefined and can mismatch downstream objectives.Heuristic methods use hand-crafted rules or scoring functions for writing, retaining, or compressing memory.Source: §5 · pp. 9–10
  • Outcome-redistributed RL signals indicate downstream success but provide limited evidence about which memory action, span, or token supports the answer.Mem-T redistributes final rewards to intermediate nodes in sampled memory-action trees.Source: §5 · p. 10
  • MemBuilder's synthetic QA yields denser action-level rewards than task success alone, but cannot identify which action, entry, or span deserves credit.AttriMem instead provides token-level feedback aligned with the content produced within memory actions.Source: §5 · p. 10

6 Discussion and Future Work

AttriMem’s discussion identifies attribution cost and quality as practical constraints, while framing the framework’s reward design and token-level representation as complementary choices with open questions.

6 Discussion and Future Work

  • Randomly masked variants add training computation, although batching them keeps the additional latency limited.All masked variants can be evaluated within a single batched model forward pass.Source: §6 · p. 10
  • Attribution quality may be more challenging in exceptionally long contexts or when using weaker answer models.LongMemEval is exceptionally long-context, but the experiments suggest the issue is limited; the text notes modern agent systems typically use capable answer models.Source: §6 · p. 10
  • Attribution-based process rewards complement rather than replace global outcome rewards, and the cited concerns primarily arise from the estimator.AttriMem is not tied to ContextCite and can use more robust long-context or white-box methods.Source: §6 · p. 10
  • The framework assumes memory-action outputs can be decomposed into tokens and aligned with attribution scores, though tokens may not be the optimal attribution unit.Answer-critical evidence may span multiple tokens or depend on interactions among them.Source: §6 · p. 11

7 Conclusion

AttriMem addresses fine-grained credit assignment in RL-based memory learning by tracing final-answer support to intermediate memory contents and turning token contributions into local process rewards. In long-horizon dialogue QA, the paper reports more effective supervision and gains in performance, memory quality, and RL stability.

7 Conclusion

  • In long-horizon dialogue question answering, attribution-guided process feedback improves overall performance, intermediate memory quality, and RL training stability.The conclusion characterizes it as more precise and effective supervision than outcome-only or action-level rewards.Source: §7 · p. 11
  • AttriMem traces final-answer support back to intermediate memory contents and converts token-level contributions into local process rewards.Source: §7 · p. 11
  • AttriMem is an attribution-guided process-feedback framework for RL-based memory learning.Source: §7 · p. 11

A ContextCite Methodology Details

The methodology adapts ContextCite to treat generated memory-action tokens as attribution sources. It estimates their contributions by masking token subsets, scoring a fixed answer, fitting a sparse linear surrogate, and mapping coefficients to local process rewards.

A ContextCite Methodology Details

  • AttriMem adapts ContextCite’s source-level attribution formulation by treating individual generated memory-action tokens as sources.Each source corresponds to a generated action-token position in the constructed textual memory context.Source: A ContextCite Methodology Details · p. 13
  • Random binary masks retain or remove action-token sources, preserving the original order of retained tokens; all sampled ablated contexts can be evaluated in one batch.Each source is independently retained with probability 1/2.Source: A ContextCite Methodology Details · p. 13
  • For each ablation, the method scores the same fixed generated answer under teacher forcing and regresses on its logit-transformed probability.Keeping the answer fixed isolates how included memory sources affect the model’s preference for the observed answer.Source: A ContextCite Methodology Details · p. 13
  • A sparse linear surrogate approximates the response score, and its fitted source coefficients serve as attribution scores.The sparsity penalty reflects the assumption that only a small subset of context materially affects a particular answer.Source: A ContextCite Methodology Details · pp. 13–14
  • A positive coefficient indicates retaining a source increases the fixed answer’s score under the surrogate; a negative coefficient indicates it decreases the score.Source: A ContextCite Methodology Details · p. 14
  • Each source coefficient becomes the corresponding action token’s process reward, with no action-level aggregation or length normalization.The attribution budget rises from 8 to 16 to 32 masks at RL steps 0, 80, and 200, reducing early cost while increasing reward fidelity as the policy strengthens.Source: A ContextCite Methodology Details · p. 14

B Implementation Details

Implementation uses a memory policy with a frozen retrieval-and-answering interface, trained with SFT and GRPO and augmented by token-level attribution rewards. The evaluation standardizes comparison conditions, tests transfer and answer-model variation, and describes its judge protocol and release plans.

B Implementation Details

  • Qwen3-4B constructs memory actions, while a frozen retriever and separate answer model provide feedback without being optimized as part of the policy.The memory store contains core, episodic, semantic, and procedural actions, with action-specific prompts and action spaces.Source: B Implementation Details · p. 14
  • Training uses a 10-epoch SFT stage followed by 400 GRPO update steps, with trajectory-level outcome advantage combined with attribution-based token-level process advantage.GRPO uses KL regularization against the reference policy, and answer rewards are normalized within each sampled group.Source: B Implementation Details · p. 14
  • Attribution computation keeps the answer fixed, scores randomly ablated contexts, fits a sparse surrogate, and maps coefficients to action tokens; the mask budget grows from 8 to 16 to 32.All masks for an estimate are evaluated in parallel as a batch.Source: B Implementation Details · p. 14
  • The comparison trains on LongMemEval and evaluates accuracy on LongMemEval, LoCoMo, and PerLTQA, including learned-memory comparisons across three answer models.For a controlled MemBuilder comparison, the designed difference is adding attribution-based token-level process rewards to the outcome reward.Source: B Implementation Details · pp. 14–15
  • The pairwise LLM judge sees the original dialogue, question, gold answer, and system outputs, with identities masked and both presentation orders averaged.Each pair is evaluated twice, once in each order.Source: B Implementation Details · p. 15
  • AttriMem training used eight NVIDIA H800 GPUs, with approximately 30 hours for SFT and five days for RL.Source: B Implementation Details · p. 15
  • The authors state that code and trained checkpoints will be released upon acceptance, and that the three evaluation datasets are publicly available.Source: B Implementation Details · p. 15

C Extended Related Work

Prior attribution approaches assign process credit at the step, turn, round, or role level, while AttriMem derives answer-conditioned token-level rewards for memory-action tokens. Related work also studies failure localization and attribution over long reasoning traces.

Related papers

Loading 2607.21106v3…