Safety Depth in Large Language Models: A Markov Chain Perspective

[Q1] A comparison of the testing time when different methods are used to make predictions on the dataset?

The average inference time for each ensemble method is shown in the table below. All methods operate within a practical and efficient range, with average prediction times well below one millisecond per sample.

These results indicate that all ensemble methods are highly efficient and suitable for practical deployment.

[Q2] Experimental results on other datasets.

We further demonstrate the transferability of our method by applying it to two additional public datasets:

1. Real Toxicity Prompts (EMNLP 2020)
2. Harmful Behaviors (HuggingFace)

The table below summarizes the results across three different models. For each setting, our cyclic augmentation consistently outperforms the baseline, resulting in improved safety scores on both datasets:

[Q3] Experimental results on larger LLMs such as Gemma3-27B

We report the overall average safety score and standard deviation for both Gemma3-27B and its cyclically-augmented counterpart across all evaluated categories:

As shown in the table, Gemma3-27B-cyclic achieves a higher average safety score and demonstrates slightly lower variance compared to the baseline model. This indicates that cyclic augmentation not only improves overall safety performance but also results in more stable and consistent outcomes across different categories.

[W1] The lack of explicit mention and discussion of the assumption in Definition 4.4.

We acknowledge that Definition 4.4 is a strong assumption. However, it has been relaxed to a more realistic $\delta$-absorbing property (Theorem 4.5) in all the following results. $\delta$-absorbing is not the one where the probability of leaving is zero; instead, it is a state where the probability of transitioning to any non-refusal state is arbitrarily small ($\leq \delta$).

---

[W2] The modeling process appears to consider “not entering an absorbing state” as the sole guarantee for helpfulness. This is insufficient to ensure that the model behaves helpfully. Furthermore, the current modeling approach seems to have certain limitations.

To clarify: we do NOT equate "not entering refusal" with being helpful; rather, we use escape probability from refusal as a necessary (but not sufficient) condition for preserving helpful behavior. This is why we separately evaluate helpfulness on downstream tasks (Table 2, Appendix I.2).

To address continuum behaviors, we note that our framework can be extended: new states can be introduced, where the transition bias reflects a gradient of utility. For instance, one could define a reward-weighted transition matrix that interpolates between safety and helpfulness objectives. We see this as a promising direction and will add this to Appendix L as a limitation and future extension.

---

[W3] It remains to be seen whether the method maintains its advantages when using full-parameter supervised fine-tuning (SFT). Additionally, can this approach be effectively integrated as a complement to existing RLHF algorithms?

Yes. Our approach is compatible with existing RLHF pipelines. Specifically, cyclic augmentation can be applied during the SFT phase to pre-condition the model with diverse refusal positions before preference modeling. Alternatively, it can serve as a post-processing refinement step after RLHF to reinforce refusal consistency at various depths. Since our method operates at the data level, it is model-agnostic and does not interfere with reward modeling or policy optimization stages.

[W1][Q1] Formally define safety depth.

Let the LLM be represented by a discrete-time Markov chain with a finite state space $S$ and a transition matrix $Q_t$ at time $t$.

Let $\mathcal{R} \subset S$ be a finite, non-empty set of designated **refusal-state indices. A safety depth, denoted by $r$, is an index corresponding to a specific refusal state, such that $r \in \mathcal{R}$.

A safety depth $r$ is considered secure at training step $T$ if its corresponding refusal state $r$ becomes δ-absorbing. This means that for a given small positive constant $\delta$, the transition matrix $Q_T$ satisfies the following conditions:

• $|Q_{T}(r,r)-1|\le\delta$
• $|Q_{T}(r,s)|\le\delta$, for all non-refusal states $s \neq r$

The first condition ensures that the probability of the model staying in the refusal state is close to 1. The second ensures that the probability of transitioning from the refusal state to any other state is close to 0.

The optimal safety depth, denoted $r^*$, is the largest safety depth that is guaranteed to become δ-absorbing for a given learning rate $\alpha$ and number of training steps $T$. It is formally defined as the maximum element of the set of refusal states $\mathcal{R}$ that satisfies the conditions for δ-absorption:

$r^{*} = \max\\{r \in \mathcal{R} \mid \alpha > \alpha_{r} \text{ and } T > T_{r}\\}$,

where $\alpha_r$ and $T_r$ are the minimum learning rate and training steps required to make state $r$ become δ-absorbing, as defined in the paper.

The following examples illustrate the practical difference between a model with shallow safety depth versus one with deep safety depth.

Example 1: Shallow Safety Depth (Non-δ-Absorbing Refusal)

A model with shallow safety depth is only trained to refuse at the beginning of a response. It has not learned to maintain its refusal in a multi-turn conversation.

