TRACE: Temporal Risk Assessment and Concept Editing for Safer Language Generation

Ankur Garg; Xuemin Yu; Hassan Sajjad; Samira Ebrahimi Kahou

cs.AIOct 11, 2026

TRACE: Temporal Risk Assessment and Concept Editing for Safer Language Generation

Authors: Ankur Garg, Xuemin Yu, Hassan Sajjad, Samira Ebrahimi Kahou

Organizations: University of Calgary · Dalhousie University · Mila - Quebec AI Institute

Abstract

Large language models (LLMs) can produce harmful content while generating a response. A useful safeguard should not only detect this risk during generation but also intervene while leaving safe responses largely unchanged and provide an auditable record of the internal patterns supporting its decisions. We propose TRACE (Temporal Risk Assessment and Concept Editing), a framework that tracks harmful-response risk during generation, selectively edits token activations, and leaves an auditable trace of its detection and steering decisions. TRACE learns a codebook of recurring activation patterns, with each codebook vector representing a latent concept. At inference time, it assigns each token activation to the nearest codebook vector to obtain the corresponding concept. A gated recurrent unit (GRU) processes the activations and their assigned concept features to predict completed-response harmfulness and a monotone cumulative streaming risk. During steering, activations assigned to harmful concepts are steered along directions defined by nearby benign targets. An external LLM labels the learned concepts post hoc, making detection sequences and harmful-to-benign edit records interpretable. Experiments across three generators and the WildGuard and S-Eval benchmarks show that TRACE improves streaming F1 over Kelp in five of six comparisons while remaining competitive on completed-response detection. TRACE's selective steering reduces Harmbench attack success with little change in refusal on originally safe responses. The post-hoc audit reveals recurring semantic structure in the learned concepts and distinct traces associated with harmful detection, safe generation, and activation edits.

Figures & tables

Figure 1: TRACE architecture. Learned discrete concepts support causal detection and concept-conditioned activation editing. The example uses Qwen layers and a Qwen/WildGuard trajectory.

		WildGuard		S-Eval
Generator	Method	Response F1/ FPR	Streaming F1/ FPR	Response F1/ FPR	Streaming F1/ FPR
Qwen3-8B	WildGuard	0.571/0.004	–	0.093/0.000	–
	DiffMean	0.525/0.255	0.525/0.255	0.820/0.362	0.820/0.362
	Kelp	0.839 /0.026	0.803/0.075	0.900/0.092	0.896/0.109
	TRACE	0.835/0.050	0.837 /0.045	0.933 /0.088	0.935 /0.079
Llama-3.1-8B	WildGuard	0.520/0.001	–	0.328/0.000	–

Table 1: Harmful-class F1/ FPR . Bold marks the highest F1 for each generator–dataset pair.

Figure 2: Median first-detection position and FPR at each method’s reported threshold.

Generator	Method	ASR $\downarrow$	Safe refusal(%)	ASR $\downarrow$	Safe refusal(%)
		WildGuard		S-Eval
Qwen3-8B	Unsteered	0.873	24.8	0.634	55.7
	Self-Reminder	0.289	$25.0\rightarrow 32.6$	0.182	$55.5\rightarrow 65.7$
	Refusal-direction addition	0.000	$24.9\rightarrow 96.5$	0.000	$55.7\rightarrow 99.6$
	CAST	0.278	$24.9\rightarrow 49.0$	0.167	$56.1\rightarrow 82.4$
	TRACE	0.306	$24.8\rightarrow 25.4$	0.062	$55.7\rightarrow 55.5$

Table 2: Steering results. Safe refusal shows matched unsteered

\rightarrow

steered refusal rates.

Figure 3: Ten most frequently named concepts in Qwen/WildGuard true-positive and true-negative traces. True-negative benign concepts are separated into refusal-related and non-refusal groups.

Figure 4: Most frequent named source and target concepts in the 283 edited harmful Qwen/WildGuard responses. Each concept is counted once per response.

Figure 5: Concept trajectories before and after intervention for harmful Qwen/WildGuard prompts. Responses are divided into four token intervals, with the mean edited-token fraction shown for each.

Appendix figures & tables18 assets

Supplementary material from the paper’s appendix.

