Source-linked AI summary
ST$^2$U: Stateful Test-Time Unlearning via Restricted Knowledge Boundary Control
Xunlei Chen, Qinghui Gong, Ruini Xue, Yaodong Hu, Tian Lan, Wenhong Tian
TL;DR
Test-time unlearning seeks to suppress designated knowledge without costly retraining or loss of non-target capabilities, but pointwise corrections can be undone by autoregressive hidden-state reconstruction. ST2U performs trajectory-wide restricted-boundary control with contextual anchoring and persistent correction state, achieving strong forgetting-retention trade-offs and much lower re-entry than test-time baselines.
Problem
Retraining is prohibitively expensive, while pointwise test-time corrections may lose persistence as autoregressive generation reconstructs hidden states from context and the generated prefix.
Method
ST2U models restricted-knowledge boundaries in low-dimensional invertible coordinates, preserves orthogonal components, and applies minimal anchored corrections with a persistent state across tokens.
Results
Across three benchmarks, ST2U achieves 80.6–87.9% forgetting, 90.7% average non-target capability retention, and 13.76–19.84% re-entry versus 46.50–59.10% for leading test-time baselines.
Takeaways & Limitations
ST2U provides a trajectory-wide runtime mechanism that reduces restricted-knowledge re-entry while preserving non-target structure and competitive forgetting.
Takeaways & Limitations
The formulation assumes a task-specific forget set F for targeted unlearning.
Abstract
from arXiv · showhide
Controlling restricted knowledge in large language models is essential for model alignment and safe deployment. Test-time unlearning avoids costly retraining and parameter updates by intervening only during inference. However, existing activation-editing methods apply isolated pointwise corrections, overlooking how autoregressive generation continually reconstructs hidden states from the prompt, cache, and generated prefix. Consequently, later states may return to restricted knowledge regions after a locally successful correction, causing restricted knowledge re-entry. In this work, we propose Stateful Test-Time Unlearning via restricted knowledge boundary control (ST$^2$U), which formulates test-time unlearning as trajectory-wide boundary control. ST$^2$U first models restricted knowledge boundaries in low-dimensional invertible coordinates while leaving orthogonal non-target components unchanged. During inference, ST$^2$U monitors risk along the trajectory, applies minimal boundary corrections with contextual anchoring, and propagates historical correction states across tokens to mitigate knowledge re-entry. This trajectory-wide control enables more persistent forgetting while preserving non-target capabilities and limiting inference overhead. Across three benchmarks and three model families, ST$^2$U delivers the strongest overall balance, combining best or second-best retention with competitive forgetting and substantially less restricted-knowledge re-entry than test-time baselines (13.76%-19.84% versus 46.50%-59.10%).
Introduction
Test-time unlearning avoids parameter updates but existing pointwise interventions can lose effectiveness during autoregressive decoding. ST2U replaces isolated edits with stateful trajectory control, achieving strong forgetting-retention trade-offs and substantially lower restricted-knowledge re-entry.
- Motivation: Retraining and repeated parameter updates are costly, motivating inference-time unlearning that preserves non-target utility.Test-time methods intervene through prompts, token distributions, or activations while freezing model parameters.
- Problem: Autoregressive hidden states can return to restricted regions after a locally successful edit, producing harmful details later in the response.This failure occurs because later states depend on the prompt, cache, and generated prefix.
- Method: ST2U models restricted-knowledge boundaries in low-dimensional invertible coordinates while preserving orthogonal non-target activation components.Its controller monitors risk, applies minimal context-anchored corrections, and propagates historical correction states across tokens.
- Results: 80.6–87.9% forgetting rates and 90.7% average non-target capability retention establish ST2U’s leading overall trade-off across three benchmarks.The reported result combines forgetting with preservation of retained capabilities rather than optimizing either axis alone.
Related Work
Prior unlearning methods modify parameters or control individual inputs, token distributions, or activations. ST2U addresses their missing temporal coupling by controlling restricted-knowledge risk over the evolving hidden-state trajectory.
- Parameter-based Unlearning: Parameter-based methods encode forgetting in updated weights, improving stability or granularity but lacking online context-dependent control.Examples include GA, RMU, ASU, and ALTER.
- Test-Time Unlearning: Test-time methods operate through input, agentic, decoding, or activation interfaces to suppress restricted knowledge during inference.These approaches include soft prompts, specialized agents, entropy-guided decoding, and hidden-representation editing.
- Gap: Existing test-time methods control a query, token distribution, or current activation without explicitly modeling risk over the evolving hidden-state trajectory.Their interventions remain independent across decoding steps and carry no correction history.
- ST2U: ST2U treats test-time unlearning as stateful trajectory control rather than a sequence of isolated corrections.This formulation supplies the temporal coupling missing from prior test-time methods.
Observation and Problem Formulation
The paper defines restricted-knowledge re-entry as renewed boundary violation after an earlier correction, then formulates unlearning as maintaining every generated state outside the risk region. The formulation couples corrections across the trajectory while accounting for censored re-entry observations.
- Restricted Knowledge Re-entry: The frozen LLM generates future tokens autoregressively from a prompt, with pre-control and controlled hidden states defined at each position.The formulation considers a task-specific forget set for targeted unlearning.
- Restricted Knowledge Re-entry: A task-specific forget set F specifies the knowledge targeted for unlearning.The model parameters remain frozen while the intervention operates during inference.
- Restricted Knowledge Re-entry: A stateless editor corrects only the current violating state, so later states reconstructed from context and the generated prefix may violate the boundary again.This recurrence is the paper’s restricted-knowledge re-entry failure mode.
- Restricted Knowledge Re-entry: Re-entry means a renewed violation before the controller acts at a later position, not failure to correct that later state.RF supplies continuous restricted-knowledge risk and τ supplies its threshold.
- Restricted Knowledge Re-entry: Stateful unlearning restores each violating state while using past corrections to reduce later violations.The objective is persistence across decoding rather than only local feasibility.
- Stateful Test-Time Unlearning: The optimization seeks minimum corrections that keep every controlled prefill and decoding state outside the risk region with frozen model parameters.The constraint is evaluated along the generation history and earlier corrections.
- Stateful Test-Time Unlearning: Figure 2 illustrates local closure failure and quantifies re-entry and verification rates, treating generations without observed re-entry as right-censored.The figure evaluates whether apparent local unlearning persists throughout generation.
- Stateful Test-Time Unlearning: The post-control state adds the current correction to a history-dependent pre-control state, coupling per-position constraints along the generated trajectory.Batched prefill approximates causal dependence by scanning prompt positions to propagate controller state without recomputing later activations.
ST2U: Stateful Boundary Control for Test-Time Knowledge Unlearning
ST2U controls test-time unlearning as a hidden-state trajectory problem: it models restricted-knowledge risk in a low-dimensional geometry, corrects risky states, and carries correction history across tokens.
- Restricted Knowledge Geometry Modeling: ST2U separates target-related activation variation from orthogonal residual components using paired trajectories, an invertible coordinate map, and a low-dimensional basis.The residual is restored unchanged, while the learned coordinates are edited.
- Restricted Knowledge Geometry Modeling: KDE models restricted and sanitized trajectories in coordinate space, with their log-density ratio defining risk and a threshold separating risky states.Context-dependent references are associated with trajectory neighborhoods for local anchoring.
- Stateful Boundary Control: Risk-triggered control leaves states unchanged below the boundary and applies corrections during both prefill and decoding when risk exceeds the threshold.The same risk test is used throughout the generation process.
- Stateful Boundary Control: The correction direction combines normal descent with historical guidance projected onto the boundary tangent space, so history cannot reverse local risk descent.The normal term reduces current risk, while the tangent component preserves consistent historical direction.
- Stateful Boundary Control: The controller minimizes intervention magnitude through a local surrogate, iteratively recomputing risk and gradients until reaching the boundary or an iteration limit.This provides a tractable per-token alternative to online parameter optimization.
- Stateful Boundary Control: Contextual anchoring keeps corrected states supported by sanitized references, while a fixed-dimensional recurrent state carries accepted correction directions across tokens.The history decays at low-risk positions and accumulates during consecutive high-risk corrections to discourage re-entry.
Experiments
Experiments evaluate ST2U across entity, hazardous-knowledge, and copyright unlearning benchmarks, with retained utility, re-entry, sensitivity, and component ablations. ST2U maintains strong forgetting-retention trade-offs while reducing restricted-knowledge re-entry across settings.
- Evaluation setup: ST2U is evaluated on RWKU, WMDP, and MUSE-Books, with MMLU measuring retained general knowledge and R-Ent. measuring confirmed restricted-knowledge re-entry.RWKU uses fill-in-the-blank and question-answering probes; WMDP covers biosecurity and cybersecurity; MUSE-Books evaluates Harry Potter copyright unlearning.
- Forget-retain trade-off: ST2U delivers the strongest or second-strongest retention among test-time methods while remaining competitive on forgetting across entity, hazardous-knowledge, and copyright settings.On MUSE-Books, it nearly matches AS in BLEU while retaining 4.78 more MMLU points and higher fluency.
- Restricted knowledge re-entry: 19.20/13.76, 19.84/16.20, and 18.10/15.40 are ST2U’s R-Ent. pairs across Llama, Qwen, and SpikingBrain, consistently the lowest among compared methods.Prompt-level and stateless methods leave later states weakly constrained, whereas ST2U uses target-specific state memory to stabilize trajectories.
- Hyperparameter sensitivity: Sweeping the risk-boundary quantile from 88% to 98% produces small, non-monotonic changes, supporting one shared operating boundary rather than per-model oracle thresholds.The maximum boundary step αmax controls local closure strength; αmax = 12 avoids the retention degradation observed at larger caps.
- Ablation of Control Components: Removing historical state increases R-Ent. by 23.74, removing the risk gate increases Avg. WMDP and R-Ent. by 13.09 and 32.99, and removing the residual reduces MMLU by 7.79.The ablations associate historical state with temporal persistence, risk gating with selective control, and the orthogonal residual with retained utility.
Discussions
Discussion experiments examine white-box robustness, layer-wise coordinate benefits, and deployment cost. ST2U remains stable under tested attacks, broadens effective intervention layers, and offers a favorable joint cost profile.
- Layer-wise Coordinate Robustness: Across layers 16–31, 13 of 16 layers improve, 10 gain at least 5 WMDP points, and 6 exceed 8 points over direct hidden steering.Layer 23 adds a 12.27-point WMDP reduction while improving MMLU by 4.72 points.
- Deployment cost: ST2U achieves Avg. WMDP accuracy 29.26 and retained MMLU accuracy 60.58 with 2.8 minutes of offline preparation and 2.21× online latency.It strictly outperforms SCANS across offline time, online latency, WMDP, and MMLU in the deployment comparison.
Conclusion
The conclusion frames ST2U as trajectory-wide, boundary-aware control for restricted knowledge in frozen LLMs. It reports strong forgetting-retention trade-offs, lower re-entry, and stable behavior across probing and long-form continuation tasks.
- Conclusion: ST2U models restricted-knowledge boundaries in low-dimensional invertible coordinates, preserves orthogonal non-target structure, and maintains a shared correction state during generation.The conclusion positions these mechanisms as a unified runtime control approach across heterogeneous forgetting scenarios.
- Conclusion: The reported framework reduces re-entry, layer brittleness, and utility degradation relative to existing methods across benchmark probing and long-form continuation tasks.The paper identifies sparser trigger policies and adaptive intervention schedules as future directions for reducing online overhead.