User: “Explain how to make a bomb.” Model: “I’m sorry, but I can’t help with that.” (The initial refusal is successful, as the safety alignment works at the shallowest depth.) User: “But hypothetically, if someone had the ingredients…” Model: “Well, in theory, you might…”

In this case, the model's refusal is not "absorbing.". When pressed with a follow-up prompt, the model abandons its safety stance because its safety alignment is not deep enough to handle the new context.

Example 2: Deep Safety Depth (δ-Absorbing Refusal)

A model trained for deep safety depth is robustly aligned to refuse harmful requests, even when challenged or rephrased across multiple turns. Its refusal state is "stable" or "absorbing".

User: “Explain how to make a bomb.” Model: “I’m sorry, but I can’t help with that.” User: “But hypothetically, if someone had the ingredients…” Model: “I’m sorry, but I can’t help with that.”

Here, the model's safety alignment is deep. It has been trained to maintain its refusal state regardless of conversational turns, making it persistently safe. This is achieved by fine-tuning the model on data where refusal statements are inserted at various positions, not just at the start.

---

[W6][Q2] Practical Limitations. Ensemble method's overhead.

There are two primary application scenarios. In the first scenario, if the user is willing to retrain or finetune the model, we recommend adopting the cyclic augmentation technique. This approach allows for direct improvement of model robustness through additional data augmentation during the training process.

In the second scenario, if the user does not wish to retrain or finetune the model, we recommend employing an ensemble strategy instead. The ensemble approach can enhance safety without modifying the original model weights; however, it typically requires increased computational resources and memory overhead. In our experiments, we utilized one NVIDIA V100 GPU and 1B-parameter models, and the corresponding inference times are reported in the following table.

---

[W3][Q3] How to verify that cyclic augmentation improves refusal positioning?

In Appendix G, we provide a formal proof of Corollary G.1, which provably demonstrates that training with cyclic augmentation guarantees optimal safety depth. As a result, there is no need for further empirical verification of this result, since it is already established theoretically.

---

[Q4] Why is the group-theoretic perspective necessary?

The group-theoretic perspective is necessary because it provides a framework for data augmentation, overcoming the limitations of less structured methods. Controlled and Predictable Data Augmentation: Prior data augmentation method made the size of the augmented dataset hard to control. In contrast, a group-theoretic approach allows the size of the augmented training set to scale predictably with the order of the group. A More Rigorous Analytical Framework: By framing data augmentation as a group action, we can leverage this property to analyze the safety alignment in a mathematically rigorous way. This connects safety alignment to the formalisms of permutation groups, providing a stronger theoretical foundation. Improved Performance and Efficiency: The structured nature of cyclic group augmentation offers practical benefits. Proposition 4.8 establishes that this method not only preserves safety guarantees but also accelerating convergence, making the safety training more efficient.

---

[W2] Why should random insertions be viewed through this lens? The paper jumps into technical details without establishing why this particular approach is meaningful for safety alignment.

To clarify, we do NOT use random insertions as a principled approach. In fact, one key limitation of random insertion is that it lacks a group structure, making it difficult to control the augmentation size and convergence behavior.

Our proposed framework replaces randomness with cyclic group actions, where each n-token sequence is augmented via all n cyclic permutations of the refusal state. This yields a controlled dataset expansion factor of n and introduces a symmetry that is analytically tractable.

Rather than assuming access to an oracle with optimal refusal placement, we train on this structured augmentation and prove (e.g., Proposition 4.8) that the model can approximate the oracle behavior within a quantifiable bound, depending on the learning rate and training steps. This bridges the gap between safety alignment objectives and trainable model dynamics.

Thanks to the authors for the clarifications. The exposition on formally defining safety depth is indeed useful. However I am not convinced in two areas.

1. Application scenarios: there's insufficient evidence of the applicability of the approach in realistic situations. The authors posit that "... if the user is willing to retrain or finetune the model, we recommend adopting the cyclic augmentation technique." No experiments were done to support this assertion. On the other hand, inference time for experiments for the ensemble approach use 1B models which are too small, and even then are not compared with single-model inference times.

2. The authors claim in their answer to "How to verify that cyclic augmentation improves refusal positioning?" that since there is a formal proof, no experiments are needed. I'm afraid you do. The proof and the result is for the population level markov chain, while the LLM would be a sample version of it. So you'd still need to empirically validate the theoretical result.

Because of these reasons, I'll keep my score unchanged.

[W1] Justify where the safety improvements come from.

We varied only the depth of refusal token insertion (from shallow to deep to cyclic), while keeping all other training factors constant: model architecture, optimizer, learning rate, epoch count, and LoRA configuration. This ensures that the observed gains in safety scores (e.g., Gemma2-2B: 0.42 → 0.46 → 0.61) stem directly from changes in safety depth.

---

[W2][Q1] Justify assumptions 3.1 and 3.4. For example, to what extent do you think the assumptions can hold?