Appendix

Generator	Dataset	Detector train / val	VQ train / val / test	Steering val safe / harmful	Official test safe / harmful
Qwen3-8B	WildGuard	10,719 / 1,191	9,647 / 536 / 536	200 / 200	1,441 / 284
	S-Eval	8,100 / 900	7,290 / 404 / 406	200 / 200	467 / 533
Llama-3.1-8B	WildGuard	8,327 / 925	7,494 / 416 / 417	200 / 200	1,519 / 206
	S-Eval	8,100 / 900	7,290 / 406 / 404	200 / 200	547 / 453
InternLM3-8B	WildGuard	8,556 / 951	7,701 / 428 / 427	200 / 200	1,486 / 239
	S-Eval	8,100 / 900	7,290 / 404 / 406	200 / 200	615 / 385

Appendix

Table 3: Dataset sizes used for detector training, VQ , steering validation, and official testing.

Figure 6: Five-fold WildGuard probe accuracy by read layer. Diamonds mark the selected layers.

Generator	Dataset	Hidden size	Learning rate	Weight decay	Dropout	Streaming weight	Harmfulness weight
Qwen3-8B	WildGuard	64	$2.44\times 10^{-4}$	0	0.40	1.832	1.00
	S-Eval	64	$4.37\times 10^{-4}$	$10^{-6}$	0.35	0.827	1.75
Llama-3.1-8B	WildGuard	64	$6.50\times 10^{-5}$	$10^{-2}$	0.15	0.738	1.25
	S-Eval	64	$7.15\times 10^{-5}$	$10^{-4}$	0.15	1.060	1.50
InternLM3-8B	WildGuard	64	$4.48\times 10^{-5}$	$10^{-5}$	0.05	0.838	0.75
	S-Eval	256	$2.90\times 10^{-4}$	$10^{-2}$	0.35	1.727	0.50

Appendix

Table 4: Hyperparameters used for final causal-detector training.

Outcome	$n$	Harmful	Benign	Neutral	Ambiguous	No agreed name
True positive	251	5.95	1.01	13.51	3.77	75.75
False positive	65	4.95	1.31	13.76	3.64	76.34
False negative	33	1.66	8.34	21.73	3.42	64.85
True negative	1,376	0.23	26.60	27.63	3.89	41.65

Appendix

Table 5: Mean Qwen/WildGuard concept-category share by detection outcome (%).

Figure 7: Most frequent named concepts in Qwen/WildGuard detection errors at threshold 0.5. False-negative concepts are separated into refusal-related and non-refusal groups.

Harmful source concept	Benign target concept	Responses
Instructional guides for cyberattacks and data exploitation	Information request refusal	30
Instructional content for harmful activities	References to external legal or professional resources	29
Instructional guides for illicit or harmful activities	External source and regulatory citation	27
Deceptive framing in fictional or hypothetical scenarios	Refusal of harmful requests	22
Fictional narrative framing for harmful content	Refusal of harmful requests	22

Appendix

Table 6: Named harmful-to-benign concept pairs in edited Qwen/WildGuard responses. Counts are over distinct responses.

Figure 8: Concept trajectories for 1,441 safe Qwen/WildGuard prompts, constructed as in Figure 5 .

		WildGuard		S-Eval
Generator	Method	Response	Streaming	Response	Streaming
Qwen3-8B	DiffMean	0.851/0.489	0.851/0.489	0.806/0.770	0.806/0.770
	Kelp	0.980 / 0.894	0.975/0.873	0.962/0.970	0.959/0.966
	TRACE	0.975/0.891	0.976 / 0.895	0.984 / 0.986	0.983 / 0.985
Llama-3.1-8B	DiffMean	0.897/0.520	0.897/0.520	0.964/0.948	0.964/0.948
	Kelp	0.963/0.836	0.957/0.807	0.977/0.969	0.978/0.971

Appendix

Table 7: Official-test response and streaming AUROC / AP . Bold marks the best value in each comparison.

Figure 9: Cumulative first-detection coverage at thresholds selected for a 5% validation FPR . Rows are datasets and columns are generators.

