TraceShield: Inference-Time Safety Filter

• TraceShield is an inference-time safety filter that traces generated completions to unsafe training spans and computes a Belief Conflict Index (BCI) to enforce refusals.
• It employs a scalable suffix-array based TraceIndex with frequency filtering to link unsafe spans with detailed provenance metadata for auditability.
• Evaluations show TraceShield significantly reduces alignment drift while maintaining low runtime overhead and minimal false-positive rates.

TraceShield is the inference-time defense module of TraceAlign, a provenance-grounded framework for analyzing and mitigating alignment drift in LLMs. Its defining operation is to trace generated completions back to an unsafe training index, compute a span-level risk score called the Belief Conflict Index (BCI), and refuse completions whose highest-risk retrieved span exceeds a calibrated threshold. In the paper’s formulation, TraceShield is an “inference-time safety filter,” an “inference-time belief-guided refusal,” and a mechanism that “blocks completions grounded in unsafe memory” under jailbreaks, adversarial prompts, decoding perturbations, and paraphrased policy evasion (Das et al., 4 Aug 2025).

1. Position within the TraceAlign framework

TraceShield is not presented as an independent theory detached from the rest of TraceAlign. It is one of three coordinated interventions introduced in the framework: TraceShield for inference-time filtering, Contrastive Belief Deconfliction (CBD) Loss for training-time regularization, and Prov-Decode for provenance-aware decoding. The paper’s central claim is that many unsafe completions are not merely behavioral anomalies, but the resurfacing of training-time belief fragments retrieved from unsafe or permissive data. TraceShield operationalizes that claim at inference time by converting provenance retrieval into a refusal policy (Das et al., 4 Aug 2025).

The module targets alignment drift specifically. In the paper’s usage, this includes unsafe completions, policy-violating outputs, jailbreak robustness failures, and unsafe reactivation of memorized pretraining beliefs. Its action space is intentionally narrow: it does not rewrite responses, redact partial spans, or re-rank a beam. Instead, it performs span-level attribution and then applies a completion-level hard refusal rule. This makes it the runtime component of a larger provenance stack rather than a generic safety classifier.

Within TraceAlign, the division of labor is explicit. CBD Loss alters the training objective, while Prov-Decode vetoes risky continuations during beam expansion. TraceShield, by contrast, requires no retraining and is primarily described as a post-generation or completion-time filter. A plausible implication is that the design deliberately separates provenance diagnosis, inference-time enforcement, and training-time deconfliction so that each mechanism can be deployed independently or in combination.

2. Provenance machinery: TraceIndex and unsafe-span attribution

TraceShield depends on a provenance index called TraceIndex. The index is suffix-array-based and is built over an unsafe corpus, denoted in the paper as $\mathcal{D}_{\text{unsafe}}$. In one formulation, the indexed suffix set is

$$\Sigma = \bigcup_{i=1}^N \bigcup_{j=1}^{L_i} T_i[j:L_i],$$

where $T_i$ is the tokenization of document $d_i$. A query span $q = [t_1,\dots,t_k]$ is matched by searching for suffixes $s \in \Sigma$ such that $s[1:k] = q$ (Das et al., 4 Aug 2025).

The appendix gives a generalized suffix-array view over the concatenated corpus token stream and reports retrieval complexity as $\mathcal{O}(k + \log n)$, while the main text also states substring retrieval of length $k$ in $O(k \log S)$, with $$\Sigma = \bigcup_{i=1}^N \bigcup_{j=1}^{L_i} T_i[j:L_i],$$0 the total number of suffixes. The paper therefore treats TraceIndex as a scalable provenance substrate rather than a small ad hoc lookup table. The retrieved output is not only a matched string but attached provenance metadata, including source document, offsets, and origin information. This is what gives TraceShield its auditable character: refusal can be tied to a specific matched unsafe span and its source.

A second important statistic is the span match count

$$\Sigma = \bigcup_{i=1}^N \bigcup_{j=1}^{L_i} T_i[j:L_i],$$1

This quantifies how frequently a candidate span occurs in the indexed corpus. The paper uses a maximum-frequency filter $$\Sigma = \bigcup_{i=1}^N \bigcup_{j=1}^{L_i} T_i[j:L_i],$$2 to keep rare, attributable spans and discard boilerplate or noisy matches. It reports that this filter reduced noisy matches by 62% and improved trace precision by 18% on the Alignment Drift Benchmark (ADB) (Das et al., 4 Aug 2025).

