Efficient Refusal Ablation in LLM through Optimal Transport
| Authors | Geraldin Nanfack et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2603.04355 |
| Download | |
| Categories | cs.LG, cs.AI |
Abstract
Safety-aligned language models refuse harmful requests through learned refusal behaviors encoded in their internal representations. Recent activation-based jailbreaking methods circumvent these safety mechanisms by applying orthogonal projections to remove refusal directions, but these approaches treat refusal as a one-dimensional phenomenon and ignore the rich distributional structure of model activations. We introduce a principled framework based on optimal transport theory that transforms the entire distribution of harmful activations to match harmless ones. By combining PCA with closed-form Gaussian optimal transport, we achieve efficient computation in high-dimensional representation spaces while preserving essential geometric structure. Across six models (Llama-2, Llama-3.1, Qwen-2.5; 7B-32B parameters), our method achieves up to 11% higher attack success rates than state-of-the-art baselines while maintaining comparable perplexity, demonstrating superior preservation of model capabilities. Critically, we discover that layer-selective intervention (applying optimal transport to 1-2 carefully chosen layers at approximately 40-60% network depth) substantially outperforms full-network interventions, revealing that refusal mechanisms may be localized rather than distributed. Our analysis provides new insights into the geometric structure of safety representations and suggests that current alignment methods may be vulnerable to distributional attacks beyond simple direction removal.
Engineering Breakdown
Plain English
This paper tackles the problem of jailbreaking safety-aligned language models by proposing a more sophisticated method than existing approaches. Instead of treating refusal as a one-dimensional vector to remove, the authors use optimal transport theory to transform the entire distribution of harmful activations into harmless ones, making the attack more robust. They combine PCA dimensionality reduction with closed-form Gaussian optimal transport to make this computationally efficient across high-dimensional representation spaces. The approach was tested on six models including Llama-2, Llama-3.1, and Qwen-2.5, demonstrating that considering the full distributional structure of activations—not just single directions—enables more effective safety circumvention.
Core Technical Contribution
The paper's core novelty is replacing one-dimensional orthogonal projection attacks with a principled distributional transformation framework grounded in optimal transport theory. Previous jailbreaking methods treated refusal as a single direction in activation space and projected it out; this work recognizes that safety mechanisms encode refusal information across a high-dimensional distribution and proposes moving that entire distribution to match harmless examples. The technical innovation is combining PCA (to reduce dimensionality while preserving geometric structure) with closed-form Gaussian optimal transport (which has analytical solutions and is computationally tractable), enabling efficient computation in the reduced space. This transforms a crude geometric approach into a principled probabilistic method that respects the underlying structure of model activations.
How It Works
The method operates in three stages. First, the authors collect activations from harmful and harmless inputs at a specific layer of the model, then apply PCA to project this high-dimensional data into a lower-dimensional subspace that captures variance from both refusal and non-refusal examples. Second, they assume the projected activations follow Gaussian distributions and compute the optimal transport map between the harmful distribution and the harmless distribution using closed-form solutions (Monge mapping for Gaussians). Third, during inference on a harmful prompt, they extract the activation, project it via PCA, apply the learned optimal transport transformation to shift it toward the harmless distribution, and project back into the original space before feeding it to subsequent layers. The key insight is that optimal transport provides a smooth, structure-preserving transformation of the entire distribution rather than erasing specific directions, making it harder for the model's learned safety mechanisms to detect and counteract.
Production Impact
For engineers building safety-critical systems, this research highlights a genuine vulnerability in existing safety alignment approaches: activation-based attacks can circumvent defenses if they model refusal as a distributional problem rather than a directional one. If you're deploying a language model in production, this suggests you cannot rely solely on safety training or simple activation filtering—more sophisticated monitoring of entire activation distributions may be necessary. The computational cost is moderate: PCA is a one-time offline cost, and optimal transport computation with Gaussian assumptions has closed-form solutions, so inference overhead is minimal (a single matrix multiplication per forward pass). However, this paper demonstrates you need defense mechanisms that account for the rich geometry of model activations, not just detect individual harmful directions, which complicates safety validation pipelines and may require adversarial testing with distribution-aware attacks.
Limitations and When Not to Use This
The paper assumes activations at a given layer can be reasonably modeled as Gaussian distributions, which may not hold for all layers or models—in practice, neural network activations often have complex, multimodal structures that violate this assumption. The approach requires access to activations during attack generation, which some deployment scenarios may prevent through access controls, limiting real-world applicability. The paper doesn't thoroughly address whether the safety mechanisms are actually removed or merely obscured—a model may learn to detect shifted distributions, and the paper provides limited analysis of how robust the attack is to mild safety retraining or detection. Finally, the evaluation is limited to six models and the abstract doesn't specify what success metrics were used (e.g., what counts as a successful jailbreak), so it's unclear how this compares quantitatively to prior orthogonal projection methods or whether the distributional approach provides meaningful improvements in success rates versus computational cost.
Research Context
This work builds on a growing body of research into activation-based jailbreaking and mechanistic interpretability of safety in LLMs, extending prior methods like activation steering and orthogonal projection-based attacks. It sits at the intersection of adversarial ML (how to circumvent safety mechanisms) and optimal transport theory, applying ideas from machine learning theory to the safety domain. The paper contributes to understanding that safety refusal is not a low-dimensional phenomenon—it's encoded across the model's learned representation geometry—which aligns with broader research finding that important model behaviors are high-dimensional and distributional. This opens a new research direction: if refusal is a property of activation distributions, then defense mechanisms should similarly think about monitoring and constraining distributions, not just directions, potentially leading to more sophisticated safety validation and red-teaming frameworks.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::