PaperlayerView exact revision on arXiv

Source-linked AI summary

Can Computation from Earlier Problems Help LLMs Solve New Ones?

Jipei He, Wenhui Tan, Xiaoyi Yu, Enver Sangineto, Fiorenzo Parascandolo, Rita Cucchiara, Ruihua Song

arXiv:2609.39394v1cs.AIcs.CLcs.LG

TL;DR

The paper asks whether computation from earlier problems can help solve a new independent problem. It introduces STAIR, which changes how current queries read stored earlier states, and reports higher later-turn accuracy across Qwen models and benchmarks.

  • Problem

    The paper asks whether computation left by earlier problems remains useful after a task switch, when each new independent problem supplies its own task-specific information.

    Source: §1 · p. 1

  • Method

    STAIR stores earlier response keys and values in a fixed bank and learns query reflections that change how the current prompt reads them while the model weights stay fixed.

    Source: §1 · p. 2

  • Results

    Across three Qwen models and four benchmarks, STAIR improves mean T2–T4 Avg@4 over Native by up to 11.67 percentage points while training 12,288 parameters.

    Source: §7 · p. 9

  • Takeaways & Limitations

    The results support treating computation left by completed problems as a resource whose readout can be learned for subsequent reasoning.

    Source: §7 · p. 9

    Current evidence covers four-turn sessions in the Qwen family, with reflector directions fitted on mathematical problems; broader training mixtures could test transfer.

    Source: §6 · p. 9

Abstract

from arXiv · show

Large language models often solve independent problems in the same conversation. Can computation from earlier problems help them solve new ones? To answer this question, we first conduct preliminary experiments showing that retained history can raise or lower later-turn accuracy, even within the same domain. To understand these effects, we use controlled replay to isolate internal state changes specific to each problem-history pairing. Across different histories, these changes preserve similar relationships among current problems. To improve reasoning under retained history, we introduce STAIR (Stale-Token Attention for Inter-query Reuse). STAIR captures keys and values from earlier response generation in a fixed bank. It learns to redirect current queries when they read this bank during prompt processing. The base model remains frozen; only 12,288 parameters are trained. Across three Qwen models and four benchmarks, STAIR improves average later-turn accuracy by up to 11.67 percentage points over the unmodified model with history.

1 INTRODUCTION

The introduction asks whether computation left by earlier problems can help with a new independent problem, since conversation history exposes both prior text and internal computations. It examines history effects and introduces STAIR, which learns to change how current queries read earlier states.

1 INTRODUCTION

  • On MATH-500, Native T4 Avg@4 falls from 95.00% under Vanilla to 93.75%, while GPQA-Diamond rises from 63.64% to 75.25% at Native T2, T3, and T4.The direction of the retained-history effect depends on the model and task.Source: §1 · p. 1
  • Replay shows persistent attention to earlier assistant responses, history-dependent shifts in current-problem representations, and similar nearest-neighbor relations across distinct histories beyond a shuffled-identity control.Paired Native and Vanilla answers show both corrected errors and lost correct answers.Source: §1 · p. 1
  • STAIR stores keys and values from completed assistant responses in a read-only bank and learns per-query-head reflections that change attention to those states while model weights stay fixed.It acts during prompt prefill; subsequent generation follows the original decoding path.Source: §1 · p. 2
  • The authors identify a recurring, problem-dependent component in representation changes conditioned on different histories.Source: §1 · p. 2
  • Across three Qwen models and four benchmarks, STAIR improves mean T2–T4 Avg@4 over Native by up to 11.67 percentage points while training 12,288 parameters.The backbone remains frozen.Source: §1 · p. 2

2 HISTORICAL COMPUTATION AFTER TASK SWITCHES

Controlled replay holds each current problem fixed while varying its source history. The results show persistent attention to earlier responses, recurring problem-dependent structure in history-induced changes, and answer changes in both directions.

2 HISTORICAL COMPUTATION AFTER TASK SWITCHES

  • Replay keeps the current problem fixed while varying completed three-turn source histories that do not contain that problem, using two groups with disjoint problems and history sources.The model reprocesses the saved conversation and current problem with its weights fixed.Source: §2 · p. 3, Figure 2 · p. 3

