WMF-AM: Probing LLM Working Memory via Depth-Parameterized Cumulative State Tracking

Dengzhe Hou; Lingyu Jiang; Deng Li; Zirui Li; Fangzhou Lin; Kazunori D Yamada

cs.AIMar 28, 2026

WMF-AM: Probing LLM Working Memory via Depth-Parameterized Cumulative State Tracking

Authors: Dengzhe Hou, Lingyu Jiang, Deng Li, Zirui Li, Fangzhou Lin, Kazunori D Yamada

Organizations: Tohoku University · Lappeenranta-Lahti University of Technology LUT · Texas A&M University · Worcester Polytechnic Institute

Abstract

Existing large language models (LLMs) evaluations use fixed-difficulty benchmarks that cannot adapt as models improve, and rarely isolate specific cognitive processes. We introduce Working Memory Fidelity-Active Manipulation (WMF-AM), a probe of cumulative state tracking, the ability to maintain and update intermediate results across K sequential operations within a single query, without a scratchpad. Unlike multi-step agent benchmarks that stress task orchestration, WMF-AM isolates within-pass cumulative load by parameterizing depth K. The core probe uses arithmetic accumulation on 28 models from 12 families (0.5B to frontier); a matched non-arithmetic extension (permissions, schedules, inventories) confirms the design generalizes beyond arithmetic. Three construct-isolation ablations confirm that cumulative load, not arithmetic skill or entity tracking, drives difficulty. We release WMF-AM as a lightweight, recalibratable diagnostic for characterizing where models degrade under cumulative load. Code and data can be accessed at https://github.com/dengzhe-hou/WMF-AM

Figures & tables

Figure 1: WMF-AM framework. (a) Cognitive analogy: a model maintains and updates a running state across $K$ sequential operations and reports only the final answer, without scratchpad. (b) Probe design: the input prompt specifies an initial state and $K$ cumulative updates; the LLM must track hidden internal state as $K$ increases. (c) Radar profiles for 10 representative models at $K{=}3/5/7$ with agent battery and yoked control scores, illustrating cross-model variation.

Phase	Probe	Depths	Items	Scoring	Key Result
1. Outcome	100-item battery	—	100	Exact match (0/1)	—
2. WMF-AM	State tracking	$K{=}3,5,7$	15/depth	Exact match on final state	$\tau{=}0.595$ †
3. Yoked	Cancellation ops	$K{=}2,4,6,8,12$	20/depth	Exact match on initial state	$\tau{=}0.381$
4. Template	3 prompt wrappers	$K{=}3,5,7$	15/depth	Rank stability ( $\tau$ )	$\tau{=}0.631$
5. Agent	10 multi-step tasks	—	10	Deterministic (0–1)	criterion
6. K-sweep	Extended depths	$K{=}3\text{--}100$	20/depth	Exact match on final state	$K_{\text{crit}}$ 1.3–55.3

Table 1: Study design and key results (

N{=}28

). All phases: temperature

=0

, greedy decoding. WMF-AM uses 4 random seeds; all other phases use a single seed.

Figure 2: WMF-AM vs. Agent Battery Score ( $N{=}28$ , 12 families). WMF-AM predicts downstream agent performance ( $\tau{=}0.595$ , $p{<}0.001$ ). Blue circles = standard models; red squares = LRM (reasoning) models; orange diamonds = LRM-distill models; black edge = API models. All 28 models are labeled. Note: DeepSeek-R1 (671B, “R1-full”) achieves perfect WMF-AM (1.000) but low agent score (0.50), discussed in Section 5 . Claude and o3-mini overlap at (1.0, 0.9); points are slightly jittered for visibility.

Predictor	$\tau$	$p$	$N$	Ceiling
MMLU (published, 5-shot)	0.480	$<$ 0.001	27	0%
GSM8K (published, CoT)	0.588	$<$ 0.001	20	35%
WMF-AM (this paper)	0.595	$<$ 0.001	28	14%
Yoked cancellation accuracy	0.381	0.007	28	—

Table 2: Predictor comparison: WMF-AM vs. alternative predictors of ABS. Published MMLU/GSM8K scores from official model cards and technical reports (5-shot MMLU, CoT GSM8K).

