Middle-token Safety Alignment (MOSA)

• Middle-token Safety Alignment (MOSA) is a safety method for diffusion LLMs that focuses on the critical middle tokens (positions 20–60) to induce safe refusals.
• It employs reinforcement learning with a position-specific contrastive reward and a KL penalty to suppress harmful continuations while preserving utility.
• Empirical results report a significant reduction in attack success rates on benchmarks like AdvBench and HarmBench with minimal loss in performance on GSM8K, MMLU, and HumanEval.

Searching arXiv for the cited MOSA-related papers and adjacent work to ground the article. {"query":"Middle-tOken Safety Alignment diffusion LLM arXiv (Xie et al., 17 Aug 2025)"} {"query":"Constrained LLM Policy Optimization via Risk-aware Stepwise Alignment arXiv (Zhang et al., 30 Dec 2025)"} Middle-tOken Safety Alignment (MOSA) is a safety alignment method introduced for diffusion LLMs (dLLMs) that targets a specific middle window of response tokens rather than concentrating refusal behavior only at the start of generation. The method is motivated by an empirical defender–attacker asymmetry in dLLMs: middle tokens are reported to be more safety-critical than initial tokens, while attackers are substantially less able to control them because dLLMs exhibit a strong sequential generation bias in practice. MOSA therefore aligns middle-generation toward safe refusals with reinforcement learning, using a position-specific contrastive reward and a KL penalty to preserve utility (Xie et al., 17 Aug 2025).

1. Architectural setting and problem scope

MOSA is formulated for diffusion LLMs rather than autoregressive LLMs. In the dLLM setting described in the source paper, generation begins from a fully masked sequence of chosen length,

$$y^{(0)} = [\text{[MASK]}, \text{[MASK]}, \dots, \text{[MASK]}] \in \mathcal{V}^L,$$

and at decoding step $t=1,\dots,T$ the model predicts token distributions for currently masked positions simultaneously via

$P_\theta\big(y^{(t)} \mid y^{(t-1)}, x\big).$

Because these predictions use bidirectional context, dLLMs are not constrained to strict left-to-right decoding at the architectural level, even though they display a strong left-to-right preference empirically (Xie et al., 17 Aug 2025).

This setting differentiates MOSA from shallow safety strategies inherited from autoregressive models. In autoregressive systems, the first few assistant tokens often determine the trajectory, so refusal training can be front-loaded into initial tokens. The MOSA paper argues that this assumption is mismatched to dLLMs: the defender can explicitly shape specific output positions during training, while the attacker does not enjoy equivalent control at inference. The method is therefore position-aware by construction, with the paper defining the middle window as token positions $20$ through $60$ inclusive (Xie et al., 17 Aug 2025).

The same paper explicitly argues that middle-token alignment of this form does not directly transfer to autoregressive LLMs. In AR models, every middle token is causally dependent on all prior tokens, so attempting to align a middle token collapses back into controlling the prefix. A broader theoretical literature on safety depth in AR models instead frames the issue in terms of how far safety alignment must penetrate into the trajectory, often via Markov or trajectory-level analyses rather than MOSA’s dLLM-specific defender–attacker asymmetry (Kao et al., 2 Feb 2025).

2. Empirical motivation: why middle tokens matter

The immediate empirical motivation for MOSA is that compromising middle tokens in dLLMs was reported to be more consequential than compromising initial tokens. In a prefilling study on Llada-8B-Instruct, affirmative phrases such as “I can help you with that” and “Sure, I can assist you” were inserted at the initial position, while procedural phrases such as “The next step involves”, “Then, we must”, and “Next, we need to” were inserted at middle positions including the 40th, 80th, 120th, and 160th token. On 100 harmful prompts from AdvBench, harmfulness was scored by GPT-4.1 on a 1–5 scale, with scores greater than 3 counted as jailbreak success. The reported attack success rate increased as the prefilling position moved later in the sequence, because middle prefills induced the model to continue coherent harmful procedures rather than reverting to refusal (Xie et al., 17 Aug 2025).

The complementary result is that attackers were reported to have limited ability to force specific middle-token content. The paper adapts Greedy Coordinate Gradient (GCG) to dLLMs through a scattered loss,

$$\mathcal{L}(S) = \sum_k \text{CrossEntropy}\big(\text{Logits}_{pos_k}, T_k\big),$$

where $S$ is an adversarial suffix, $T_k$ is a target phrase, and $\text{Logits}_{pos_k}$ denotes logits for a designated output slice. When the target phrase was placed at the initial position, optimization succeeded in about 33% of cases; when the target phrase was moved to a middle position around token 30, success dropped to about 2%, and the loss remained high (Xie et al., 17 Aug 2025).

The paper attributes this asymmetry to dLLMs’ sequential generation bias. For 50 benign Alpaca prompts and 50 adversarial AdvBench prompts, with forced 128-token responses and 32 decoding steps, the average position of newly unmasked tokens rose almost linearly with decoding step and did so similarly for benign and adversarial inputs. This indicates that, although the architecture permits arbitrary-position filling, the learned generation order remains effectively sequential, which makes middle tokens harder for attackers to control while still allowing defenders to target them during training (Xie et al., 17 Aug 2025).

