MemAgent: Reshaping Long-Context LLM with Multi-Conv RL-based Memory Agent

TL;DR

Long-context LLMs break down on truly long inputs: attention is quadratic, extrapolation tricks degrade, and bolt-on memory modules disrupt the normal generation path. MemAgent reframes the problem as an agent loop: the model reads the document one chunk at a time, and after each chunk it overwrites a small fixed-length “memory” (1024 tokens of plain text) with an updated set of notes. Because the memory never grows, the active context window is constant, so cost is strictly linear in document length and the model can stream a document of any size. The hard part — teaching the model what to keep and what to throw away — is solved with reinforcement learning: only the final answer gets a reward, and that reward is pushed back through every memory-update step using a new “Multi-Conv DAPO” algorithm. The headline result: a 7B/14B model with an 8K window, trained only on 32K-token documents, holds ~75-78% accuracy out to 3.5 million tokens, while 1M-context baselines collapse to 0% well before their advertised limit.

Problem & Motivation

The concrete pain: you want an LLM to reason over an entire book, a massive codebase, a multi-year chat history, or an agent’s accumulated memory — inputs that are tens of thousands to millions of tokens. Three existing families all fall short:

1. Length extrapolation (RoPE tricks + continued pretraining) — methods like NTK, PI, YaRN, DCA reshape positional embeddings to stretch the window. But attention is still O(n²), so processing is slow, and accuracy visibly degrades once you push past the trained length. In this paper’s own table, a Qwen2.5-1M model that advertises a 1-million-token window drops to 0% accuracy at 896K — far short of its theoretical ceiling.
2. Efficient attention (sparse / linear / SSM) — sliding windows, linear attention, and state-space models get to O(n), but they typically require training a new architecture from scratch, sparse patterns are hand-designed heuristics, and linear attention is awkward to train in parallel. Worse, their compressed “memory” lives in opaque hidden-state space — you can’t read it or fix it.
3. Context compression / external memory plugins — condense tokens or attach a memory database. These struggle to extrapolate, and the extra modules disrupt the standard generation process, hurting compatibility and parallelization.

The authors frame the real goal as a trilemma every solution must satisfy at once: (1) arbitrary input length, (2) no performance cliff as length grows, (3) linear decoding cost. Nothing above hits all three. Their inspiration is human reading: we don’t memorize every word of a book — we take terse notes on what matters and discard the rest. MemAgent operationalizes exactly that.

What’s New (Core Contribution)

Three genuine contributions, not repackaging:

• A token-space, fixed-length, overwrite-based memory loop. Before: memory was either the full KV-cache (grows unbounded), an opaque compressed feature vector, or an external module that breaks the generation path. Now: the “memory” is just ordinary tokens — a human-readable note — that the same base LLM rewrites after every chunk. The window stays constant, so cost is linear and the base decoder is completely unmodified (no new attention kernels, no positional-embedding surgery).
• End-to-end RL training of a multi-context agent workflow (“Multi-Conv DAPO”). Before: multi-turn agent RL (Search-R1, Agent-R1, GiGPO) concatenated all turns into one conversation with attention masking, or used sliding windows — fine for tool-calling but not for a workflow that spawns many independent conversations per query. Now: each chunk’s read-write step is its own independent conversation, the reward from the final answer is broadcast as the advantage to all the intermediate conversations, and the policy gradient is computed over a (group, conversation, token) tensor. This is what lets you train “keep the useful notes” behavior even though only the last step produces a checkable answer.
• A clean theoretical reframing. They show MemAgent is mathematically a latent-variable recurrent factorization of the language model: the memory is a latent state, each chunk is a read step p(c_k | m_{k-1}) and a write step p(m_k | c_k, m_{k-1}). This turns a transformer into an RNN whose state size you choose — and explains why RL is required (the discrete overwrite of latent tokens can’t be learned by plain backprop).

How It Works (Technically)

The loop, in words