2.1 ATTENTION TO EARLIER RESPONSES AND STATE DISPLACEMENT

  • Earlier responses received 39.49% and 38.29% of attention mass at layer 3 and 12.41% and 12.87% at layer 19, while state displacement persisted into later layers.Attention mass is the share of current-problem prompt queries assigned to earlier assistant-response positions; displacement compares the representation under history with its no-history counterpart.Source: §2.1 · pp. 3–4, Figure 2 · p. 3

2.2 INTERACTION RESPONSES RETAIN PROBLEM-DEPENDENT STRUCTURE

  • Cross-history neighborhood overlap was 48.06% and 46.50%, versus 11.73% and 12.08% under shuffled identities, despite interaction responses accounting for only 0.69% and 0.66% of displacement energy.Figure 2C compares the eight nearest candidate problems across histories with a Question-identity shuffle control.Source: Figure 2 · p. 3, §2.2 · p. 4

2.3 NATIVE CONTINUOUS REASONING CHANGES ANSWERS IN BOTH DIRECTIONS

  • At T4, 35 answers changed from wrong to correct and 60 from correct to wrong, while Avg@4 fell 1.25 percentage points; Native also fell below Vanilla after correct T1 answers.The paired transitions are normalized by 2,000 pairs per turn.Source: Figure 2 · p. 3, §2.3 · p. 4

3 STAIR RE-ADDRESSES FROZEN HISTORICAL COMPUTATION

STAIR adds a read-only bank of earlier computation and learns to redirect current queries over it while leaving the backbone and stored states fixed. A differential read adjusts the current prompt computation, and the auxiliary branch is disabled during subsequent generation.

3 STAIR RE-ADDRESSES FROZEN HISTORICAL COMPUTATION

  • During each new prompt’s prefill, STAIR reads the separate fixed bank; after prefill, the auxiliary branch is disabled and decoding follows the original path.The visible conversation is processed again at each turn, and STAIR leaves stored entries unchanged.Source: §3 · p. 4, §3.3 · p. 6

3.1 A READ-ONLY HISTORICAL K/V BANK

  • Earlier response-generation states are stored as read-only keys and values, with bank contents drawn from reasoning bodies in Qwen3.5 and assistant-answer bodies in Instruct models.New states are detached when captured and appended for later turns.Source: §3.1 · p. 5

3.2 QUERY-SIDE RE-ADDRESSING

  • After auxiliary position encoding, learned reflectors reverse each controlled query head’s component along a normal while preserving query length.The controlled layers are fixed at [3, 11, 19] across the three model configurations.Source: §3.2 · pp. 5–6

3.3 DIFFERENTIAL READOUT AND LEARNING

  • STAIR subtracts the reference read from the reflected read, isolating the content change caused by redistributing attention over the same bank.The differential readout passes through the frozen output projection and is added to native attention; Qwen3.5 also applies its native gate.Source: Figure 3 · p. 5, §3.3 · p. 6
  • Training uses frozen-backbone responses arranged into four-turn sessions, with teacher forcing at T2–T4 and losses that update only the reflector normals.The objective combines response-token loss with KL divergence from the Native teacher, using λ = 1.Source: §3.3 · p. 6

4 EXPERIMENTS

Across three Qwen configurations and four reasoning benchmarks, the experiments compare STAIR with Native under retained histories and test shared-history controls, LoRA joint training, and bank-length costs. Results vary by benchmark and model, with the largest reported gains on AIME 2025.

4.1 EXPERIMENTAL SETUP

  • Evaluation covers three Qwen configurations and MATH-500, AIME 2025, AMC23, and GPQA-Diamond, with matched schedules and sample seeds for Native and STAIR.Controllers train for one epoch on responses from their own backbone; GPQA-Diamond tests transfer from mathematical training to scientific reasoning.Source: §4.1 · pp. 6–7

4.2 CONTINUOUS REASONING

  • On AIME 2025, mean T2–T4 Avg@4 rises over Native by 11.67 points for Qwen3.5-4B, 6.11 for Qwen3.5-9B, and 3.61 for Qwen3-4B Instruct.All three improve Avg@4 at each later turn and mean Pass@4 on AIME; results compare Native and STAIR averaged over T2–T4, with Vanilla as the no-history T1 reference.Source: Table 2 · p. 7, §4.2 · p. 7

