ST$^2$U: Stateful Test-Time Unlearning via Restricted Knowledge Boundary Control

Xunlei Chen; Qinghui Gong; Ruini Xue; Yaodong Hu; Tian Lan; Wenhong Tian

cs.LGAug 24, 2026

ST $^2$ U: Stateful Test-Time Unlearning via Restricted Knowledge Boundary Control

Authors: Xunlei Chen, Qinghui Gong, Ruini Xue, Yaodong Hu, Tian Lan, Wenhong Tian

Organizations: University of Electronic Science and Technology of China · Southwest Jiaotong University

Abstract

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%).

Figures & tables

Figure 1: Example of restricted knowledge re-entry during decoding. A locally effective intervention initially yields a safe high-level explanation, but later decoding returns to the restricted topic and produces harmful details.

Figure 2: Empirical characterization of restricted knowledge re-entry. Panel (a) illustrates that local closure after knowledge unlearning does not persist throughout generation, while panel (b) quantifies the re-entry and verification rates across intervention strategies. We treat generations without an observed re-entry event as right-censored.

Figure 3: Overview of ST 2 U. Offline paired trajectories define low-dimensional invertible coordinates, KDE-based risk geometry, and context-dependent references. During prefill and decoding, risk triggered control combines normal descent with tangent-projected history, steps toward the risk boundary, and anchors feasible states to local references. A fixed-dimensional state propagates accepted correction directions across tokens to mitigate restricted knowledge re-entry.

Baseline	RWKU					WMDP
Baseline	FB $\downarrow$	QA $\downarrow$	Gen. $\uparrow$	Flu. $\uparrow$	R-Ent. $\downarrow$	Bio. $\downarrow$	Cyber $\downarrow$	Gen. $\uparrow$	R-Ent. $\downarrow$
Llama3.1-8B-Instruct ( 7 )
Base	88.31	74.92	64.38	4.20	89.14 E	72.74	47.35	64.38	91.14 E
RMU $\dagger$	33.10	27.84	56.71	3.19	21.30	35.68	38.82	51.31	19.80
AS $\dagger$	15.90	19.62	57.40	3.43	59.80	27.72	33.82	57.74	57.40
ICUL $\ddagger$	50.20	43.77	60.80	4.07	83.50	55.90	44.30	59.60	82.80

Table 1: Main results on RWKU and WMDP.

\dagger

\ddagger

, and

\star

denote parameter-updating, in-context, and test-time intervention methods. R-Ent. denotes confirmed re-entry. Base E has no preceding closure, we report Entry Rate instead of confirmed R-Ent.

Method	Forget Perf.		Retain Perf.
Method	BLEU $\downarrow$	R-L $\downarrow$	MMLU $\uparrow$	Flu. $\uparrow$
Original	74.76	85.14	46.39	4.03
Retain	1.91	9.83	47.82	2.87
WHP	23.55	17.93	40.40	2.52
AS	13.72	11.68	37.56	3.08
ICUL	38.47	32.49	39.86	3.60

Table 2: Copyright unlearning results on MUSE-Books.

Figure 4: Parameter sensitivity analysis on WMDP.

Variant	WMDP $\downarrow$	MMLU $\uparrow$	R-Ent. $\downarrow$
Full ST 2 U	29.26	60.58	13.76
w/o historical state	37.23	62.10	37.50
w/o risk gate	42.35	55.55	46.75
w/o residual preservation	35.45	52.79	18.75

Table 3: Ablation on Avg.WMDP (Llama-3.1-8B-Instruct).

Figure 5: Layer-wise gains from low-dimensional invertible coordinates over same-layer direct hidden steering on Llama-3.1-8B. Positive values denote additional Avg. WMDP accuracy reduction. Colors indicate MMLU accuracy changes, and bubble sizes encode the standard deviation across runs.

Method	Recovery $\downarrow$	$\Delta$ MMLU
RMU	4.87	-17.32
SCANS	15.84	+1.28
ST 2 U	1.16	-1.92

Table 4: White-box test-time attack (Llama3.1-8B-Instruct).

Figure 6: Relative deployment cost (Llama3.1-8B-Instruct).

LLM Unlearning

Large language model unlearning studies how to remove specified knowledge or behaviors while preserving a model’s unrelated capabilities.

See all 191 papers

Test-Time Unlearning via Sparse Autoencoder
Sep 14, 2026Pingzhi Li, Jinhao Duan, Vaishnav Tadiparthi +6Machine Unlearning LLM Unlearning
University of North Carolina at Chapel Hill · 2Honda Research Institute · 3Michigan State University
Attention Smoothing Is All You Need For Unlearning
Mar 1, 2026Saleh Zare Zade, Xiangyu Zhou, Sijia Liu +1Machine Unlearning LLM Unlearning
Department of Computer Science, Wayne State University · Department of Computer Science and Engineering, Michigan State University · Institute for AI and Data Science, Wayne State University
Not Every Change Is Necessary: Recoverable Drift in Large Language Model Unlearning
Oct 8, 2026Xunlei Chen, Qinghui Gong, Jingkun Xue +3Machine Unlearning LLM Unlearning
University of Electronic Science and Technology of China · Southwest Jiaotong University

Sep 14, 2026J/K move · Enter open · S save

Test-Time Unlearning via Sparse Autoencoder

Pingzhi Li, Jinhao Duan, Vaishnav Tadiparthi, Nakul Agarwal+5

University of North Carolina at Chapel Hill · 2Honda Research Institute · 3Michigan State University

Machine Unlearning LLM Unlearning

Open paper AskPDF arXiv

Machine unlearning aims to remove specific knowledge from a trained large language model (LLM) without retraining from scratch. Existing methods modify model weights via gradient ascent and its advances. While effective on certain benchmarks, these weight-based approaches exhibit a sharp forget-utility trade-off, where stronger forgetting of target knowledge can degrade model utility, and unlearned knowledge may reappear under post-unlearning fine-tuning or prompt attacks. We propose ARIA (autoencoder-gated inference-time unlearning), a test-time unlearning method that leaves model weights intact and gates access to unwanted knowledge only when generation enters a forget-related state. ARIA uses sparse autoencoder (SAE) latents to train a lightweight linear detector, then applies an interpretable intervention on triggered states with negligible test-time overhead. Empirical evaluations on TOFU, R-TOFU, and WMDP show that ARIA improves the forget-retain trade-off over weight-based baselines across both a thinking model (DeepSeek-R1-Distilled-Qwen-1.5B) and an instruction model (Gemma-3-1B-it), e.g., reducing WMDP-cyber forget-set accuracy significantly while keeping MMLU within 1% of the pre-unlearning model. We further introduce three post-unlearning adversarial attacks targeting weight-space and decoding-space recovery, and find that ARIA remains robust under all three, with forgetting changing by less than 1% under attack. A feature-level case study leveraging the interpretability of ARIA suggests that some retain degradation may reflect response styles underlying the unlearning data rather than leakage of the targeted knowledge itself, highlighting a potential source of bias in unlearning task construction.

Abstract

LLM Unlearning

Explore similar work

Test-Time Unlearning via Sparse Autoencoder

Attention Smoothing Is All You Need For Unlearning

Not Every Change Is Necessary: Recoverable Drift in Large Language Model Unlearning