At every step the model sees exactly two things: the next chunk of the document and the current memory. It produces an updated memory (the “write”). When the document runs out, a separate answer step consults only the question plus the final memory and emits a \boxed{} answer. Two prompt templates do all the work (Table 1 in the paper):

• Context-processing prompt: “Here is a problem, a section of the article, and your previous memory. Update the memory with new info that helps answer the problem, keep relevant prior details.” → outputs Updated memory:
• Answer prompt: “Here is a problem and your memory. Answer and put it in \boxed{}.” → outputs the answer.

That’s the whole inference mechanism. The cleverness is entirely in training the model to write good memories.

The math, demystified

Why it’s linear. The memory has a fixed size M (1024 tokens). Each chunk has size ≤ C (5000 tokens). Per step you run a normal transformer over a context of size C + M — a constant. With K = N/C chunks for a document of N tokens, total cost is O(K · (C+M)) = O(N). No term grows with total length because the window the model actually attends over never grows. Compare quadratic attention, where the cost of the last token alone is O(N).

The recurrent factorization (Eq. 8). A standard LLM factorizes p(x_1:N) = ∏ p(x_n | x_1:n-1) — every token conditions on all prior tokens, which is the quadratic bottleneck. MemAgent inserts a latent memory sequence m_1:K-1 and rewrites it as:

p(x_1:N) = Σ_{m_1:K-1}  ∏_k  p(c_k | m_{k-1})  ·  p(m_k | c_k, m_{k-1})
                              └── read ──┘       └──── write ────┘

In plain English: the probability of the whole document equals reading each chunk given the previous note, then writing the next note given the chunk and the previous note. Set m_0 = ∅. This is literally the equation of a recurrent net — the transformer becomes an RNN whose hidden state is a chunk of readable text whose length you pick. The key consequence: the memory is in token space, so it’s inspectable and even editable, unlike the opaque latent vectors of linear-attention models.

Why RL is essential. The memory tokens are a latent variable updated by a discrete overwrite. You only get a training signal at the very end (was the final answer correct?). Plain backpropagation can’t flow through “the model chose to write these particular discrete tokens as its note.” RL handles exactly this: treat each read-write-read transition as an action, and reward memories that lead to a correct final answer. That bridges explicit supervision (answers) and implicit structure (good notes).

The RL objective — building up from GRPO. Start with GRPO (Group Relative Policy Optimization), the algorithm behind DeepSeek-R1’s reasoning training. For one input, the old policy samples a group of G candidate outputs. Each gets a scalar reward R_i. The “advantage” — how much better than average this output is — is just the group-normalized reward:

Â_i = (R_i − mean(R)) / std(R)          # Eq. 1

No separate value/critic network needed — the group average is your baseline. The policy is then updated with a clipped objective (Eq. 2) that increases the probability of high-advantage outputs while clipping how far the update can move, plus a KL penalty β·D_KL(π_θ || π_ref) keeping it near the reference model so it doesn’t degenerate.

The MemAgent extension (Multi-Conv DAPO). The wrinkle: one query produces multiple conversations (one per chunk + one for the answer), not one. You can’t just attention-mask them into a single sequence. So:

• Compute one outcome reward R_i per sample, derived from the final (answer) conversation.
• Broadcast the group-normalized advantage to every conversation from that sample (Eq. 4): Â_{i,j,t} = R_i − mean(R). Following Dr.GRPO, they drop the std-normalization (it can bias toward easy questions).
• Extend the loss from a (group, token) average to a (group, conversation, token) average (Eq. 5), borrowing DAPO’s token-level averaging and asymmetric clip bounds ε_low, ε_high.

The intuition that matters: if the final answer was right, reinforce the token choices in all the memory-writing steps that led there; if wrong, suppress them. The model thereby learns, across many episodes, which note-taking habits produce correct answers downstream.

Reward modeling. A rule-based verifier, no reward model. For single-answer QA with equivalent ground truths: R = max over y in Y of is_equiv(y, ŷ) (Eq. 6) — credit if the prediction matches any acceptable answer. For “list-all” tasks (e.g., multi-value needle-in-haystack): R = |{y in Y : y in ŷ}| / |Y| (Eq. 7) — partial credit for the fraction of required items recalled.