This provenance layer is the main technical distinction between TraceShield and classifier-only refusal. The filter does not merely decide whether a completion “looks unsafe.” It asks whether the completion contains retrievable, low-frequency, provenance-linked spans associated with unsafe training data. That design choice also explains the system’s dependence on corpus access: without an indexable unsafe corpus or a surrogate thereof, TraceShield loses the mechanism that makes its refusal grounded rather than purely behavioral.

3. Belief Conflict Index and the refusal criterion

Given a completion

$$\Sigma = \bigcup_{i=1}^N \bigcup_{j=1}^{L_i} T_i[j:L_i],$$3

TraceShield extracts spans, retrieves matching spans $$\Sigma = \bigcup_{i=1}^N \bigcup_{j=1}^{L_i} T_i[j:L_i],$$4 from the unsafe index, and computes a BCI score for each retrieved span. The paper defines BCI for a matched span $$\Sigma = \bigcup_{i=1}^N \bigcup_{j=1}^{L_i} T_i[j:L_i],$$5 as

$$\Sigma = \bigcup_{i=1}^N \bigcup_{j=1}^{L_i} T_i[j:L_i],$$6

Its normalized form is

$$\Sigma = \bigcup_{i=1}^N \bigcup_{j=1}^{L_i} T_i[j:L_i],$$7

The operational intuition is that BCI measures rarity, lexical specificity, memorization likelihood, and what the paper calls “belief density.” Long rare spans are treated as especially risky because their reproduction is taken to be more indicative of memorized unsafe fragments than of ordinary generative behavior (Das et al., 4 Aug 2025).

The paper further interprets normalized BCI as cross-entropy against the training unigram distribution. Let $$\Sigma = \bigcup_{i=1}^N \bigcup_{j=1}^{L_i} T_i[j:L_i],$$8 be the empirical unigram distribution over span $$\Sigma = \bigcup_{i=1}^N \bigcup_{j=1}^{L_i} T_i[j:L_i],$$9. Then

$T_i$0

with decomposition

$T_i$1

The appendix also invokes Pinsker’s inequality,

$T_i$2

to motivate the claim that high nBCI corresponds to statistically atypical spans and elevated epistemic risk. The paper does not formalize BCI as a direct semantic similarity to a policy document; rather, policy conflict enters through the provenance decision rule, which treats high-BCI spans retrieved from unsafe sources as alignment-relevant conflict evidence.

TraceShield aggregates span scores at the completion level using the maximum:

$T_i$3

Its refusal rule is then

$T_i$4

with $T_i$5, described as empirically calibrated (Das et al., 4 Aug 2025).

The worked example in the paper is the span

$T_i$6

with token frequencies

$T_i$7

The paper computes $T_i$8, and therefore above threshold; the normalized value given in the appendix is $T_i$9 (Das et al., 4 Aug 2025).

4. Inference-time operation and relation to the other TraceAlign modules

TraceShield’s runtime procedure is concise. Given a candidate completion, it retrieves top-$d_i$0 matched spans from the unsafe index, computes BCI for each, forms $d_i$1, and refuses if the maximum exceeds threshold. The primitive scoring object is token-additive, the principal analysis unit is the span, and the final decision object is the completion. This is why the paper characterizes TraceShield as span-level analysis with completion-level enforcement (Das et al., 4 Aug 2025).

The module is therefore hard-refusal rather than soft filtering. The paper describes it with phrases such as “blocks completions,” “refuses the response,” and “block completions containing high-BCI spans.” It does not specify a canonical outward refusal template, nor does it describe regeneration as part of the core mechanism. The appendix notes that refusal can include the matched span, the BCI score, and a source pointer for auditability, but that is framed as interpretability support rather than as a mandatory interaction design.

Its relation to the other two TraceAlign interventions clarifies the architectural intent. CBD Loss adds a regularization term during DPO-style fine-tuning,

$d_i$2

and total objective

$d_i$3

Prov-Decode acts earlier, modifying beam search by vetoing candidate continuations predicted to form high-BCI spans:

$d_i$4

TraceShield differs from both by requiring no retraining and by acting after a candidate response is available. This suggests that it is the most deployment-friendly element of the TraceAlign stack, particularly when only inference-time access is available.

5. Benchmarking and reported performance