Figure 10: Harmful recall and safe-response FPR by retained-length quintile. For each method, harmful and safe responses are sorted separately by length, and each quintile contains roughly one fifth of the corresponding responses. Error bars are 95% Wilson intervals.

Target generator	Training responses	TRACE		Kelp
		Response	Streaming	Response	Streaming
Qwen3-8B	Qwen (same generator)	0.933	0.935	0.900	0.896
	Llama	0.910	0.886	0.824	0.802
	InternLM	0.933	0.928	0.875	0.887
Llama-3.1-8B	Llama (same generator)	0.966	0.961	0.916	0.919
	Qwen	0.952	0.943	0.779	0.909

Appendix

Table 8: S-Eval F1 when each target detector is trained on responses from the target generator or another generator. Activations and detector components are target-specific. Bold marks the higher value for each row.

Comparison	Variant	WildGuard		S-Eval
		Response	Streaming	Response	Streaming
Detector input	Activation only	0.836	0.844	0.929	0.934
	Concept only	0.817	0.795	0.916	0.916
	Activation and concept	0.835	0.837	0.933	0.935
Codebook size	400	0.820	0.820	0.930	0.928
	800	0.835	0.837	0.933	0.935

Appendix

Table 9: Qwen3-8B harmful-class F1 ablations on the official test sets at threshold 0.5. Bold marks the best value within each comparison.

Generator	Dataset	Output	F1	FPR	AUROC	AP
Qwen3-8B	WildGuard	Response	$0.8382\pm 0.0039$	$0.0354\pm 0.0127$	$0.9753\pm 0.0009$	$0.8976\pm 0.0058$
	WildGuard	Streaming	$0.8357\pm 0.0062$	$0.0333\pm 0.0103$	$0.9754\pm 0.0013$	$0.9010\pm 0.0065$
	S-Eval	Response	$0.9296\pm 0.0112$	$0.0607\pm 0.0279$	$0.9837\pm 0.0008$	$0.9859\pm 0.0002$
	S-Eval	Streaming	$0.9312\pm 0.0105$	$0.0592\pm 0.0216$	$0.9834\pm 0.0008$	$0.9860\pm 0.0008$
Llama-3.1-8B	WildGuard	Response	$0.7908\pm 0.0048$	$0.0369\pm 0.0040$	$0.9738\pm 0.0016$	$0.8677\pm 0.0027$
	WildGuard	Streaming	$0.7777\pm 0.0286$	$0.0424\pm 0.0130$	$0.9684\pm 0.0038$	$0.8642\pm 0.0079$

Appendix

Table 10: Detector results across seeds 42–44 (mean

\pm

sample standard deviation). F1 and FPR use threshold 0.5.

Figure 11: Validation HarmBench ASR across steering strengths. Stars mark settings used on the official test sets.

Generator / dataset	$\lambda$	ASR	Non- repetition	$\Delta$ NLL	Edited fraction	Safe broken	Use
Qwen3-8B WildGuard	0.25	0.790	0.82	0.42	0.87	3/200	Sensitivity
	0.50	0.410	0.54	0.64	0.75	10/200	Default
	0.75	0.300	0.51	0.74	0.67	11/200	Selected
	1.00	0.210	0.47	0.92	0.63	17/200	Not used
Qwen3-8B S-Eval	0.25	0.340	0.88	0.38	0.77	8/200	Sensitivity
	0.50	0.100	0.68	0.58	0.64	13/200	Selected

Appendix

Table 11: Validation ASR and quality checks for norm-preserving steering on 200 harmful and 200 safe responses. Underlining marks failed quality limits.

Dataset	Update	$\lambda$	ASR $\downarrow$	Non-repetition $\uparrow$	$\Delta$ NLL $\downarrow$	Edited fraction $\downarrow$	Safe broken $\downarrow$
WildGuard	Norm-preserving	0.75	0.300	0.51	0.74	0.67	11/200
	Direct addition	0.50	0.865	0.93	0.02	0.91	3/200
	Direct addition	1.00	0.835	0.88	0.12	0.89	3/200
	Direct addition	2.00	0.740	0.78	0.45	0.84	1/200
	Direct addition	4.00	0.455	0.57	0.68	0.74	4/200
S-Eval	Norm-preserving	0.50	0.100	0.68	0.58	0.64	13/200

