TRACEALIGN -- Tracing the Drift: Attributing Alignment Failures to
  Training-Time Belief Sources in LLMs
Large Language Models (LLMs) fine-tuned to align with human values often
exhibit alignment drift, producing unsafe or policy-violating completions when
exposed to adversarial prompts, decoding perturbations, or paraphrased
jailbreaks. While prior work has behaviorally characterized alignment failure,
little is known about the training-time belief sources underlying these
failures. We introduce TraceAlign, a unified framework for tracing unsafe
completions back to their root causes in the model's training corpus. Central
to our approach is the Belief Conflict Index (BCI), which quantifies semantic
inconsistency between generated spans and aligned policies, based on retrieved
training documents using suffix-array matching. We propose three complementary
interventions: (i) TraceShield, an inference-time safety filter that refuses
completions with high-BCI spans, (ii) Contrastive Belief Deconfliction Loss, a
contrastive fine-tuning objective penalizing high-BCI continuations during DPO,
and (iii) Prov-Decode, a provenance-aware decoding strategy that vetoes beam
expansions predicted to yield high-BCI spans. Together, these defenses reduce
alignment drift by up to 85% on our curated Alignment Drift Benchmark (ADB)
while preserving utility on standard tasks, with delta less than 0.2 and
improved refusal quality. We further derive a theoretical upper bound on drift
likelihood via suffix-array span statistics, linking memorization frequency and
length to adversarial reactivation risk. TraceAlign thus provides the first
scalable, traceable, and grounded toolkit for understanding and mitigating
alignment failures at source. To encourage further exploration and development,
we open-source our implementation at:
https://anonymous.4open.science/r/tracealign-2DA7
6892c3aa8da45ffb0a2b242c
Amitava Das
6892c3aa8da45ffb0a2b242d
Vinija Jain
6892c3aa8da45ffb0a2b242e
Aman Chadha
claimed_verified
2025-08-06T19:17:14.301Z
63a4754927f1f64ed7238dac
https://cdn-avatars.huggingface.co/v1/production/uploads/63a4754927f1f64ed7238dac/aH-eJF-31g4vof9jv2gmI.jpeg
Aman Chadha
user
amanchadha
TRACEALIGN -- Tracing the Drift: Attributing Alignment Failures to
  Training-Time Belief Sources in LLMs
Large Language Models (LLMs) fine-tuned to align with human values often
exhibit alignment drift, producing unsafe or policy-violating completions when
exposed to adversarial prompts, decoding perturbations, or paraphrased
jailbreaks. While prior work has behaviorally characterized alignment failure,
little is known about the training-time belief sources underlying these
failures. We introduce TraceAlign, a unified framework for tracing unsafe
completions back to their root causes in the model's training corpus. Central
to our approach is the Belief Conflict Index (BCI), which quantifies semantic
inconsistency between generated spans and aligned policies, based on retrieved
training documents using suffix-array matching. We propose three complementary
interventions: (i) TraceShield, an inference-time safety filter that refuses
completions with high-BCI spans, (ii) Contrastive Belief Deconfliction Loss, a
contrastive fine-tuning objective penalizing high-BCI continuations during DPO,
and (iii) Prov-Decode, a provenance-aware decoding strategy that vetoes beam
expansions predicted to yield high-BCI spans. Together, these defenses reduce
alignment drift by up to 85% on our curated Alignment Drift Benchmark (ADB)
while preserving utility on standard tasks, with delta less than 0.2 and
improved refusal quality. We further derive a theoretical upper bound on drift
likelihood via suffix-array span statistics, linking memorization frequency and
length to adversarial reactivation risk. TraceAlign thus provides the first
scalable, traceable, and grounded toolkit for understanding and mitigating
alignment failures at source. To encourage further exploration and development,
we open-source our implementation at:
https://anonymous.4open.science/r/tracealign-2DA7
TraceAlign is a framework that identifies and mitigates alignment drift in LLMs by tracing unsafe completions to their training sources and applying interventions to reduce drift while maintaining utility.
Large Language Models alignment drift adversarial prompts decoding perturbations paraphrased jailbreaks TraceAlign Belief Conflict Index suffix-array matching TraceShield Contrastive Belief Deconfliction Loss Prov-Decode Alignment Drift Benchmark memorization frequency adversarial reactivation risk