4.3 HISTORICAL READOUT FROM A SHARED T1 HISTORY

  • With T1 held fixed, STAIR reaches 48.33% Avg@4 against 21.67% for Native, a 26.67-point gain; Pass@4 rises from 56.67% to 80.00%.Disrupting historical K–V pairings lowers STAIR to 33.06% across three seeds; the separately trained Bank-free controller reaches 36.67%.Source: §4.3 · p. 8, Table 3 · p. 8

4.4 LORA COMPARISON AND JOINT TRAINING

  • For Qwen3-4B Instruct on AIME 2025, STAIR improves Avg@4 over Native by 3.61 points, and adding STAIR to LoRA improves Avg@4 over LoRA by 3.89 points.LoRA trains 336 times as many parameters; Pass@4 rises by 4.44 points in both comparisons.Source: §4.4 · p. 8, Table 4 · p. 8

4.5 BANK PREFIX AND RUNTIME

  • On shared-T1 AIME 2025 inputs, T2 Avg@4 is 35.00%, 43.33%, and 39.17% with the first 2K, 8K, and 32K bank tokens, versus 48.33% with all saved tokens.In a separate fixed-length runtime test, full-bank reading adds a median 0.81 s and 6.60 GiB peak memory over Native.Source: §4.5 · p. 8

5 RELATED WORK

The paper situates STAIR among approaches that reuse computation across text segments, retrieve or cache prior states, and adapt reasoning in frozen or modified models. It distinguishes STAIR by changing how current queries read a fixed historical bank during prefill.

Related work on historical computation

  • Transformer-XL and Recurrent Memory Transformer carry information across segments through recurrence, while Memorizing Transformers retrieve past K–V pairs before attending over memory.Source: §5 · p. 9

Related work on prompt reuse

  • Prompt Cache and CacheBlend reuse computation for repeated prompt text, with CacheBlend selectively recomputing interactions with preceding text.Source: §5 · p. 9

Related work on reasoning adaptation

  • Unlike latent-reasoning and cache-augmentation approaches, STAIR subtracts a reference read from a query-reflected read over the same historical bank during prefill.The passage also surveys prompt tuning, pause tokens, latent reasoning, and sampling guidance [Tan et al., 2026a].Source: §5 · p. 9

Related work on task switches

  • Unlike the incremental task instructions studied by Laban et al. (2026), these sessions present a complete, independent problem at each turn.Source: §5 · p. 9

6 DISCUSSION AND LIMITATIONS

The discussion identifies uncertainty about how learned reading rules transfer across domains and bounds current evidence to four-turn Qwen sessions. It also notes growing bank costs in longer sessions and proposes broader training as a direction for future work.

Discussion and limitations

  • Current evidence covers four-turn sessions in the Qwen family, and broader training mixtures are proposed to test how reflector directions transfer beyond mathematical problems.The effect of reflector directions depends on the current query and historical keys.Source: §6 · p. 9
  • As completed turns add states, longer sessions increase bank storage and prefill reading costs; selective retention could limit bank growth.Source: §6 · p. 9
  • Pretraining on diverse task switches could learn reflector directions alongside the backbone; supervised fine-tuning or reinforcement learning could further test adaptation as reasoning behavior changes.Source: §6 · p. 9

7 CONCLUSION

The paper concludes that STAIR can improve later-turn reasoning by learning how current queries read earlier response states while keeping the backbone fixed. Its replay experiments use balanced problem–history pairings and measure attention and representation changes across transformer layers.

7 CONCLUSION

  • Up to 11.67 percentage points in mean T2–T4 Avg@4 over Native across four benchmarks and three Qwen models, with 12,288 parameters trained.The authors conclude that computation from completed problems can be treated as a resource whose readout can be learned for subsequent reasoning.Source: §7 · p. 9

A.1 REPLAY CONSTRUCTION

  • The replay design varies problem and history coverage, with 8,192 problem–history cells in the three-turn problem-coverage condition.The designs use disjoint current-problem sets and history-source pools across two data groups; the study also computes an empty-history reference for each of 256 current problems.Source: §A.1 · p. 15