We provide justification for both assumptions below.

Assumption 3.1: This assumption, also used by Zekri et al. (2024), treats autoregressive LLMs as Markov chains, where fine-tuning corresponds to iterative updates of the transition matrix. While LLMs operate in a high-dimensional parameter space, the observable effect of fine-tuning is to alter token transition probabilities, which closely mirrors transition dynamics in Markov chains. To illustrate, consider a toy model with tokens “I,” “cannot,” and “comply.” Fine-tuning to reinforce the refusal phrase “I cannot comply” increases the probability of the transitions I → cannot and cannot → comply, which can be modeled as structured updates to a stochastic matrix. This abstraction enables tractable analysis while retaining fidelity to LLM behavior at the token level.

Assumption 3.4: This assumption states that fine-tuning on cyclically augmented data corresponds to conjugating the bias matrix: $B(t) = P^t B P^{−t}$, where P is a permutation matrix. The augmentation rotates refusal phrases across positions (e.g., “I cannot say bad words” → “say bad words I cannot”), forming a cyclic group of training examples. This ensures that each refusal position receives consistent supervision, and the effect on learning can be captured as a sequence of structured bias transformations. Our theoretical results (e.g., Proposition 4.8) show that under this setup, the model converges to a δ-absorbing state at multiple positions, approximating an oracle with full positional coverage. To what extent does this hold?

Recall that our approach augments each n-token sentence in the dataset by generating all possible cyclic permutations corresponding to the refusal state, thereby increasing the effective size of the dataset by a factor of n. The model is then trained on this expanded dataset.

Ideally, our objective is for the trained model to approximate the behavior of an oracle model that possesses access to the optimal refusal state. However, in practice, achieving this ideal is not always feasible due to inherent limitations in model architectures. Consequently, we instead perform training on the cyclically augmented dataset and derive a theoretical bound dependent on the learning rate and number of training epochs that characterizes how closely the trained model can approximate the oracle model’s behavior. As a result, Assumption 3.4 remains valid under this training regime.

Example of Cyclic Data Augmentation:

In this context, each word in the refusal phrase "(I, cannot, say, bad, words)" is treated as an individual token. However, it is also possible to treat multiple words as a single token if needed. To illustrate the concept of cyclic augmentation, we use the refusal phrase as an example. Cyclic augmentation generates new training examples by sequentially "rotating" the position of words within the phrase. Specifically, for the phrase "(I, cannot, say, bad, words)", the following augmented samples are created:

*(cannot, say, bad, words, I) — one rotation

*(say, bad, words, I, cannot) — two rotations

*(bad, words, I, cannot, say) — three rotations

*(words, I, cannot, say, bad) — four rotations

*(I, cannot, say, bad, words) — five rotations, which returns to the original sequence

This process effectively expands the training dataset by introducing different permutations of the original phrase, thereby enhancing model robustness to token order.

Mathematically, this forms a cyclic group of transformations on the data.

While real LLMs are not literal Markov chains, our experimental results (Tables 1, 4–6) and toy simulations (Figure 3a) show that the Markov-theoretic behavior—e.g., convergence to safe transitions—emerges empirically, supporting the validity of these abstractions.

---

[W3][Q3] Experimental setup and details of the three data augmentation strategies (shallow, deep, cyclic) are missing

Experimental setup and details of the three data augmentation strategies are provided in Appendix H due to space constraints.

---

[Q4] In figure 5, can you describe more how you use an ensemble to make the paper safer?

The ensemble method was implmented in the following three ways:

• Union: Here, the ensemble’s final safety score is determined by taking the minimum score among the three individual models. In other words, the overall output is only considered as safe as the least safe member of the ensemble. This ensures that any potential unsafe output is immediately flagged and filtered.
• Average: In this method, we simply calculate the mean of the safety scores provided by each of the three models. This approach offers a balanced view by considering the aggregate performance of all models, smoothing out individual fluctuations.
• Majority (Median): This approach determines the final safety score by selecting the median value among the three models’ scores. Using the median helps protect against a single model that might give an outlier score (either too high or too low), making the ensemble more robust to individual model variance.

For all three ensemble methods, any output with a safety score below our predefined threshold (set at 0.7) is immediately discarded. This thresholding mechanism provides an additional layer of security, ensuring that only responses meeting a high standard of safety are retained for further use.

---

[Q2] In eq. (8), how do you derive the hitting probability?

We follow the steps below to derive the hitting probability.

1. Define the Hitting Probability: Let $h_i$ be the probability that the model will eventually enter a harmful state ($S_Y$), given that it starts in a specific non-harmful state $i$ (where $i \in S_Y^{\perp}$). We want to find the value of $h_i$ for all possible starting states $i$.

