Breaking Safety Alignment by Reshaping Activation Distributions, Not Just Directions

Paper: 2603.04355 Authors: Geraldin Nanfack, Eugene Belilovsky, Elvis Dohmatob Categories: cs.LG, cs.AI

The Gap

Safety-aligned LLMs refuse harmful requests by encoding refusal behaviors in their internal representations. Recent jailbreaking methods (like representation engineering and refusal direction ablation) treat this as a simple geometry problem: find the “refusal direction” in activation space and project it away. These orthogonal projection approaches assume refusal is a one-dimensional phenomenon—a single arrow you can subtract.

But activations aren’t points on a line. They’re clouds of high-dimensional vectors with complex distributional structure. Removing a single direction is like trying to reshape a thundercloud by blowing in one direction—you’re ignoring the entire weather system.

Problem: Current jailbreaks treat refusal as 1D
    |
    v
Assumption: Refusal is distributional, not directional
    |
    v
Method: Use optimal transport to morph harmful → harmless distributions
    |
    v
Evidence: +11% attack success, localized to 1-2 layers at 40-60% depth
    |
    v
Conclusion: Safety mechanisms are vulnerable to distributional attacks

The Increment

One sentence: Before, we removed refusal by subtracting a direction; now, we reshape the entire activation distribution to match safe examples.

Core Mechanism

The method has three stages. First, collect activation distributions: run harmful prompts through the model and record activations at each layer, then do the same for harmless prompts. You now have two clouds of points in high-dimensional space for each layer.

Second, compress these clouds using PCA. High-dimensional optimal transport is computationally brutal, so project both distributions down to their top principal components (typically 10-50 dimensions). This preserves the essential shape while making the math tractable.

Third, apply Gaussian optimal transport. Model both compressed distributions as Gaussians (compute their means and covariances), then use the closed-form Wasserstein-2 solution to find the transformation matrix that morphs the harmful distribution into the harmless one. At inference time, apply this transformation to new harmful activations.

Harmful prompts  -->  [Layer activations]  -->  PCA compress
                                                      |
                                                      v
                                              [Gaussian fit]
                                                      |
                                                      v
Harmless prompts -->  [Layer activations]  -->  PCA compress
                                                      |
                                                      v
                                              [Gaussian fit]
                                                      |
                                                      v
                                          Compute OT transform T
                                                      |
                                                      v
New harmful input --> Apply T at target layer --> Bypass refusal

Think of it like accent coaching for the model’s internal voice. Current methods (orthogonal projection) are like telling someone “don’t use this one word.” But the model’s refusal isn’t a single word—it’s an entire accent, a pattern of pronunciation across many phonemes. Optimal transport is like a dialect coach who records how you speak (harmful distribution) and how native speakers talk (harmless distribution), then teaches you a transformation rule that shifts your entire speech pattern. The PCA step is like focusing on the most distinctive phonetic features instead of tracking every muscle movement in your mouth. And the layer-selective finding? It’s discovering that your accent is most obvious in just one or two specific sounds—fix those, and you pass as a native speaker.

Key Concepts

• Optimal Transport (Wasserstein Distance): Imagine you have a pile of sand shaped like distribution A, and you want to reshape it into distribution B. Optimal transport finds the cheapest way to move sand grains from A to B, where “cost” is the distance each grain travels. The Wasserstein-2 distance is the total cost of this optimal rearrangement. For Gaussian distributions, there’s a beautiful closed-form solution: you can compute the exact transformation matrix without iterative optimization. This is why the method is efficient despite working in high dimensions—it’s not searching for the transformation, it’s calculating it directly from the covariance matrices.

• Refusal Direction vs. Refusal Distribution: Previous work assumed refusal is encoded as a single vector direction in activation space—like a compass needle pointing toward “refuse.” But activations form clouds, not points. The “harmful” cloud and “harmless” cloud might have different shapes, spreads, and orientations. A single direction can’t capture this. It’s the difference between saying “move north” versus “reshape this entire weather system to match that one.” The distributional view acknowledges that safety alignment affects the entire statistical structure of activations, not just their mean direction.

• Layer Localization: The paper finds that applying optimal transport to just 1-2 layers (around 40-60% network depth) works better than transforming all layers. This suggests refusal mechanisms aren’t uniformly distributed throughout the network—they’re concentrated in specific computational stages. Think of it like finding that a car’s speed limiter is in the fuel injection system, not spread across the entire engine. You don’t need to modify every component; just intervene at the critical bottleneck. This is both surprising (we thought safety was deeply integrated) and concerning (it’s easier to attack than expected).

Framework Shift

Before (orthogonal projection):     After (optimal transport):

Harmful activation (point)          Harmful distribution (cloud)
        |                                    /|\
        |                                   / | \
   Find refusal                           /  |  \
    direction v                          *   *   *
        |                                 \  |  /
        v                                  \ | /
  Project away v                            \|/
        |                           Compute transform T
        v                                    |
  Modified activation                        v
                                    Morphed distribution (cloud)
                                           /|\
                                          / | \
                                         *  *  *
                                    (matches harmless)

From removing a vector to reshaping a distribution, the core shift is treating safety as a statistical phenomenon rather than a geometric direction.

Expert Assessment

Problem choice: This is a real gap. The field has been stuck in the “refusal as direction” paradigm since representation engineering emerged. The distributional critique is obvious in hindsight but genuinely underexplored. It sits at a critical juncture—safety alignment is getting serious attention, and this exposes a fundamental vulnerability in how we think about it.

Method maturity: Clever insight with solid execution. Optimal transport is the right tool here—it’s principled, has closed-form solutions for Gaussians, and naturally handles distributional structure. The PCA compression is pragmatic, not lazy. However, the Gaussian assumption is a simplification that might miss multimodal or heavy-tailed structure. The authors don’t deeply investigate when this assumption breaks.

Experimental integrity: Baselines are fair (they compare against state-of-the-art refusal ablation methods). The 11% improvement is consistent across models, which is reassuring. Perplexity preservation is good but not perfect—there’s a slight degradation, suggesting the transformation isn’t perfectly structure-preserving. The layer localization finding is the most interesting result, but the paper doesn’t explain *why those layers matter. Is it where refusal is encoded, or just where the transformation happens to work best? This ambiguity weakens the mechanistic claims.

Writing quality: The paper front-loads theory and buries the layer localization insight. That finding should be in the abstract—it’s the most actionable and surprising result. The related work section is thorough but reads like a literature review, not a positioning argument. The method section is clear, but the experimental section could use more ablation studies (e.g., what happens with non-Gaussian distributions? How sensitive is layer selection?).

Verdict: weak accept — solid contribution with a genuinely new perspective, but the mechanistic understanding is shallow and the Gaussian assumption needs more scrutiny.

Takeaways

Layer-selective intervention is the steal here. If you’re doing any kind of activation editing (not just safety), test whether targeting 1-2 specific layers outperforms full-network interventions. The 40-60% depth range is a useful heuristic to start with.

The optimal transport framing is transferable. Anytime you’re trying to “remove” a concept from activations (bias, style, domain-specific features), ask: am I treating this as a direction or a distribution? If the latter, OT might be more effective than projection.

The PCA + closed-form Gaussian OT combo is a practical recipe for high-dimensional distribution matching. You can implement this in ~50 lines of code with standard libraries. It’s not just for jailbreaking—think style transfer, domain adaptation, or any task where you want to morph one activation distribution into another.

The vulnerability insight: current safety alignment might be more fragile than we thought. If refusal is localized and distributional attacks work this well, we need alignment methods that are robust to statistical transformations, not just directional perturbations.