The main evaluation benchmark is ADB, a curated Alignment Drift Benchmark of 5,200 adversarial prompts across five high-risk domains: Explosives / Weapons (1,000), Cybercrime / Hacking (1,200), Self-Harm / Medical Misinformation (1,000), Hate / Harassment (1,000), and Illicit Finance / Fraud (1,000). The benchmark is built by starting from safe alignment-era prompts, adversarially rewriting them with GPT-4 using fictional, academic, or indirect framing, filtering with G-Eval for safety risk, validating across multiple aligned LLMs, and annotating domain, severity, and provenance metadata (Das et al., 4 Aug 2025).

On the OLMo-2 ablation reported in the paper, the standalone “T only” condition for TraceShield reduces drift from 41.8% to 14.6%, with $d_i$5, refusal quality 4.3, and false-positive rate 2.1%. The main text also summarizes the same module as cutting alignment drift from 42.1% to 14.6% while preserving utility. In the same ablation, TraceShield alone performs better than CBD alone on drift, though slightly worse than Prov-Decode alone; the combined $d_i$6 condition reaches drift 6.2% with refusal quality 4.7 and FPR 2.7% (Das et al., 4 Aug 2025).

Across model families, the paper also reports attack-success-rate reduction using TraceShield alone:

Model	Baseline Drift Rate	ASR Before	ASR After
LLaMA-2-Chat-13B	43.5%	76.8%	28.7%
OLMo-2-32B-Instruct	41.8%	75.2%	26.1%
GPT-NeoX-Aligned	40.9%	73.9%	27.5%

The runtime overhead reported for TraceShield is under 80 ms per 100-token output on CPU, while a FAQ-style appendix characterizes this more loosely as about 100 ms latency per query. The paper presents this as a low-latency inference-time defense that substantially reduces drift without retraining and with only minimal perplexity impact (Das et al., 4 Aug 2025).

6. Limitations, failure modes, and broader shielding context

TraceShield’s limitations follow directly from its provenance design. The paper identifies lexical rigidity in TraceIndex as a major failure mode: semantically equivalent but lexically divergent paraphrases may evade exact or prefix-based suffix-array matching. It may also miss unsafe generations that are novel recombinations rather than retrievable substrings. Conversely, the unigram-based BCI can over-penalize rare but benign technical phrases; the example given is “lithium carbonate titration curve.” The reported T-only false-positive rate of 2.1% should therefore be read as low but not negligible (Das et al., 4 Aug 2025).

A second limitation is deployment dependence on corpus access. The paper states that TraceAlign relies on span-level access to training data to construct TraceIndex and compute BCI. For closed-source or API-only models, exact provenance grounding is unavailable unless one substitutes a surrogate unsafe corpus. The same section also notes that the framework does not cleanly distinguish whether a problematic belief source arose in pretraining or alignment tuning. The theoretical discussion is also suggestive rather than fully optimality-proving: the paper motivates the threshold rule via memorization pressure, span length, and information-theoretic rarity, but it does not prove that the hard threshold at $d_i$7 is optimal.

Within the broader shielding literature, TraceShield occupies a distinct niche. Dynamic shielding in control systems computes runtime-adaptive safety enforcers over parametric safety specifications and updates them by controller product plus nonblocking pruning rather than recomputation (Corsi et al., 28 May 2025). SHIELD on multi-tenant FPGAs is an adaptive power-trace flattening defense that uses ring oscillators to detect victim-induced power fluctuations and inject compensating noise (Ahmadi et al., 2023). ShieldMMU suppresses page-fault trace leakage in SGX-like systems by authenticating PTE metadata inside the MMU and restoring attacked present bits before the OS can observe the fault (Liu et al., 4 Sep 2025). SurrogateShield mediates prompt/response flows locally by replacing PII with type-consistent surrogates and restoring originals after inference (Jathanna, 28 Jun 2026). Local shield synthesis for Dec-POMDPs compiles finite-state safety processes into history-sensitive local Mealy-machine shields that depend on observation history rather than only current observation (Yang et al., 8 Apr 2026). Against that background, TraceShield is specifically an LLM inference-time refusal mechanism grounded in training-data provenance, not a generic runtime shield, trace obfuscator, or privacy proxy.

The conceptual commonality across these systems is that “shielding” denotes intervention between an untrusted generator of actions or outputs and a protected execution environment. What distinguishes TraceShield is the form of evidence it uses: not physical dynamics, not page-table integrity, not prompt redaction, and not decentralized control history, but retrieved unsafe training spans scored by a provenance-based conflict index. This places it squarely in the emerging class of memory- and provenance-grounded LLM safety mechanisms.