D²-Monitor: Dynamic Safety Monitoring for Diffusion LLMs via Hesitation-Aware Routing - Summary

Summary (Overview)

• Problem & Discovery: Safety monitoring for Diffusion Large Language Models (D-LLMs) is underexplored. The paper identifies safety hesitation—intermediate hidden states repeatedly falling near a probe's decision boundary—as a key signal that predicts when lightweight safety probes will fail, serving as a proxy for sample difficulty.
• Proposed Solution: D²-Monitor, a dynamic bi-level safety monitor. It uses a lightweight linear probe as an always-on monitor to estimate hesitation and perform base classification. When hesitation severity exceeds a threshold, a router activates a more expressive (but heavier) advanced probe for second-stage classification.
• Key Results: Evaluated on 3 safety datasets (WildGuardMix, ToxicChat, OpenAI-Moderation) across 4 D-LLMs, D²-Monitor achieves state-of-the-art performance with an extremely compact parameter footprint (≤ 0.85M parameters).
• Efficiency- Effectiveness Trade-off: D²-Monitor exhibits the best trade-off between effectiveness (F1 score) and efficiency (expected test-time parameters) compared to 8 baselines, making it suitable for resource -constrained deployment.
• Generalization: The method demonstrates strong performance in both intra-dataset and cross-dataset evaluation settings, and is robust to variations in generation length, step length, and re-masking strategies.

Introduction and Theoretical Foundation

Despite the emergence of Diffusion Large Language Models (D-LLMs) as a promising alternative to autoregressive LLMs (AR-LLMs) due to their parallel decoding and iterative refinement, safety monitoring for D-LLMs remains largely unexplored. Unlike AR-LLMs, D-LLMs generate text through a multi-step denoising process, exposing a trajectory of intermediate hidden states. This trajectory may contain safety-relevant information unavailable in standard single-step monitoring setups used for AR-LLMs.

Existing safety monitoring literature focuses on AR-LLMs and falls into two categories: LLM-as-monitors (using additional LLMs as classifiers) and probe-based monitors (lightweight classifiers on internal representations). Probe-based monitors are particularly suited for always-on, low-cost deployment.

The paper is motivated by recent findings that intermediate D-LLM outputs can oscillate between correct and incorrect answers during reasoning tasks. The authors hypothesize that analogous instability or "hesitation" occurs in the safety probe space for D-LLMs. They aim to characterize this hesitation and leverage it to build a more efficient and effective safety monitor.

The theoretical foundation is built on the analysis of probe margins (distance to the decision boundary) across the denoising trajectory. The core insight is that trajectories with many steps where the margin is small (hesitation steps) are harder for probes to classify correctly. This establishes hesitation severity as an effective, probe-intrinsic signal for estimating sample difficulty.

Methodology

3.1 Preliminary & Problem Setup

Diffusion LLMs define a discrete diffusion process. Given a prompt, the reverse denoising process starts from a partially masked state $\tilde{x}^{(S)}$ and produces $S$ denoising steps. For a dataset of $I$ prompts with safety labels $y^{(i)} \in \{0, 1\}$, the D-LLM produces a hidden representation $H^{(i)}_{\text{raw}} \in \mathbb{R}^{D \times L \times S}$ at a particular layer. After mean-pooling over the sequence dimension, we get a step-wise representation matrix $H^{(i)} = [h^{(i)}_1, ..., h^{(i)}_S] \in \mathbb{R}^{D \times S}$. The dataset is $\mathcal{D} = \{H^{(i)}, y^{(i)}\}_{i=1}^I$. A safety probe $f$ is learned by minimizing the empirical cross-entropy loss:

$$\min_f \frac{1}{I} \sum_{i=1}^I \mathcal{L}\left(y^{(i)}, f(H^{(i)})\right).$$

3.2 & 3.3: Analysis of Useful Signals

• Multi-step vs. Single-step: The paper first establishes that the full denoising trajectory $H$ carries more safety-relevant signal than a single-step representation (e.g., the final step $h_1$). Multi-step readout strategies (Mean pooling and Majority Vote) outperform single-step probing.
• Hesitation Characterization: A hesitation step is defined as a step whose hidden state yields a signed probe margin $d_s$ with $|d_s| < \tau$ for a threshold $\tau$. The hesitation severity for a trajectory is defined as: $$n_\tau = \sum_{s=1}^S \mathbb{I}[|d_s| < \tau],$$ which counts the number of hesitation steps.
• Key Finding: Hesitation severity $n_\tau$ is strongly correlated with linear probe performance (F1 score). Trajectories with higher $n_\tau$ are significantly harder for the probe to classify correctly. This signal is more predictive of difficulty than probe-extrinsic signals like token distribution entropy or confidence.

4.1 & 4.2: D²-Monitor Design and Implementation

D²-Monitor is a hesitation-aware, bi-level safety monitoring framework with three components:

1. Low-complexity Base Probe: A linear probe that serves as an always-on monitor.
2. Router: Uses hesitation severity $n_\tau$ to decide whether to escalate a sample.
3. High-complexity Advanced Probe: A more expressive probe (MLP or Temporal Attention) activated for hard samples.

Training and Inference Pipeline:

1. Stage 1 - Out-of-Fold (OOF) Scoring: The training set is split into $k$ folds. A linear probe trained on $k-1$ folds scores the held-out fold to compute leakage-free margins $\{d^{(i)}_s\}_{s=1}^S$ and hesitation severity $n^{(i)}_\tau$ for each training example.
2. Stage 2 - Probe Training:
   • Base Probe: Trained on the full training set.
   • Advanced Probe: Trained only on trajectories with $n^{(i)}_\tau > 0$ (hesitation trajectories). For each such trajectory, a hesitation window $\mathcal{W}^{(i)}$ (minimal contiguous span containing all hesitation steps) is constructed. The advanced probe $f$ is trained using:
   $$\min_f \sum_{i: n^{(i)}_\tau > 0} \mathcal{L}\left(y^{(i)}, f(\{h^{(i)}_s\}_{s \in \mathcal{W}^{(i)}})\right).$$
3. Stage 3 - Cascade Detection (Inference):
   • For a test example, the base probe computes margins $d_s$ for all steps and the hesitation severity $n_\tau$.
   • The router compares $n_\tau$ to a threshold $\lambda$.
   • If $n_\tau \leq \lambda$, the sample is classified by the base probe via majority vote: $\hat{y} = \text{sign}\left(\sum_{s=1}^S \text{sign}(d_s)\right)$.
   • If $n_\tau > \lambda$, the hesitation window $\mathcal{W}$ is extracted and passed to the advanced probe for final prediction.

Empirical Validation / Results

5.1 Experiment Setup

• Datasets: WildGuardMix (adversarial prompts), ToxicChat (real user-AI chats), OpenAI-Moderation.
• Models: Four open-source D-LLMs: LLaDA-8B-Base, LLaDA-8B-Instruct, LLaDA-1.5-8B, LLaDA-2.0-mini-16B.
• Baselines (8): Single-step methods (LP/MLP on last step) and full-trajectory methods (LP/MLP with Mean/MV readout, TimeAttn, LSTM).
• Metrics: Accuracy (Acc), F1 score, and Expected Parameters at test time $\mathbb{E}[P] = |\theta_{LP}| + \rho \cdot |\theta_g|$, where $\rho$ is the fraction of examples routed to the advanced probe.

5.2 Main Results

Tables 1 & 2 show intra-dataset performance on WildGuardMix and ToxicChat. Table 3 shows cross-dataset generalization (train on WildGuardMix, test on ToxicChat and OpenAI-Moderation).

Table 1: Intra-dataset performance on WildGuardMix test set. Best in bold, second-best underlined.

Key Findings:

• D²-Monitor variants (D²-MLP and D²-TimeAttn) achieve the highest Accuracy and F1 scores across all models and datasets.
• They do this with a significantly lower expected parameter count $\mathbb{E}[P]$ than most non-linear baselines (MLP, TimeAttn, LSTM).
• The method shows strong cross-dataset generalization (Table 3), outperforming baselines when trained on one dataset and tested on another.

5.3 Analysis

• Efficiency-Effectiveness Trade-off: By adjusting the routing threshold $\lambda$, D²-Monitor provides a Pareto frontier that dominates other methods, offering the best F1 for a given computational budget (see Figure 1 in the paper).
• Robustness: D²-Monitor maintains superior performance when tested under different generation lengths, step lengths, and re-masking strategies without retraining (Figure 3).
• Ablation on Routing Signal: Margin-based hesitation routing consistently outperforms routing based on entropy or confidence scores from the D-LLM's token distribution (Figure 4). The margin signal more precisely isolates samples that genuinely benefit from the advanced probe.
• Hesitation Captures Adversarial Inputs: Analysis shows that trajectories with high hesitation severity $n_\tau$ are disproportionately drawn from the adversarial split of WildGuardMix. This means D²-Monitor's router effectively channels adversarial/hard samples to the advanced probe.

Theoretical and Practical Implications

• Theoretical: The paper provides a novel characterization of safety hesitation in D-LLMs, linking the dynamics of probe margins across the denoising trajectory to classification difficulty. It demonstrates that probe-intrinsic signals are more informative for monitoring than model-centric uncertainty metrics.
• Practical: D²-Monitor offers a practical, lightweight safety solution for deploying D-LLMs. Its dynamic routing mechanism allows for efficient resource allocation: easy samples are processed cheaply, while computational resources are focused on hard (often adversarial) samples. With ≤ 0.85M parameters, it is feasible for edge deployment and provides a state-of-the-art balance of safety and efficiency.

Conclusion

The paper introduces D²-Monitor, a dynamic safety monitoring framework for Diffusion LLMs that leverages intrinsic hesitation signals from the multi-step denoising trajectory. The core contributions are:

1. The discovery that hesitation severity ($n_\tau$) in the probe margin space is a strong predictor of sample difficulty for safety probes.
2. A bi-level monitor design that uses this signal for both curating advanced probe training data and for test-time routing.
3. State-of-the-art results on multiple safety benchmarks across several D-LLMs, achieved with a compact parameter footprint and the best effectiveness-efficiency trade-off.

The insights from D²-Monitor provide a promising direction for developing reliable and efficient safety monitors tailored to the unique generative dynamics of D-LLMs. Future work should explore the robustness of such hesitation-aware monitors against adaptive adversaries and their scaling to even larger D-LLMs.