A.2 ATTENTION AND STATE DISPLACEMENT

  • At layer 19, median state displacement is approximately 17.76 in both groups; at layer 35 it is 62.07 and 61.97.Replays use a frozen backbone, and the empty-history reference retains the system prompt and current-turn template.Source: §A.2 · p. 15, §A.1 · p. 15
  • Attention profiles track four separately aggregated source categories across layers, so their plotted values need not sum to exactly 100% at each layer.Figure 4 uses common scales; solid and dashed lines distinguish the two data groups.Source: §A.2 · p. 16, Figure 4 · p. 16

A.3 INTERACTION ENERGY AND CROSS-HISTORY GEOMETRY

The analysis separates problem–history interaction responses from additive problem and history effects, then tests whether their neighborhood structure recurs across histories. It also reports answer transitions, conditional accuracy groupings, and the association between historical reads and next-layer MLP responses.

A.3 INTERACTION ENERGY AND CROSS-HISTORY GEOMETRY

  • At k = 8, interaction-response overlap is 48.06% and 46.50%, versus raw displacement at 73.83% and 69.82%; additive reconstruction yields 100% in both groups.At layer 19 with three-turn histories, interaction energy is 0.69% in group 1 and 0.66% in group 2.Source: §A.3 · pp. 16–17

A.4 PAIRED ANSWER TRANSITIONS

  • Both wrong-to-correct and correct-to-wrong transitions occur at every later turn relative to Vanilla T1.Net changes are computed as wrong-to-correct counts minus correct-to-wrong counts, divided by 2,000 and multiplied by 100.Source: §A.4 · p. 17, Table 6 · p. 17

A.5 CONDITIONING ON THE FIRST HISTORY ANSWER

  • Table 1 pools T2–T4 accuracy by actual T1 correctness across three models on MATH-500, AIME 2025, and AMC23†.The correct-T1 and wrong-T1 pools contain 18,879 and 1,425 later responses, respectively; GPQA-Diamond is reported separately.Source: §A.5 · p. 18, Table 7 · p. 18, Table 8 · p. 19

A.6 HISTORICAL READS AND DOWNSTREAM MLP RESPONSES

  • At layer 3, adding earlier-assistant reads yields incremental test R2 of 0.0177 and 0.0170 in the two fit/test directions.Shuffling assistant reads across histories within each problem reduces the increment to 0.00304 and −0.00087; equal-width earlier-user reads carry comparable predictive information.Source: §A.6 · p. 21, Figure 6 · p. 21

B STAIR IMPLEMENTATION AND TRAINING

STAIR keeps captured historical keys and values in a separate bank and applies a trainable query-reflection branch during prompt processing. Training supervises later-turn responses while detaching bank entries and updating only reflector normals.

B.1 STATE CAPTURE AND BANK ACCUMULATION

  • Earlier assistant-generation keys and values are captured detached and appended in turn order, while prior bank entries retain their original captured values.The ordinary decoding cache is local to a turn; the auxiliary historical bank persists separately across turns.Source: §B.1 · pp. 21–22

B.2 AUXILIARY COORDINATES AND EXACT READS

  • The auxiliary branch assigns bank positions, applies the model’s RoPE configuration, reflects positioned queries, and centers logit changes using a reference-weighted key mean.The centering term is constant across bank positions for a fixed query and cancels under softmax, yielding the reweighting with a smoothed reference distribution.Source: §B.2 · p. 22

B.3 OUTPUT PATH AND PREFILL SCOPE

  • The auxiliary contribution is added to the native attention output before residual addition and the block’s MLP, and the branch is limited to current-prompt processing.It is disabled for earlier-history and target-response positions and turned off after prompt prefill during inference.Source: §B.3 · p. 22

B.4 SESSION REPLAY AND GRADIENT FLOW

  • Teacher-forced response losses train the current-turn prefix computation, while detached captured bank tensors remain fixed inputs for later turns.Training uses T1 for history and response supervision on T2–T4; the current prompt suffix is controlled by the auxiliary branch.Source: §B.4 · p. 23, Figure 7 · p. 23