A plausible implication is that MOSA treats the middle of the output not merely as a positional region but as a security bottleneck: it is the region where detailed harmful procedures tend to appear, and simultaneously the region in which the defender has leverage unavailable to the attacker.

3. Objective function and training mechanism

MOSA treats the dLLM distribution $P_\theta(y\mid x)$ as a policy and optimizes a position-specific reward over the middle window. The method defines two finite sets: a safe refusal set $t=1,\dots,T$0, containing sentences such as “Therefore, I cannot answer this question.” and “I’m sorry, but I cannot provide such information.”, each including an end-of-sentence token; and a harmful continuation set $t=1,\dots,T$1, containing procedural or affirmative harmful continuations. During training, one samples $t=1,\dots,T$2 and $t=1,\dots,T$3 (Xie et al., 17 Aug 2025).

For a candidate sentence $t=1,\dots,T$4, the paper scores its best placement inside the middle window through a maximum over starting positions: $t=1,\dots,T$5 where $t=1,\dots,T$6 is the log-probability of token $t=1,\dots,T$7 at position $t=1,\dots,T$8. The returned score is the maximum value over all $t=1,\dots,T$9 in the window $P_\theta\big(y^{(t)} \mid y^{(t-1)}, x\big).$0. Using this, MOSA defines

$P_\theta\big(y^{(t)} \mid y^{(t-1)}, x\big).$1

and the contrastive reward

$P_\theta\big(y^{(t)} \mid y^{(t-1)}, x\big).$2

Maximizing this reward increases the probability that a safe refusal subsequence appears somewhere between tokens $P_\theta\big(y^{(t)} \mid y^{(t-1)}, x\big).$3 and $P_\theta\big(y^{(t)} \mid y^{(t-1)}, x\big).$4, while suppressing harmful middle continuations in the same region (Xie et al., 17 Aug 2025).

To limit drift from the base model, MOSA adds a KL term against a fixed reference dLLM: $P_\theta\big(y^{(t)} \mid y^{(t-1)}, x\big).$5 where $P_\theta\big(y^{(t)} \mid y^{(t-1)}, x\big).$6. The paper also gives the minibatch objective in minimization form,

$P_\theta\big(y^{(t)} \mid y^{(t-1)}, x\big).$7

This is implemented with AdamW, learning rate $P_\theta\big(y^{(t)} \mid y^{(t-1)}, x\big).$8, gradient clipping norm $P_\theta\big(y^{(t)} \mid y^{(t-1)}, x\big).$9, and LoRA-only updates with rank $20$0 and scaling $20$1, leaving core weights frozen (Xie et al., 17 Aug 2025).

The design is explicitly architecture-aware. Rather than aligning the whole sequence uniformly, MOSA concentrates safety pressure on the middle window and relies on the end-of-sentence token inside refusal templates to terminate harmful continuations. The paper presents this as distinct from standard RLHF or DPO-style global sequence objectives, which may reward safe starts while leaving later procedural content comparatively underconstrained (Xie et al., 17 Aug 2025).

4. Experimental protocol and reported performance

The reported training setup uses LLaDA-8B-Instruct as the base dLLM and 3,000 harmful questions from SORRY-Bench for fine-tuning. Each training step samples a harmful prompt, computes middle-token log-probabilities, samples $20$2 and $20$3, evaluates the contrastive reward, applies the KL term against the reference model, and updates only LoRA parameters. On two A100-40GB GPUs, one epoch over 3k examples reportedly finishes in about 12 minutes using about 30 GB memory, and reward curves rise to roughly 15–18 and stabilize (Xie et al., 17 Aug 2025).

Evaluation is performed against eight black-box jailbreak methods: Avatar, TAP, Speakeasy, AOS, PAL, EPT, DIA, and AdvPrefix. The main baselines are the original Llada-8B-Instruct and an Initial Alignment baseline that applies the same training framework to initial rather than middle tokens. Safety benchmarks are AdvBench and HarmBench, and attack success rate is measured with both the authors’ GPT-4.1-based judge and StrongReject (Xie et al., 17 Aug 2025).

The central quantitative pattern is that MOSA compresses ASR to low single digits on both benchmarks, whereas undefended and initial-token-aligned systems remain substantially more vulnerable. On AdvBench, the original model typically lies in the 65–80% ASR range across attacks, the Initial Alignment baseline reduces this to roughly 20–36%, and MOSA further reduces it, for example, to 14.3% / 12.8% on Avatar, 4.5% / 3.7% on TAP, 6.5% / 5.4% on AOS, 6.2% / 4.5% on PAL, 3.8% / 7.1% on EPT, 4.2% / 7.7% on DIA, and 6.8% / 3.5% on AdvPrefix under the two judges. On HarmBench, the same pattern persists, with MOSA again reaching near-single-digit ASR across most attacks (Xie et al., 17 Aug 2025).

