Table of Contents
Fetching ...

Safety Alignment Depth in Large Language Models: A Markov Chain Perspective

Ching-Chia Kao, Chia-Mu Yu, Chun-Shien Lu, Chu-Song Chen

TL;DR

This work reframes LLM safety as a Markov-chain problem by introducing safety depth, the designated refusal state at which generation of harmful content is prevented. It shows that iterative fine-tuning with a structured bias can render refusal states arbitrarily close to absorbing, providing formal guarantees for safe behavior. By incorporating cyclic group actions for data augmentation, the authors prove that safety depth bounds can be maintained under symmetry-induced perturbations, and that broader ensembles can compensate for shallower alignments. The combination of a theoretical foundation and empirical validation on toy chains and open-source LLMs offers a practical blueprint for designing scalable, robust safety mechanisms that complement existing alignment methods.

Abstract

Large Language Models (LLMs) are increasingly adopted in high-stakes scenarios, yet their safety mechanisms often remain fragile. Simple jailbreak prompts or even benign fine-tuning can bypass these protocols, underscoring the need to understand where and how they fail. Recent findings suggest that vulnerabilities emerge when alignment is confined to only the initial output tokens. Unfortunately, even with the introduction of deep safety alignment, determining the optimal safety depth remains an unresolved challenge. By leveraging the equivalence between autoregressive language models and Markov chains, this paper offers the first theoretical result on how to identify the ideal depth for safety alignment, and demonstrates how permutation-based data augmentation can tighten these bounds. Crucially, we reveal a fundamental interaction between alignment depth and ensemble width-indicating that broader ensembles can compensate for shallower alignments. These insights provide a theoretical foundation for designing more robust, scalable safety strategies that complement existing alignment approaches, opening new avenues for research into safer, more reliable LLMs.

Safety Alignment Depth in Large Language Models: A Markov Chain Perspective

TL;DR

This work reframes LLM safety as a Markov-chain problem by introducing safety depth, the designated refusal state at which generation of harmful content is prevented. It shows that iterative fine-tuning with a structured bias can render refusal states arbitrarily close to absorbing, providing formal guarantees for safe behavior. By incorporating cyclic group actions for data augmentation, the authors prove that safety depth bounds can be maintained under symmetry-induced perturbations, and that broader ensembles can compensate for shallower alignments. The combination of a theoretical foundation and empirical validation on toy chains and open-source LLMs offers a practical blueprint for designing scalable, robust safety mechanisms that complement existing alignment methods.

Abstract

Large Language Models (LLMs) are increasingly adopted in high-stakes scenarios, yet their safety mechanisms often remain fragile. Simple jailbreak prompts or even benign fine-tuning can bypass these protocols, underscoring the need to understand where and how they fail. Recent findings suggest that vulnerabilities emerge when alignment is confined to only the initial output tokens. Unfortunately, even with the introduction of deep safety alignment, determining the optimal safety depth remains an unresolved challenge. By leveraging the equivalence between autoregressive language models and Markov chains, this paper offers the first theoretical result on how to identify the ideal depth for safety alignment, and demonstrates how permutation-based data augmentation can tighten these bounds. Crucially, we reveal a fundamental interaction between alignment depth and ensemble width-indicating that broader ensembles can compensate for shallower alignments. These insights provide a theoretical foundation for designing more robust, scalable safety strategies that complement existing alignment approaches, opening new avenues for research into safer, more reliable LLMs.

Paper Structure

This paper contains 58 sections, 6 theorems, 132 equations, 8 figures, 4 tables.

Key Result

Proposition 3.3

For a finite group $G$, its action on the training set can be represented as a permutation $\sigma_g$:

Figures (8)

  • Figure 1: Permutations of phrases used for data augmentation. The top row represents a cyclic group, while the bottom row, as proposed by qi2024safetyalignmentjusttokens, is non-cyclic.
  • Figure 2: Visualization of $\delta$-absorbing. At $t=0$, all states (both refusal and regular) can transition relatively freely between each other. The transition probabilities are determined by the initial matrix $Q_0$. At $t=k$ (where $k$ satisfies the Theorem \ref{['thm:safety']} and Corollary \ref{['thm:safety2']}), the refusal states have a thicker self-loop, which means a very high probability.
  • Figure 3: Single model convergence showing exponential decay in blue line with confidence intervals over 50 bias applications, demonstrating reliable convergence to safe behavior. Cyclic group action convergence is displayed in a red line, illustrating stable convergence despite periodic fluctuations.
  • Figure 4: Comparison of ensemble combination methods (Union, Average, and Majority) showing escape probabilities, where box plots indicate the distribution of outcomes and individual points show specific results.
  • Figure 5: Gemma safety score comparison. Each bar indicates the model’s average safety score for that category.
  • ...and 3 more figures

Theorems & Definitions (28)

  • Definition 3.2: Augmented Training Set
  • Proposition 3.3
  • Remark 3.5
  • Definition 4.1: Safety Alignment
  • Definition 4.2: Safe Alignment in Markov View
  • Remark 4.3
  • Theorem 4.5: $\delta$-absorbing
  • proof
  • Remark 4.6
  • Corollary 4.7: Largest Safety Depth That Becomes $\delta$-Absorbing
  • ...and 18 more