Spontaneous Context Restoration: How Language Models Recover from Corrupted Inputs

Introduction

• Across architectures and tasks, repair effects shift from corrupted positions in early layers toward clean context and prediction positions later, while hidden states predict repair outcomes.The first-block probe reaches mean AUC 0.78; moderate-corruption finetuning raises the controlled model’s 50%-accuracy corruption tolerance from 40% to 90%.Source: Introduction · pp. 1–2

Related Work

• Prior corruption studies report performance degradation and behavioral mitigation, but do not characterize internal processes associated with successful performance under corruption.Related work also covers clean-input circuits, compensation after internal ablation, representation geometry, and corruption finetuning.Source: Related Work · p. 2
• Earlier corruption and prompt-perturbation studies emphasize behavioral degradation, while circuit, compensation, and representation studies address related but distinct internal questions.This work focuses on recovery from external, label-preserving corruption rather than compensation after internal ablation.Source: Related Work · p. 2

Experimental Setup

• The study asks when models preserve correct task-relevant computation under label-preserving corruption and which internal representations distinguish restoration from failure.In ARC and CommonsenseQA, question text is corrupted; factual QA corruption includes context and question, while formatting instructions and answer choices remain protected.Source: Experimental Setup · p. 3

Synthetic Symbolic Setting: Controlled Attention-only transformers on arithmetic

• Two attention-only decoder transformers, trained exclusively on clean arithmetic sequences, test restoration in a controlled setting.The 4-layer model checks depth-independence and replicates main findings from the 10-layer model.Source: Synthetic Symbolic Setting: Controlled… · p. 3

Natural-Language Setting: Pretrained language models on NLP tasks

• Five instruction-tuned LLMs are evaluated on ARC-Challenge, factual QA, and CommonsenseQA under four corruption modes.Prompts protect task scaffolds and answer choices; accuracy degrades with corruption, with dropout and replacement most destructive and factual QA most fragile.Source: Natural-Language Setting: Pretrained… · p. 3

Activation patching

• Activation patching replaces corrupted-run activations with clean counterparts at corrupted or clean positions and measures the change in the target-token logit.Pretrained-model patching is restricted to replacement and spelling corruptions because they preserve sequence length.Source: Analysis Methods · p. 3

Cosine similarity and repair probes

• Cosine similarity measures alignment between corrupted and clean activations at the prediction position, while cosine and logistic-regression probes test whether repair outcomes are predictable from hidden states.The logistic regression probe uses the full corrupted activation with five-fold stratified cross-validation.Source: Analysis Methods · pp. 3–4

Linearization test

• The linearization test compares the true output-logit change along a corruption direction with a first-order Taylor prediction; the residual measures departure from linearity.Small residuals indicate a locally linear region of the logit map; at α = 1.0, the residual captures nonlinearity beyond second order.Source: Analysis Methods · p. 4

Repair Signals Shift from Corrupted to Clean Positions

• Repair effects shift from corrupted positions early to clean positions later, accumulating through the residual stream and concentrating at the prediction position.At clean positions, attention adds roughly uniform increments while skip-connection effects grow across layers.Source: Where does repair happen? Repair… · p. 4
• Ablating 4, 6, or 8 of 8 heads at one target layer leaves accuracy close to baseline, so repair does not rely on an individual layer-local head subset.Source: Where does repair happen? Repair… · p. 4

Failures Encounter High Curvature Along Corruption Trajectories

• 6−8× larger final-layer linearization residuals occur for failed than repaired examples at α = 1.0 across pretrained models.At α = 0.25, first-order predictions closely match true logit changes; the failed-to-repaired residual ratio rises from ≈2× to ≈4× as α increases from 0.25 to 1.0.Source: Why does repair fail? Failures… · pp. 4–5, Figure 3 · p. 5

Repair Outcome Is Decodable from Residual-State Representations

• Cosine alignment predicts repair strongly in controlled models, but is insufficient in some pretrained settings, where linear probes capture information beyond global angular alignment.The pretrained-model separation is not simply corruption detection, which is near-ceiling from the earliest layers.Source: What distinguishes repair from… · pp. 5–6, Figure 4 · p. 5