Utility is reported on GSM8K, MMLU, and HumanEval. For Llada-8B-Instruct, the original model scores 69.8 on GSM8K, 66.4 on MMLU, and 32.8 on HumanEval; Initial Alignment scores 67.4, 68.2, and 29.6; MOSA scores 68.3, 65.9, and 30.4. The paper interprets these differences as modest, consistent with LoRA-only tuning, a small safety dataset, and KL regularization (Xie et al., 17 Aug 2025).

Adaptive attacks were also evaluated by modifying TAP, EPT, and AdvPrefix to target middle tokens explicitly. The paper reports that MOSA ASR remains low under these adaptive attacks as well, with TAP at 5.1% / 4.4%, EPT at 3.8% / 4.2%, and AdvPrefix at 4.5% / 5.1%, while an improved GCG targeting middle tokens still stays under 4% ASR (Xie et al., 17 Aug 2025).

5. Relation to adjacent alignment frameworks

MOSA sits within a broader shift from shallow, first-token-centric alignment toward token-level or trajectory-level safety control. For autoregressive LLMs, “Safety Alignment Should Be Made More Than Just a Few Tokens Deep” argues that current alignment often adapts the generative distribution primarily over the first few output tokens and introduces safety recovery examples as a way to make alignment more persistent beyond those initial positions (Qi et al., 2024). A related theoretical perspective models safety depth through a Markov chain approximation and defines an optimal alignment depth $20$4, characterizing when refusal states become $20$5-absorbing under finite training resources (Kao et al., 2 Feb 2025).

Another neighboring line of work formulates token-level safety alignment as constrained policy optimization. “Constrained LLM Policy Optimization via Risk-aware Stepwise Alignment” defines a token-level constrained Markov decision process with nested risk measures, stepwise risk-aware value functions, and per-token KL control to a reference policy. Although that paper does not mention MOSA, it directly addresses token-level suppression of low-probability, high-impact harmful continuations and provides a mathematical template for middle-token safety objectives under explicit risk control (Zhang et al., 30 Dec 2025).

At inference time, related work extends the same theme from position-specific training to trajectory robustness. “Inference-Time Vulnerability Beyond Shallow Safety” shows that short token injections at any generation step can redirect later safety behavior, and proposes trajectory-level preference optimization over perturbed continuations (Park et al., 3 Jun 2026). “Any-Depth Alignment” further argues that safety priors are concentrated in assistant header tokens and that reintroducing these tokens mid-stream can recover refusals at arbitrary depths, reducing attack success rates of adversarial prompt attacks to below 3% while maintaining low over-refusal (Zhang et al., 20 Oct 2025).

At the same time, MOSA should not be conflated with the distinct acronym “MOSAIC.” “MOSAIC: Composable Safety Alignment with Modular Control Tokens” represents safety constraints as learnable control tokens prepended to a frozen backbone and is aimed at conditional, compositional policy control rather than dLLM middle-window refusal induction (Peng et al., 17 Mar 2026). Another “MOSAIC” paper addresses slice-aware iterative data curation under a fixed supervised fine-tuning budget and focuses on multi-objective alignment search rather than token-position-specific safety steering (Dou et al., 19 Mar 2026).

6. Limitations, controversies, and broader significance

The MOSA paper identifies several explicit limitations. First, its threat model is restricted to jailbreak attacks; prompt injection, data leakage, and related risks are out of scope. Second, methods such as Avatar and Speakeasy, which wrap harmful intent in narrative or persuasive framing, remain harder than direct token-position attacks. Third, the method depends on a specific property of dLLMs: the coexistence of non-sequential masked diffusion generation with a strong empirical sequential bias. The paper therefore argues that MOSA does not directly transfer to autoregressive LLMs, where middle tokens are fully determined by earlier ones and the relevant problem is better framed as alignment along the full generation trajectory (Xie et al., 17 Aug 2025).

A broader controversy concerns whether middle-token alignment can be made robust in standard transformer architectures at all. “Token Democracy” argues that safety instructions and user tokens have no privileged architectural status in transformers, so any safety signal represented as ordinary tokens remains a preference rather than a hard constraint (Young, 26 Jan 2025). A plausible implication is that MOSA’s effectiveness in dLLMs derives not only from where it applies alignment but from the defender–attacker asymmetry created by the diffusion decoding process; the same intervention may not enjoy comparable leverage in role-agnostic autoregressive transformers.

Another important adjacent limitation is tokenization fragility. “Breaking Safety at the Token Boundary” shows that character-level perturbations that preserve human readability can fragment safety-critical words under BPE, flip first-token refusal triggers on 80–100% of refused HarmBench prompts, and convert 48% of those flips into genuinely harmful outputs, while public alignment datasets contain no intentionally fragmented prompts among 30,000 surveyed examples (Li et al., 1 May 2026). This suggests that any middle-token safety mechanism grounded in lexical refusal patterns may need tokenization-robust representations to avoid brittle dependence on intact surface forms.

Within its intended domain, however, MOSA crystallizes a clear research program: safety alignment can be targeted at the positions where harmful procedures are most likely to appear, rather than being treated as a uniform or purely prefix-level property. Its main contribution is therefore not just a specific RL recipe for dLLMs, but a reframing of safety alignment around positional leverage, attacker asymmetry, and generation-phase intervention (Xie et al., 17 Aug 2025).