# TRACEALIGN -- Tracing the Drift: Attributing Alignment Failures to Training-Time Belief Sources in LLMs

arXiv: https://arxiv.org/abs/2508.02063 · HTML: https://mlsanity.com/paper/2508.02063/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in LLM preference optimization methods, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 36 of its 70 references that are in the map.
- First posted on arXiv 4 Aug 2025.

## Placed in

- LLM alignment / RLHF / preference optimisation / [LLM preference optimization methods](https://mlsanity.com/alignment/llm-preference-optimization-methods/)
- LLM safety and security / [Safety alignment, refusal and guardrails](https://mlsanity.com/safety/safety-alignment-refusal-and-guardrails/)

## Builds on (most cited first, 30 of 36)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [Universal and Transferable Adversarial Attacks on Aligned Language Models](https://mlsanity.com/paper/2307.15043/) (2023)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [A Baseline for Detecting Misclassified and Out-of-Distribution Examples in Neural Networks](https://mlsanity.com/paper/1610.02136/) (2016)
- [G-Eval: NLG Evaluation using GPT-4 with Better Human Alignment](https://mlsanity.com/paper/2303.16634/) (2023)
- [SimCSE: Simple Contrastive Learning of Sentence Embeddings](https://mlsanity.com/paper/2104.08821/) (2021)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [A Simple Unified Framework for Detecting Out-of-Distribution Samples and Adversarial Attacks](https://mlsanity.com/paper/1807.03888/) (2018)
- [Enhancing The Reliability of Out-of-distribution Image Detection in Neural Networks](https://mlsanity.com/paper/1706.02690/) (2017)
- [Energy-based Out-of-distribution Detection](https://mlsanity.com/paper/2010.03759/) (2020)
- [RealToxicityPrompts: Evaluating Neural Toxic Degeneration in Language Models](https://mlsanity.com/paper/2009.11462/) (2020)
- [LIMA: Less Is More for Alignment](https://mlsanity.com/paper/2305.11206/) (2023)
- [A General Language Assistant as a Laboratory for Alignment](https://mlsanity.com/paper/2112.00861/) (2021)
- [Mass-Editing Memory in a Transformer](https://mlsanity.com/paper/2210.07229/) (2022)
- [Adversarial NLI: A New Benchmark for Natural Language Understanding](https://mlsanity.com/paper/1910.14599/) (2019)
- [Extracting Training Data from Diffusion Models](https://mlsanity.com/paper/2301.13188/) (2023)
- [Universal Adversarial Triggers for Attacking and Analyzing NLP](https://mlsanity.com/paper/1908.07125/) (2019)
- [Latent Retrieval for Weakly Supervised Open Domain Question Answering](https://mlsanity.com/paper/1906.00300/) (2019)
- [RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignment](https://mlsanity.com/paper/2304.06767/) (2023)
- [Documenting Large Webtext Corpora: A Case Study on the Colossal Clean Crawled Corpus](https://mlsanity.com/paper/2104.08758/) (2021)
- [Does Learning Require Memorization? A Short Tale about a Long Tail](https://mlsanity.com/paper/1906.05271/) (2019)
- [GeDi: Generative Discriminator Guided Sequence Generation](https://mlsanity.com/paper/2009.06367/) (2020)
- [Diverse Beam Search: Decoding Diverse Solutions from Neural Sequence Models](https://mlsanity.com/paper/1610.02424/) (2016)

## Abstract (arXiv)

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