Early Hidden States Support Deployable Failure Triage

• Layer-0 probes achieve mean ROC-AUC 0.749–0.797 across five pretrained models and flag 16.4–33.4% of failures within the top-risk 10%.The probe uses the corrupted prompt’s early hidden state without a clean reference or completed generation.Source: Figure 5 · p. 6, Can failure be predicted before… · p. 6

Moderate-Corruption Finetuning Improves Robustness

• FT@50 raises controlled-model corruption tolerance at 50% accuracy from 40% to 90%, while clean accuracy remains 0.997 vs. 0.999.Pretrained-model gains vary by task and corruption; no single finetuning severity is uniformly optimal.Source: Can repair robustness be improved?… · pp. 6–7, Figure 6 · p. 7

Conclusions

• Context restoration emerges spontaneously in attention-only models trained exclusively on clean sequences, and repair outcomes relate to residual-state alignment and downstream nonlinearity.The controlled setting removes MLP sublayers; pretrained models show linear probes outperforming cosine alignment, though the roles of pretraining and MLPs are not isolated.Source: Conclusions · p. 7
• Across attention-only transformers and five pretrained LLMs, early layers localize repair effects at corrupted positions before late layers concentrate them at the prediction position.Head ablation leaves repair accuracy close to baseline in the controlled model.Source: Conclusions · p. 7
• Failed examples show linearization residuals up to 6–8× larger than repaired examples; moderate-corruption finetuning raises τ50 from 40% to 90% while preserving clean performance.The robustness gain accompanies reduced displacement-normalized linearization error along corruption directions.Source: Conclusions · p. 7
• A corrupted prompt’s early hidden state can flag elevated failure risk before generation for clarification, rerunning, or more expensive verification.The authors also report that LLM finetuning gains depend on task and corruption type.Source: Conclusions · p. 7
• Together, the findings connect restoration to depth-wise routing and residual-representation geometry, and show that corruption-aware finetuning can modify robustness and local response.These are representation-level diagnostics for repair failure within the studied settings.Source: Conclusions · p. 7

Limitations

• The attention-only arithmetic setting cannot establish direct transfer of its quantitative findings to full-architecture transformers, multilingual settings, or tasks whose targets are not preserved under local corruption.The analysis removes MLP sublayers and uses procedurally generated arithmetic sequences; five instruction-tuned decoder-only LLMs are evaluated on three English-language task families and four label-preserving corruption modes.Source: Limitations · p. 7
• Generality beyond the evaluated architectures remains untested, including mixture-of-experts, state-space, encoder-decoder, and retrieval-augmented models.The authors also note that the attention-only analysis cannot fully characterize how MLP-shaped features reshape repair geometry.Source: Limitations · p. 7

S1.1 Problem Setup

• The setup asks whether a model predicts the correct target from a label-preserving corrupted input and traces where recovery occurs in the network.Predictions are read from the final sequence position.Source: S1.1 Problem Setup · p. 10

S1.2 Data Generation

• Procedurally generated arithmetic sequences provide unambiguous next-token targets and known corruption positions for comparing clean and corrupted internal trajectories.Four sequence families are used; their starting values and step sizes are sampled from family-specific distributions.Source: S1.2 Data Generation · p. 10
• The training corpus contains 39,360 sequences, and the tokenizer makes each integer a single token while randomly permuting integer token indices.The vocabulary contains 13,008 tokens; test sequences vary in length while training sequences use a fixed length.Source: S1.2 Data Generation · p. 10, S1.3 Tokeniser · p. 10

S1.4 Model Architecture

• Two decoder-only attention-only transformers are trained on clean sequences with next-token cross-entropy and no feedforward sublayers.The 10-layer model is used for main mechanistic analyses, while a 4-layer model checks that findings are not depth-specific.Source: S1.4 Model Architecture · p. 10, S1.5 Training · p. 11

S1.6 Corruption Regimes (Test Time)

• Test corruptions replace uniformly selected positions with zero, in-range values, or out-of-range values while preserving the prediction target.Evaluation measures exact-match next-token accuracy on held-out sequences across corruption regimes, rates, lengths, and sequence families.Source: S1.6 Corruption Regimes (Test Time) · p. 11, S1.7 Evaluation Protocol · p. 12

