A Memory Mirage? Auditing and Improving Working Memory in Language Agents
Pranjal Sharma ⋅
Abstract
MemGen [1] proposes that a language model can carry state across turns by injecting learned latent vectors $z_{\text{mem}}$ into its hidden states, rather than relying on the state being re-exposed in the prompt. We evaluate this claim on the JHU Working Memory Benchmark [2], loading the official MemGen checkpoint (Qwen2.5-1.5B backbone) under a zero-text-leakage protocol: committed states are never re-provided in text, so correct behavior can only come from the injected representation itself. Working memory, in the cognitive-science sense, is the capacity to hold and manipulate state without external rehearsal [3]. We test three tasks constructed to isolate this capacity: Number Guessing, in which a secret $x^{*}\in[1,10]$ must be tracked across five yes/no queries; Yes-No Deduction, in which a rule chain $A{\to}B{\to}C{\to}D$ must be tracked across turns; and Math Magic, in which a running total is tracked through four arithmetic operations. The injected representation does not support state maintenance on any of the three tasks. The trained Weaver module outputs vectors with $\ell_2$-norm $\approx 220.0$, roughly $49\times$ the backbone's typical hidden-state norm ($\approx 4.5$). This magnitude saturates self-attention logits and collapses generation onto a single output: across $N{=}100$ trials the model answers "Yes'' to every query (mean rate 10.00/10). Scaling the injected vector by $\gamma{=}0.1$ removes this collapse but produces an equally rigid "No'' response instead (mean rate 0.00). Injection magnitude determines which fixed response the model converges to, independent of query content - a form of gain-dependent bias consistent with how large-magnitude activations are known to dominate attention computation in Transformers, and with the role gain modulation plays in biological working-memory circuits [4]. Baseline, MemGen, and the reduced-gain variant all fail Games 1 and 3 outright; only the Game 2 contradiction rate separates the four conditions, indicating that state maintenance is the primary bottleneck across every configuration tested. We introduce Latent Activation Splicing (LAS), which requires no additional training: we cache the base model's own key and value projections at the token where state is committed, and prepend them to the KV-cache prefix of later turns. LAS reduces Game 2 contradictions from 100\% to 30.0\% with no text-level leakage. The remaining error traces to RoPE mismatch: cached keys retain the rotary encoding of their original position, which distorts relative attention distances once spliced into a new position. We address this with RA-LAS, which inverse-rotates keys at extraction and reapplies rotation at splice time, and propose DCLGM, a gain controller that normalizes injected-vector magnitude against the model's own activation scale at each step: $$\gamma_t=\min\left(1,\alpha\frac{|h_t|2}{|W_r z{\text{mem}}|_2}\right)$$ These results indicate that latent-vector memory injection is structurally unstable without explicit magnitude control, and that position-aware activation caching recovers a substantial fraction of state-tracking performance with no additional training.
Chat is not available.
Successful Page Load