Inference-Time Vulnerability Beyond Shallow Safety: Alignment Along Generation Trajectories

The rapid development of large language models (LLMs) like GPT-3, Llama, and others has transformed AI applications, enabling high-quality text generation. Early safety efforts, such as RLHF and Constitutional AI, improved alignment by fine-tuning models to avoid harmful outputs. However, these methods mainly focus on static output responses, leaving models vulnerable during the inference process itself. Recent studies reveal that models tend to be safe only in initial tokens, with safety diminishing as generation proceeds, especially under adversarial interventions like suffix or prefilling attacks. Existing defenses at the internal representation level or input filtering are insufficient against dynamic, mid-sequence perturbations. This gap underscores the need for training models to understand and recover from disruptions during generation, ensuring safety throughout the process.

The core challenge is that current safety alignment techniques do not account for the dynamic nature of generation. Models are susceptible to short token injections at arbitrary steps, which can redirect their output toward harmful content. This vulnerability is exacerbated by the fact that internal representations, such as refusal states, do not reliably predict robustness. Consequently, models can occupy a refusal-like internal state yet still be redirected by minimal perturbations. Addressing this requires a paradigm shift from static output-focused training to dynamic trajectory-aware training, enabling models to recognize and recover from mid-generation disruptions effectively.

The key innovations include: 1) constructing training data by simulating mid-sequence token injections to generate diverse trajectories; 2) defining refusal directions in hidden state space to guide recovery; 3) employing a preference optimization (SimPO) framework that trains models to favor safe trajectories over unsafe ones; 4) implementing iterative augmentation, where each training cycle refines the model’s ability to resist and recover from perturbations. This approach fundamentally differs from prior methods by focusing on the entire generation process, not just output or static internal states, thus enabling models to maintain safety dynamically.

• �� Analyze hidden states during generation to identify refusal and accept directions via prototypes.
• �� Monitor cosine similarity between hidden states and these prototypes during decoding.
• �� Trigger token injection when similarity exceeds thresholds, simulating harmful or refusal perturbations.
• �� Construct augmented trajectories by inserting tokens at critical steps, capturing transitions from safe to unsafe behaviors.
• �� Use a preference-based loss (SimPO) to train models to recover safe trajectories post-perturbation.
• �� Perform multiple iterative rounds, recalculating prototypes and augmenting data, to progressively enhance robustness.

Experiments involve three models—Llama-3.1-8B, Mistral-7B, Qwen-7B—evaluated on in-domain AdvBench and out-of-domain benchmarks. Attack success rates (ASR) are measured under various injection scenarios. Baselines include existing defenses like Circuit Breakers, SafeProbing, and DPO. The training involves simulating mid-sequence injections at points where models enter refusal states, then fine-tuning with the augmented trajectories. Ablation studies assess the impact of injection timing, phrase semantics, and multiple injections. Results show that trajectory alignment reduces ASR from over 90% to below 5%, with iterative training further improving robustness across attack types.

The proposed method achieves a dramatic reduction in attack success rates, demonstrating that models can be trained to resist mid-sequence perturbations effectively. Across multiple models and attack scenarios, ASR drops from over 90% to near zero. Iterative training enhances resilience, with subsequent rounds further lowering ASR and improving out-of-domain performance. The approach outperforms existing defenses, especially in complex, unseen attack types, validating the importance of training on dynamic trajectories rather than static states. Hidden state analysis confirms that robustness correlates with reduced sensitivity to perturbations, not just internal state alignment.

This trajectory-based safety alignment can be integrated into deployment pipelines for AI assistants, content moderation, and sensitive decision-making systems. It provides a scalable way to ensure safety during real-time generation, especially in adversarial environments. Long-term, this approach can be extended to multi-modal models, multi-task learning, and reinforcement learning settings, fostering the development of AI systems that are inherently resilient against manipulation and malicious interventions.

The primary limitation is the high computational cost of simulating and training on numerous trajectories, especially for very large models. The method currently targets known attack types, and its effectiveness against novel or unforeseen attacks needs further validation. Additionally, extreme perturbations may still challenge the recovery mechanism, indicating room for integrating stronger robustness techniques. Future work should focus on reducing training overhead and broadening attack coverage.