Architecture & data flow

flowchart LR
  D[Long document N tokens] --> SPLIT[Split into chunks c1..cK]
  SPLIT --> L1[LLM read+write: chunk c1 + empty memory]
  L1 -->|memory m1| L2[LLM read+write: chunk c2 + m1]
  L2 -->|memory m2| LD[... iterate over all chunks ...]
  LD -->|memory mK| ANS[LLM answer step: question + mK]
  ANS --> OUT["boxed answer ŷ"]
  Q[Question] --> L1
  Q --> L2
  Q --> ANS
  OUT --> V[Rule-based verifier vs ground truth]
  V -->|reward R| RL[Multi-Conv DAPO: broadcast advantage to ALL read/write steps]
  RL -.updates.-> L1
  RL -.updates.-> L2
  RL -.updates.-> ANS

Schematic of the streaming memory loop: each chunk is read alongside the fixed-size memory, which is overwritten before the next chunk. Watch the context window (chunk + memory) stay constant size while the document scrolls past — that constant box is why cost is linear. The final step answers from memory alone.

Schematic cost-vs-length curves: quadratic attention (full context) versus MemAgent's linear streaming. Toggle to see why a fixed window defeats the "performance cliff" that hits extrapolated baselines. Illustrative, not the paper's exact numbers.

The algorithm, simplified

# Inference: stream a document of any length through a fixed window.
def memagent_answer(llm, question, document, chunk_size=5000):
    memory = ""                                   # latent state, in plain tokens
    for chunk in split(document, chunk_size):     # K = len(document)/chunk_size steps
        memory = llm(CTX_PROMPT.format(           # read chunk + old memory -> overwrite memory
            problem=question, memory=memory, section=chunk))
        memory = truncate(memory, max_tokens=1024)  # memory NEVER grows -> O(1) per step
    return llm(ANS_PROMPT.format(                 # answer from memory only, ignore raw doc
        problem=question, memory=memory))         # -> "\boxed{...}"


# Training one query with Multi-Conv DAPO (the contribution).
def train_step(policy, question, gold, document, G=16):
    samples = []
    for _ in range(G):                            # GRPO: a GROUP of rollouts per query
        convs = []                                # each rollout = many independent conversations
        memory = ""
        for chunk in split(document):
            out = policy.sample(CTX_PROMPT.format(problem=question, memory=memory, section=chunk))
            convs.append(out); memory = truncate(out, 1024)   # one conversation per chunk
        answer = policy.sample(ANS_PROMPT.format(problem=question, memory=memory))
        convs.append(answer)
        R = rule_verifier(extract_boxed(answer), gold)        # reward ONLY from final answer
        samples.append((convs, R))

    Rs = [R for _, R in samples]
    baseline = sum(Rs) / len(Rs)                  # group mean = critic-free baseline (Dr.GRPO: no /std)
    loss = 0
    for convs, R in samples:
        adv = R - baseline                        # SAME advantage broadcast to every conversation
        for conv in convs:                        # <-- the novelty: reward the note-taking steps too
            loss += clipped_pg(policy, conv, adv) # token-level clipped PG (DAPO) + KL to ref model
    return loss / total_tokens(samples)           # averaged over (group, conversation, token)

Built on Prior Work

Results & Evidence

Setup. Base models: Qwen2.5-7B-Instruct and -14B-Instruct, trained with verl. Deliberately constrained to an 8K window (1024 query + 5000 chunk + 1024 memory + 1024 output). Training data: ~32K-token synthetic multi-hop QA from HotpotQA. Test: RULER HotpotQA from 7K up to 3.5M tokens, plus OOD RULER tasks (needle-in-haystack variants, variable tracking, frequent-words, SQuAD QA).

Headline numbers (RULER QA accuracy %):