S1.8 Fine-tuning on Corrupted Data

• Finetuning uses corrupted sequences as inputs and their original clean next tokens as labels, then evaluates the models across the same held-out conditions as the base model.The procedure starts from converged checkpoints and reports accuracy-versus-ablation curves, τ50, mean accuracy, and clean accuracy.Source: S1.8 Fine-tuning on Corrupted Data · p. 12

S1.11 Activation Patching

• Activation patching compares clean and corrupted runs by replacing corrupted or clean position groups with clean-run activations and measuring target-logit changes.Individual-position patching maps effects across layers and sequence positions; all-position interpolation varies the mix of clean and corrupted activations.Source: S1.11 Activation Patching · pp. 12–13

S1.12 Cosine Similarity and Repair Probes

• Cosine similarity and logistic-regression probes test whether prediction-position activations distinguish repaired from failed inputs, including at matched corruption rates.The linearization test measures curvature of the target-logit function along corruption-induced displacement while holding other positions and downstream computation fixed.Source: S1.12 Cosine Similarity and Repair… · p. 13, S1.13 Linearisation Scale Test · p. 13

2. Corrupted forward pass on input ˜x. A hook captures hcorr

• The test measures linearization residuals between true and first-order predicted logit changes, with a directional curvature proxy and RMSE used to quantify departures from linearity.The score at α = 1 is −|r(1)|, where higher values indicate a more linear response and greater repairability.Source: §2 · p. 14
• Residual comparisons condition on examples whose clean prediction is correct, then distinguish repaired from failed examples by whether the corrupted prediction remains correct.Plots compare curvature across those corrupted-run outcome labels.Source: §2 · p. 14
• At each layer, example, corruption regime, and ablation rate, the procedure constructs perturbed states and obtains true logit changes with a patched forward pass.The hook modifies only position p at layer ℓ before the forward pass continues through later layers.Source: §2 · pp. 13–14

S1.14 Head Ablation

• Head ablation replaces selected heads’ post-softmax attention weights with uniform causal attention, preserving the causal mask while removing content-dependent attention.Non-ablated heads remain unchanged; subsets of 4, 6, or 8 of 8 heads are tested across layers.Source: S1.14 Head Ablation · p. 14
• For each head-ablation configuration, the study records layer-wise residual token agreement and final accuracy on corrupted inputs at 10%, 30%, and 50% ablation rates.Three random subsets are sampled for each subset size.Source: S1.14 Head Ablation · p. 14

S2.1 Problem Setup

• The evaluation asks whether pretrained instruction-tuned models retain the target answer under label-preserving prompt corruption and traces recovery internally using greedy decoding without a system prompt.The evaluated model set is listed in Table 4.Source: S2.1 Problem Setup · p. 14, Table 4 · p. 14, S2.2 Models · p. 14

S2.2 Models

• Five instruction-tuned pretrained language models are evaluated, with prompts rendered through each model’s native chat template and no system prompt.Models up to 8B use bfloat16; larger models use 8-bit quantization.Source: S2.2 Models · p. 14, Table 4 · p. 14

S2.3 Tasks and Prompt Construction

• The task suite includes ARC-Challenge science multiple choice, CommonsenseQA multiple choice, and SQuAD-derived factual QA requiring answers copied from context.Prompts separate a corruption region from protected scaffold tokens and surrounding user text.Source: S2.3 Tasks and Prompt Construction · p. 15

S2.4 Corruption Modes

• Four corruption modes—dropout, replacement, spelling, and distractor insertion—are applied at rates from 10% to 100%, with rate defined against token mass in the corruption region.Dropout deletes words, replacement substitutes function words, spelling shuffles internal characters, and distractors insert irrelevant sentences.Source: S2.4 Corruption Modes · p. 15

S2.5 Evaluation Metrics

• Outputs are evaluated with exact match, token-level F1, token agreement with the clean prediction, and a repair label defined as F1 ≥ 0.5.For multiple-choice tasks, exact match also accepts the correct option letter.Source: S2.5 Evaluation Metrics · p. 15

S2.6 Position-Group Activation Patching