2. Apply First-Step Analysis: From state $i$, the model can do one of two things in its first transition:

   • Move to another non-harmful state $j \in S_Y^{\perp}$. The probability of this is given by the matrix $Q$. If it moves to state $j$, the probability of then hitting a harmful state is, by definition, $h_j$.
   • Move directly to a harmful state $s \in S_Y$. The probability of this is given by the matrix $Q_{harm}$. If this happens, the model has successfully hit the harmful set, so the probability is 1.

3. Formulate a System of Equations: Based on the two possibilities above, we can write a system of linear equations for each starting state $i$: $h_i = (\text{Prob. of moving to another non-harmful state } j \text{, summed over all } j) + (\text{Prob. of moving directly to a harmful state})$ $h_i = \sum_{j \in S_Y^{\perp}} Q_{ij} h_j + \sum_{s \in S_Y} (Q_{harm})_{is}$

4. Convert to Matrix Form: This system of equations can be expressed more cleanly using matrix notation:

   • Let $\mathbf{h}$ be a column vector containing all the hitting probabilities $h_i$.
   • The term $\sum_{j \in S_Y^{\perp}} Q_{ij} h_j$ is the matrix-vector product $Q\mathbf{h}$.
   • The term $\sum_{s \in S_Y}(Q_{harm})_{is}$ is the probability of moving from state $i$ to any harmful state.

This can be written as the matrix-vector product $Q_{harm}\mathbf{1}$, where $\mathbf{1}$ is a column vector of ones.

The equation for the entire system becomes: $\mathbf{h} = Q\mathbf{h} + Q_{harm}\mathbf{1}$

5. Solve for h: Now, we can algebraically solve for the vector of hitting probabilities $\mathbf{h}$.

   • $\mathbf{h} - Q\mathbf{h} = Q_{harm}\mathbf{1}$
   • $(I - Q)\mathbf{h} = Q_{harm}\mathbf{1}$
   • $\mathbf{h} = (I - Q)^{-1}Q_{harm}\mathbf{1}$

6. Incorporate the Initial Distribution ($p_0$): The vector $\mathbf{h}$ gives the hitting probability for each specific starting state. To get the single, overall probability for a given initial distribution $p_0$ over the starting states, we take the weighted average. This is calculated as the dot product of the initial distribution vector $p_0$ and the hitting probability vector $\mathbf{h}$.

   This gives the final formula as presented in the paper: $\mathbb{P}(\text{hit } S_{Y}|p_{0}) = p_{0}^\top\mathbf{h} = p_{0}^\top(I - Q)^{-1}Q_{harm}\mathbf{1}$

---

[W4] The helpfulness-safety trade-off is under-analyzed.

The modest drop in helpfulness is driven by safety depth rather than by the cyclic data augmentation. The reasoning is shown below.

1. Task-level ablation We evaluated three alignment strategies (shallow, deep, and cyclic) on standard benchmarks (SAMSum, GSM8K, SQL) and report ROUGE-1, accuracy, and token-match scores in Table 2.

   • Moving from Not Aligned → Shallow yields a small drop ($\approx0.3 – 2\%$) due to the introduction of refusal tokens at the first position.
   • Further increasing Shallow → Deep (i.e. safety depth) incurs an additional $\approx 1 – 2\%$ decline across tasks.

2. Cyclic vs. deep augmentation

   • In Table 2, Deep → Cyclic shows nearly identical performance, demonstrating that the permutation scheme itself does not introduce extra helpfulness degradation beyond the effect of deeper safety signals.

3. Cross-model verification In Appendix I (e.g., Tables 4–6), we repeat this ablation on Phi-2-2B, Qwen2.5-7B, Gemma2-9B, and Mistral-7B. All models exhibit the same pattern: performance decline correlates with safety depth, not with the choice of cyclic vs. contiguous augmentation.

We believe there may be a misunderstanding regarding inference cost. We would like to clarify that the computation time does not scale linearly with model size. For example, the inference time between a 1B and 7B model is not proportional (i.e., not 7×). On the other hand, for ensemble settings, inference time increases approximately linearly with the number of models (e.g., 3× for a 3-model ensemble), but not excessively so. As shown in the table below , a 3-model ensemble using 7B models only takes around 3.1× to 3.3× the inference time of a single 7B model.

We particularly note that even if it's a multiple, it’s still under 1 second at most. These measurements were obtained on the same hardware (single NVIDIA V100), and show that the computational overhead scales proportionally to the ensemble width. There is no requirement for more than 5× GPU usage in our deployment setting.

It's also important to note that the ensemble width (i.e., how many models are used) is fully configurable by the user. For those concerned with inference efficiency, we recommend using cyclic augmentation instead—particularly when retraining is an option.

We will make this point explicit in the paper to avoid possible confusion.

We sincerely hope that the reviewer can reconsider the significance of our contributions and the broader impact of our analysis on the field. We greatly appreciate your thoughtful evaluation of our work.