• RL-MemAgent-14B: 83.6% at 7K → 78.1% at 3.5M. Essentially flat across a 500x length increase.
• RL-MemAgent-7B: 82.0% at 7K → 71.1% at 3.5M.
• Qwen2.5-14B-1M (advertised 1M window): 60% at 7K → 0% at 896K.
• DS-Distill-Qwen-32B (reasoning model, 128K): 70% at 7K → 7% at 896K.
• QwenLong-L1-32B: strong to ~60K, then collapses (13% at 448K).

So MemAgent both starts higher and barely decays, while every baseline cliffs out — often hitting 0% before its nominal limit.

Ablations that matter:

• Memory alone (no RL) helps but isn’t enough: a model given the memory loop without RL beats vanilla truncation and degrades more gracefully, but still declines with length. RL is what flattens the curve. This is the key causal claim: structure (memory) + learning (RL) are both necessary.
• OOD generalization: trained only on HotpotQA-style QA, it still transfers to other RULER task families.

Caveats / what this does NOT establish:

• Tasks are extraction/QA-flavored. RULER and needle-in-haystack reward finding and recalling specific facts. The fixed 1024-token memory is plausibly too small for tasks needing global synthesis (summarize a whole novel, reconcile contradictions spread across the document, multi-document reasoning where everything matters). The paper doesn’t stress-test that regime.
• Single benchmark family (RULER) + synthetic data. No real-world long-doc tasks (legal review, full-repo code reasoning, long-horizon agent memory) in the headline table.
• Reward is rule-based exact-match. Works for QA with checkable answers; doesn’t cover open-ended generation where you’d need a reward model.
• Latency tradeoff unstated. Linear compute is great, but you now make ~K sequential LLM calls per query (5-7 turns even on training-length docs, far more on 3.5M). That’s serialized wall-clock latency the table doesn’t surface.
• No comparison to RAG. The most obvious production baseline for “answer a question over a huge corpus” — retrieval + a normal LLM — isn’t in the comparison.

How You’d Use It

This maps cleanly onto things you already build:

• Long-document Q&A as a service without a giant-context model. You can offer “ask anything about this 500-page contract / this 2,000-page deposition bundle” using a cheap 8K-32K model instead of paying for a frontier 1M-context API. The memory loop is just an orchestration pattern around a small model.
• Agent long-term memory. This is essentially a learned memory-compaction policy for an agent. Instead of dumping a growing transcript into context (and hitting limits / paying quadratically), the agent maintains a fixed-size, human-readable scratchpad it rewrites each turn. You can drop the inference loop into an existing agent today, no training required.
• Auditable summaries. Because the memory is plain text, you can show the client what the system decided to remember at each step — a real differentiator versus opaque RAG or black-box long-context. Good for regulated domains.
• Cost moat. Linear cost + small base model = dramatically cheaper unit economics on long-doc workloads than calling a frontier long-context model per query. That’s a margin story for an AI-services offering.

The honest read: the inference pattern (chunk → rewrite memory → answer) is immediately reusable and gives you most of the structural benefit even with an off-the-shelf model. The RL training is what flattens the accuracy curve, and that’s a real (GPU-heavy) project.

Build Your Own (Minimal Recipe)

80% of the value, no training:

1. Chunker — split the document into ~4-5K-token pieces.
2. Two prompts — the context-update prompt and the answer prompt, verbatim from the paper’s Table 1 templates.
3. Loop — iterate chunks, feeding (question, current_memory, chunk); overwrite memory with the output each time; truncate memory to a fixed budget (~1K tokens).
4. Answer — final call with (question, final_memory).

That’s a one-evening build on any instruct model (qwen2.5-7b-instruct, llama-3.1-8b, or an API). You’ll get graceful-degradation behavior immediately — the “memory without RL” ablation, which already beats truncation.

The hard 20% (the RL):