• Position-group patching compares target-logit changes after replacing corrupted-position activations or clean-position activations with their clean-run counterparts.Because dropout and distractor corruption change token counts, this analysis is limited to replacement and spelling; examples with differing tokenization lengths are excluded.Source: S2.6 Position-Group Activation Patching · p. 16

S2.7 Cosine Similarity and Repair Probes

• Cosine similarity and logistic-regression probes use final-token hidden states to assess repair predictability, while a separate probe tests whether corruption identity is linearly decodable.Repair probes use the F1 ≥ 0.5 label and report five-fold stratified cross-validated ROC-AUC; evaluations also control within corruption-mode and ablation-rate combinations.Source: S2.7 Cosine Similarity and Repair Probes · p. 16

S2.8 Linearization Scale Test

• The linearization test perturbs the final-position hidden state at four scales and compares its predicted logit change with patched-forward results.Repair labels come from precomputed evaluation metrics, and the test runs at every tenth layer plus the second-to-last layer.Source: S2.8 Linearization Scale Test · p. 16

S3 Supplementary Boxes

• Figure 8 shows ARC multiple-choice prompts across corruption families, with clean prompts in green, corrupted prompts in orange at p = 0.5, and the gold answer in red.Source: Figure 8 · p. 17
• Figure 9 presents CommonsenseQA prompts across corruption families, distinguishing clean prompts, corrupted prompts at p = 0.5, and the gold answer by color.Source: Figure 9 · p. 18
• Figure 10 shows factual prompts with clean and p = 0.5 corrupted versions, highlighting the target answer in red.Source: Figure 10 · p. 19
• Figure 11 presents full symbolic toy-task sequences with clean inputs, in-range, out-of-range, and zero-ablation corruptions, while keeping the gold next-number target unchanged.Source: Figure 11 · p. 20

S4 Supplementary tables

• Table 5 reports controlled-model head-ablation accuracy under zero ablation, comparing a same-corruption-rate baseline with post-ablation accuracy averaged across eligible layers.Standard deviations are shown in parentheses.Source: Table 5 · p. 21
• Table 6 gives controlled-model is-corrupted probe AUC averaged over two recorded hook sites.The table comment records the maximum absolute difference between hook sites.Source: Table 6 · p. 21
• Table 7 reports pretrained is-corrupted probe AUC on ARC, with probed layers as rows and models as columns, averaging across four corruption modes.Cells include standard deviations.Source: Table 7 · p. 22
• Table 8 summarizes linearization residual diagnostics by model and task family, including the slope from residual norm to linearization RMSE.Source: Table 8 · p. 22

• Pretrained-model curves show dropout and replacement are generally more damaging than distractor-style corruption, alongside toy-model ablation curves and clean exact-match baselines.Source: Figure 12 · p. 23, Figure 13 · p. 23, Figure 14 · p. 24
• Patching analyses show repair effects concentrated in a small subset of positions, remaining near corrupted locations early and spreading toward output-relevant positions later; another comparison excludes the final position.Source: Figure 15 · p. 25, Figure 16 · p. 26, Figure 17 · p. 26
• Pretrained-model position-group patching plots the probability of restoring the clean prediction by patching corrupted versus clean positions across corruption modes and task families.Source: Figure 18 · p. 27
• Toy-model repaired examples remain more aligned with clean activations than failed examples, while repair-success probes compare linear and cosine readouts and corruption identity becomes linearly decodable only after several layers.Pretrained ARC probes and earliest-layer probe plots provide additional probe comparisons; the latter reports ROC-AUC and 10%-budget recall lift.Source: Figure 19 · p. 28, Figure 20 · p. 28, Figure 21 · p. 29
• Toy-model plots show larger linearization residuals for failed examples at larger interpolation steps and corruption levels, while pretrained-model plots show the clearest separation at larger steps and later layers.Source: Figure 23 · p. 31, Figure 24 · p. 32, Figure 25 · p. 33
• Finetuning on corrupted inputs differentially improves LLM robustness and reduces displacement-normalized linearization error, with the largest error reduction at moderate corruption.Source: Figure 26 · p. 34, Figure 28 · p. 35