B.5 SUPERVISION AND PARAMETER UPDATES

  • Only reflector normals are updated: 12,288 parameters per model across all query heads at zero-based layers [3, 11, 19].The backbone, vocabulary head, native output gate, and attention output projection remain frozen; the training objective also uses a detached Native distribution as its reference.Source: §B.5 · pp. 23–24

C.1 TRAINING SETTINGS

Training uses one-epoch STAIR runs with fixed model-specific decoding settings and a matched LoRA comparison. Evaluation matches problem assignments and sample seeds across conditions, then applies benchmark-specific answer scoring and completeness rules.

C.1 TRAINING SETTINGS

  • STAIR trains for one epoch with AdamW at learning rate 0.003, while all models update 12,288 reflector parameters at layers [3, 11, 19].The LoRA baseline has 4,128,768 trainable parameters, 336 times the STAIR count; compute budgets were not measured.Source: §C.1 · pp. 24–25

C.2 SAMPLING AND MATCHED CONDITIONS

  • Evaluation assigns each benchmark problem once to every turn column across sessions, with distinct problems within each session and assignments fixed across conditions.Native and STAIR use matched problems, references, turn assignments, and sample seeds, while each continues from its own generated history.Source: §C.2 · p. 25

C.3 ANSWER SCORING

  • Mathematical answers are scored from the last boxed response, while GPQA-Diamond answers are matched to the fixed ordering of each question’s choices.Turn-level scores are reported only when all four responses are present for every benchmark problem; incomplete combinations remain unreported.Source: §C.3 · pp. 25–26

C.4 SHARED-HISTORY CONTROLS AND RUNTIME MEASUREMENTS

The shared-history controls isolate how STAIR and bank-access choices affect later-turn evaluation, while runtime measurements quantify the costs of reading and constructing a historical bank.

C.4 SHARED-HISTORY CONTROLS AND RUNTIME MEASUREMENTS

  • All 120 shared-T1 inputs use fixed Native-path T1 trajectories, matched histories, current problems, seeds, and decoding budgets; only T2 answers are evaluated.The trajectories were generated separately from those used in continuous evaluation.Source: §C.4 · p. 26
  • Controls compare STAIR with ordinary-history conditions, a separately trained bank-free controller, random reflectors, direct reflected reads, and permuted K–V pairings.The K–V pairing and Random-reflector controls each cover three seeds; other conditions cover 120 responses.Source: §C.4 · p. 26
  • The bank-prefix study limits readable historical entries to 2,048, 8,192, or 32,768, or reads the full bank, while leaving visible history and query positions unchanged.Its 120 matched Qwen3.5-4B AIME 2025 inputs use four sample seeds per problem; the median bank contains about 63.4K tokens.Source: §C.4 · p. 26, Table 12 · p. 27
  • 567 ms, 0.81 s, and 6.60 GiB are the paired median increases over Native for T2 prefill, the fixed workload, and peak allocated memory, respectively.Measurements use a Full bank with a median 80.7K tokens; exact-token replay adds a median 7.29 s and is excluded from the table.Source: Table 13 · p. 27, §C.4 · p. 27

D.1 QWEN3-4B INSTRUCT

D.2 LORA COMPARISON AND JOINT TRAINING ON AIME 2025

The supplied passages identify an AIME 2025 adapter comparison for Qwen3-4B and turn-wise result tables for Qwen3.5-4B and Qwen3.5-9B.

D.2 LORA COMPARISON AND JOINT TRAINING ON AIME 2025

  • Table 15 covers Qwen3-4B Instruct on AIME 2025, comparing standalone and jointly trained adapters by turn and reporting T2–T4 means.Bold marks the highest score in each column, including ties.Source: Table 15 · p. 28

D.3 QWEN3.5-4B

  • Table 16 reports Qwen3.5-4B results by turn and gives means across T2–T4.Source: Table 16 · p. 29

D.4 QWEN3.5-9B

  • Table 17 reports Qwen3.5-9B results by turn and gives means across T2–T4.Source: Table 17 · p. 29

Related papers

Loading 2609.39394v1…