• Stand up Multi-Conv rollouts. You need a trainer that can sample a group of rollouts where each rollout is multiple conversations, then assemble a (group, conversation, token) loss. The paper uses verl; that’s the library to reach for (it has GRPO/DAPO and the rollout machinery). This is the genuinely tricky part — most RLHF stacks assume one conversation per sample.
• Reward plumbing. A rule-based verifier (exact/equivalence match for QA; set-recall for list tasks). Easy logic, but you need ground-truth-labeled long-doc QA data.
• Data synthesis. Build training docs by planting “golden” paragraphs (containing answers) among distractors, à la RULER, and filter out questions the base model already answers without context.

Reach for: verl (RL), Qwen2.5-Instruct (base), HotpotQA / RULER (data + eval), vLLM for fast rollout sampling.

How to Improve It

Limitations as leverage — concrete, testable directions:

1. Variable / hierarchical memory. Fixed 1024 tokens is a guess. Let the model request more memory for dense sections, or maintain a two-tier memory (a long-term compressed note + a short working note). Test whether this rescues synthesis-heavy tasks where one flat note is too small.
2. Reward intermediate memories directly. Right now only the final answer is rewarded; the memory steps get a broadcast signal. Add a cheap auxiliary reward (e.g., does the memory contain the golden facts, scored by a verifier or a small judge model) to give denser, faster-learning credit assignment.
3. Parallelize the read pass. The K sequential calls are a latency wall. Explore a tree/map-reduce variant: summarize chunks in parallel into partial memories, then a learned merge step combines them — trading strict recurrence for wall-clock speed.
4. MemAgent + RAG hybrid. Use retrieval to pre-select candidate chunks, then run the memory loop only over those. Cuts the number of LLM calls on huge corpora and directly addresses the missing-RAG-baseline gap.
5. Open-ended generation rewards. Replace exact-match with a reward model or LLM-judge so the same training works for long-doc summarization and drafting, not just extractive QA.
6. Editable-memory tooling. Since memory is plain text, expose a human-in-the-loop checkpoint: let a user correct the running note mid-stream. A genuinely sellable feature for high-stakes review workflows, and a free supervision signal.

Glossary

• Long-context trilemma — the three properties a good long-context method must satisfy together: arbitrary length, no accuracy cliff, linear cost.
• Memory (here) — a fixed-length sequence of ordinary text tokens (1024) that the model overwrites after each chunk; the agent’s note-to-self.
• Overwrite strategy — replacing the entire memory each step (vs. appending), which is what keeps the window constant-size.
• Chunk — a contiguous slice of the document (~5000 tokens) processed in one step.
• Extrapolation — running a model on inputs longer than it was trained on; usually causes accuracy to drop.
• RoPE / NTK / PI / YaRN / DCA — rotary positional embeddings and various tricks to rescale them for longer windows.
• RL (reinforcement learning) — training by reward signal rather than labeled targets; here, “did the final answer turn out correct?”
• Policy — the model being trained, viewed as something that chooses actions (next tokens / memory updates) to maximize reward.
• Reward — scalar score for an output; here computed by a rule-based verifier comparing the boxed answer to ground truth.
• Advantage — how much better an output is than the group’s average; tokens with positive advantage get reinforced.
• GRPO (Group Relative Policy Optimization) — RL method that uses the average reward of a sampled group as the baseline, avoiding a separate critic network.
• DAPO — a GRPO refinement with token-level loss averaging and decoupled clip bounds; MemAgent extends its loss to multiple conversations.
• Dr.GRPO — variant that removes std-normalization of the advantage to avoid biasing toward easy samples.
• KL penalty — a term keeping the trained policy close to a frozen reference model so it doesn’t drift into degenerate text.
• Clipping — capping how far one update can change a token’s probability, for stable RL.
• RLVR — Reinforcement Learning with Verifiable Rewards; using checkable (rule-based) outcomes as the reward signal.
• Latent variable — an unobserved intermediate (here, the memory) that the joint probability is factorized through.
• RULER / HotpotQA / NIAH — long-context benchmark, multi-hop QA dataset, and the “needle in a haystack” recall paradigm, respectively.
• verl — the open-source RL training framework the authors used to implement multi-conversation rollouts.