Appendix

Table 12: Qwen3-8B validation comparison of norm-preserving steering and direct addition. Bold marks the best value within each dataset and metric.

Dataset	Concept weight	ASR $\downarrow$	Non-repetition $\uparrow$	$\Delta$ NLL $\downarrow$	Edited fraction $\downarrow$	Safe broken $\downarrow$
WildGuard	Harmfulness score	0.410	0.54	0.64	0.75	10/200
WildGuard	Uniform	0.170	0.43	0.92	0.64	30/200
S-Eval	Harmfulness score	0.100	0.68	0.58	0.64	13/200
S-Eval	Uniform	0.020	0.48	0.98	0.67	49/200

Appendix

Table 13: Qwen3-8B validation ablation of concept harmfulness weighting at

\lambda=0.5

. Underlining marks failed quality limits.

Generator	Dataset	Activation cache (h)	VQ training (h)	Detector training (h)
Qwen3-8B	WildGuard	0.93	9.22	0.45
	S-Eval	1.24	5.89	0.27
Llama-3.1-8B	WildGuard	0.21	4.69	0.23
	S-Eval	0.13	3.13	0.14
InternLM3-8B	WildGuard	0.52	16.18	0.41
	S-Eval	0.35	8.44	0.38

Appendix

Table 14: Wall-clock hours for the final TRACE runs. Activation caching and VQ training use one H100 with 40 GB of available memory; detector training uses one H100 with 10 GB.

LLM Safety

Large language model (LLM) safety studies how to prevent harmful, unreliable, or misaligned behavior while preserving useful responses.

See all 426 papers

ForeSight: Enhancing Risk Monitoring via Early Safety Signal Distillation
Sep 12, 2026Hanling Wang, Chenlong Wei, Ling Xu +6LLM Safety Monitoring LLM Safety
University of California, San Diego · Xi’an Jiaotong-Liverpool University · East China Normal University +2
TRACE: Trajectory Return Attribution and Contrastive Erasure for Multi-Turn Safety
Oct 1, 2026Fengpeng Li, Kemou Li, Qizhou Wang +3LLM Fine-Tuning Adversarial Attacks on LLMs
PRADA Lab, King Abdullah University of Science and Technology · State Key Laboratory of Internet of Things for Smart City, University of Macau · Imperfect Information Learning Team, RIKEN Center for Advanced Intelligence Project +1
TRACE: An Evidence-Grounded Benchmark for Safety Evaluation of Large Reasoning Models
Aug 25, 2026Zhenyu Wu, Siyuan Chen, Changchun Yang +8LLM Safety Benchmarks Language Model Safety Evaluation
King Abdullah University of Science and Technology · Aramco · KAUST

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

ForeSight: Enhancing Risk Monitoring via Early Safety Signal Distillation

Hanling Wang, Chenlong Wei, Ling Xu, Hanyan Niu+5

University of California, San Diego · Xi’an Jiaotong-Liverpool University · East China Normal University+2

LLM Safety Monitoring LLM Safety Language Model Safety Evaluation

Open paper AskPDF arXiv

As large language models (LLMs) are increasingly deployed, the generation of harmful content has become a critical safety concern. Existing safeguards operate at the input, output, or streaming-generation stages, while early-risk methods that rely on surface tokens or output logits may suffer from weak initial signals, and internals-based detectors using dense representations may retain highly entangled and redundant safety-irrelevant information. It therefore remains unclear whether the earliest post-generation hidden states already contain reliable signals about final-response harmfulness. To address this gap, we propose ForeSight, a first-token output-risk forecasting framework that distills weak and redundant early safety signals into compact, layer-aware risk representations. Experiments on five safety benchmarks and two target models demonstrate that ForeSight achieves superior and efficient early-risk forecasting while relying solely on first-token hidden states. The code is available at: https://github.com/Scabbards1500/Foresight

Abstract

LLM Safety

Explore similar work

ForeSight: Enhancing Risk Monitoring via Early Safety Signal Distillation

TRACE: Trajectory Return Attribution and Contrastive Erasure for Multi-Turn Safety

TRACE: An Evidence-Grounded Benchmark for Safety Evaluation of Large Reasoning Models