TRACEALIGN -- Rastreando la Deriva: Atribución de Fallos de Alineación a Fuentes de Creencias Durante el Entrenamiento en LLMs

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

August 4, 2025

Autores: Amitava Das, Vinija Jain, Aman Chadha

cs.AI

Resumen

Los Modelos de Lenguaje de Gran Escala (LLMs, por sus siglas en inglés) ajustados para alinearse con los valores humanos a menudo exhiben desviación en la alineación, produciendo respuestas inseguras o que violan políticas cuando se exponen a indicaciones adversas, perturbaciones en la decodificación o variaciones de jailbreaks parafraseados. Si bien trabajos previos han caracterizado conductualmente los fallos de alineación, se sabe poco sobre las fuentes de creencias durante el entrenamiento que subyacen a estos fallos. Presentamos TraceAlign, un marco unificado para rastrear respuestas inseguras hasta sus causas raíz en el corpus de entrenamiento del modelo. Central a nuestro enfoque es el Índice de Conflicto de Creencias (BCI, por sus siglas en inglés), que cuantifica la inconsistencia semántica entre los segmentos generados y las políticas alineadas, basándose en documentos de entrenamiento recuperados mediante coincidencia de arreglos de sufijos. Proponemos tres intervenciones complementarias: (i) TraceShield, un filtro de seguridad en tiempo de inferencia que rechaza respuestas con segmentos de alto BCI, (ii) Pérdida de Desconflicto de Creencias Contrastiva, un objetivo de ajuste fino contrastivo que penaliza continuaciones de alto BCI durante el DPO, y (iii) Prov-Decode, una estrategia de decodificación consciente de la procedencia que veta expansiones de haz predichas para producir segmentos de alto BCI. Juntas, estas defensas reducen la desviación en la alineación hasta en un 85% en nuestro Punto de Referencia de Desviación de Alineación (ADB, por sus siglas en inglés) mientras preservan la utilidad en tareas estándar, con un delta menor a 0.2 y una calidad de rechazo mejorada. Además, derivamos un límite superior teórico sobre la probabilidad de desviación mediante estadísticas de segmentos de arreglos de sufijos, vinculando la frecuencia y longitud de memorización con el riesgo de reactivación adversaria. TraceAlign proporciona así el primer kit de herramientas escalable, rastreable y fundamentado para comprender y mitigar los fallos de alineación en su origen. Para fomentar una mayor exploración y desarrollo, hemos hecho público nuestro código en: https://anonymous.4open.science/r/tracealign-2DA7.

English

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