N

varies because not all models have published scores. See Appendix J for model-matched (

N{=}17

) comparison.

Figure 3: K-sweep analysis (

N{=}28

). (a) K-degradation curves: accuracy vs. depth

K

for all 28 models (gray), with five representative models highlighted (Claude-Sonnet-4, o3-mini, DeepSeek-V3, GPT-4o, DeepSeek-R1). Standard models show sigmoid-cliff collapse; DeepSeek-R1 shows non-monotonic recovery. (b)

K_{\text{crit}}

vs. Agent Battery Score (

\tau{=}0.171

p{=}0.23

, n.s.): collapse threshold does not predict agent performance. Faded points indicate unreliable sigmoid fits (

R^{2}{<}0.90

Model	Type	Sup	Unsup	$\Delta$
GPT-4o	Std	0.9	0.3	0.6
Claude Sonnet 4	Std	0.9	0.4	0.5
DeepSeek-V3	Std	0.9	0.4	0.5
qwen2.5:32b	Std	0.9	0.4	0.5
qwen2.5:14b	Std	0.9	0.5	0.4
R1 full	LRM	0.6	0.2	0.4

Table 3: Load-shift intervention (selected models, Sup

\geq 0.4

\Delta

= supported

-

unsupported. Full 28-model table in Appendix.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

$K$	Mean acc	Range	$\tau$	$p$
3	0.620	$[0.050,1.000]$	0.60	$<$ 0.001
5	0.454	$[0.000,1.000]$	0.63	$<$ 0.001
7	0.327	$[0.000,1.000]$	0.50	$<$ 0.001
10	0.304	$[0.000,1.000]$	0.43	$<$ 0.01
15	0.262	$[0.000,1.000]$	0.22	0.13
20	0.202	$[0.000,1.000]$	0.44	$<$ 0.01

Appendix

Table 4: K-sweep discriminability (

N{=}28

\tau{=}\tau(\text{probe},\text{Agent})

at each depth. Discriminability peaks at

K{=}3

–

7

and declines at higher depths as most models approach floor.

Figure 4: Model evaluation profiles across WMF-AM dimensions (

N{=}10

representative models). Left : Construct controls (K

{=}

1, non-arithmetic, CoT, supported agent, K

{=}

50). Right : Depth resilience (K

{=}

10 through K

{=}

50). Solid = API; dashed = open-weight. 10 models span the full capability range from Qwen2.5:3B to Claude-Sonnet-4.

Closed-source / API (7)
Model	Family	WMF-AM	Agent	Type
claude-sonnet-4	Anthropic	1.000	0.90	API
o3-mini	OpenAI	1.000	0.90	LRM
deepseek-r1-full	DeepSeek	1.000	0.50	LRM
deepseek-v3	DeepSeek	0.950	0.90	API
gpt-4o	OpenAI	0.800	0.90	API

Appendix

Table 5: Per-model scores (

N{=}28

, 12 families). WMF-AM = mean accuracy at

K\in\{3,5,7\}

(4 seeds, open-weight; single eval, API); Agent = 10-task deterministic agent battery completion rate. Sorted by WMF-AM within each group (API / open-weight).

\tau{=}0.595

(

p{<}0.001

Figure 5: Load-shift intervention ( $N{=}28$ ). Paired bars show supported (full history, solid color) vs. unsupported (last turn only, light color with colored edge) agent performance. $\Delta$ labels above each pair indicate the performance drop. Most models lose 20–60% (red $\Delta$ labels for drops $>0.3$ ); o3-mini and R1:14b are unaffected ( $\Delta{=}0.0$ , green labels). Models sorted by supported score; all 28 models shown.

#	Task ID	Name	Cat.	Description
1	multi_step_calc	Multi-step Calculation	ST	Decompose and compute $((17\times 3)+29)\times 2-15$ using a calculator tool
2	entity_tracking	Entity Tracking	ST	Track bank balances for 3 entities across 5 sequential transfers
3	sequential_search	Sequential File Search	ST	Search across multiple files to locate a target code number
4	uncertain_lookup	Uncertain Fact Lookup	NST	Verify uncertain information (boiling point of mercury) before answering
5	multi_source_conflict	Conflict Detection	NST	Detect conflicting facts across two source documents
6	conversation_recall	Conversation Recall	NST	Recall the first piece of information after 5 intervening tool calls

Appendix

Table 6: Agent battery: 10-task composition and classification. ST = state-tracking; NST = non-state-tracking.

Probe	$K$ -param	Ceiling	Construct iso.	Cross-domain	Agent valid.	Raw data	$N$ models
Minerva QS [ Xia et al., 2025 ]	No ( $K{=}200$ )	100% open-weight	No	No	No	No	17
Entity Tracking [ Kim and Schuster, 2023 ]	No (fixed)	High	No	No	No	No	6
n-back [ Zhang et al., 2024 ]	No (fixed)	Moderate	No	No	No	No	8
Huang et al. [ Huang et al., 2025 ]	No (fixed)	High	No	No	No	Partial	17
WMF-AM (ours)	Yes	14%	4 ablations	3 domains	10-task	Yes	28

Appendix

Table 7: WMF-AM vs. prior probes of state tracking / working memory in LLMs. Key:

K

-param = tunable depth parameter; Construct iso. = ablation controls; Cross-domain = non-arithmetic variant; Agent valid. = downstream agent battery; Raw data = per-model item-level data released.

Predictor	$\tau$	Calls/model
MMLU (published, 5-shot)	0.689	$\sim$ 14,000
WMF-AM $K{=}5$ only	0.644	20
WMF-AM $K{\in}\{3,5,7\}$	0.607	60
GSM8K (published, CoT)	0.589	$\sim$ 1,300

Appendix

Table 8: Model-matched predictor comparison (

N{=}17

). All predictors evaluated on exactly the same 17 models.

Working Memory in Language Models

See all 14 papers

Language Models Share a Readout Bottleneck in Working Memory
Oct 15, 2026Dongyu GongWorking Memory in Language Models Attention Mechanisms
Yale University
TriggerBench: Investigating Prospective Memory for Large Language Models
Jun 22, 2026Tianhua Zhang, Xinjiang Wang, Qianxi Zhang +6LLM Evaluation Long-Context Language Model Evaluation
The Chinese University of Hong Kong · Microsoft Research Asia
MemFail: Stress-Testing Failure Modes of LLM Memory Systems
May 26, 2026Ishir Garg, Neel Kolhe, Dawn Song +1LLM Agent Memory LLM Memory
University of California, Berkeley

Oct 15, 2026J/K move · Enter open · S save

Language Models Share a Readout Bottleneck in Working Memory

Dongyu Gong

Yale University

Working Memory in Language Models Attention Mechanisms LLM Interpretability

Open paper AskPDF arXiv

Large language models (LLMs) reproduce a human working-memory signature on the N-back task: performance declines as memory load N increases, and errors concentrate on lures, items that repeat the current stimulus at a neighboring but incorrect serial position. We ask whether this shared behavioral signature reflects a shared mechanism. The literature offers two candidate accounts. On a precision account, softmax attention cannot concentrate its weight on the item N positions back, so adjacent items are retrieved instead. On an interference account, attention selects the correct item, but the content of competing items corrupts the decision. We dissociate these accounts in 15 pretrained LLMs from three families (70 million to 14 billion parameters) and in minimal transformers trained from scratch. In all three families, the attention heads that carry the task attend to the target on most missed matches. Linear probes decode near-perfectly whether a matching letter occurred at the target or at a lure position, even on trials where the model false-alarms to the lure, although its own decisions barely distinguish the two. Sharpening attention yields negligible gains, whereas suppressing the letter identity of non-target items improves sensitivity (d') in nearly all model-by-load conditions. In larger models the task is spread over more heads, and removing interference helps more. In minimal transformers, precision limits a single attention layer, but a second layer removes this constraint. The working-memory limit of LLMs therefore lies in reading items out under interference rather than in locating them, consistent with interference accounts of human working memory.

Abstract

Working Memory in Language Models

Explore similar work

Language Models Share a Readout Bottleneck in Working Memory

TriggerBench: Investigating Prospective Memory for Large Language Models

MemFail: Stress-Testing Failure Modes of LLM Memory Systems