📊 statistics

Signed-Permutation Coordinate Transport for RMSNorm Transformers

This paper identifies that RMSNorm transformers possess a signed-permutation gauge ( $B_d$ ) rather than a simple permutation gauge ( $S_d$ ), and introduces a sign-marginalized Hungarian matching method to enable robust coordinate transport across checkpoints, thereby significantly improving the transferability of interpretability tools, steering vectors, and optimizer states while resolving symmetry-induced failures in existing alignment approaches.

Original authors: John Sweeney

Published 2026-07-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: John Sweeney

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have two identical houses built from the same blueprint. In one house, every room is labeled "Kitchen," "Bedroom," and "Bathroom." In the other house, the rooms are in the exact same spots, but someone has shuffled the labels.

If you want to move a specific piece of furniture (like a "steering vector" that makes the house say "no" to bad requests) from House A to House B, you need to know which label in House B corresponds to the "Kitchen" in House A. If you guess wrong, you might put the "No" button in the "Bathroom," and the house won't work right.

This paper is about solving a very specific, tricky version of this "shuffled labels" problem for modern AI models.

The Core Problem: The "Sign" Mystery

For a long time, researchers thought the only way these AI houses could be shuffled was by permuting (swapping) the rooms. If Room 1 became Room 5, you just swapped the labels.

However, the author discovered that for the most popular modern AI models (called RMSNorm models, used by giants like Llama and Qwen), there is a second, hidden way the rooms can be scrambled: flipping the signs.

Think of a room's label not just as a name, but as an arrow pointing in a direction.

Permutation: Swapping the arrow from pointing North to pointing East.
Sign Flip: Keeping the arrow pointing North, but flipping it so it points South.

In these modern models, every single room can independently flip its arrow (North $\to$ South) without breaking the house. This creates a massive amount of confusion. If you try to move a tool from one model to another using only the old "swapping" rules, you might accidentally flip half the arrows.

The "Broken Compass" Analogy

The paper argues that current tools for aligning AI models are like a compass that only knows "North" and "East," but is blind to "South."

The Old Way (Permutation Only): You try to match the rooms by looking at how they behave. If Room A in Model 1 is very similar to Room B in Model 2, you match them. But if Room A actually behaves like the opposite of Room B (because of a sign flip), the old compass thinks they are totally different and refuses to match them.
The Result: When researchers tried to move "steering tools" (like a button to make the AI refuse harmful requests) using the old method, the tools often broke completely. The "Refusal" button was flipped to "Accept," or the tool just stopped working.

The Solution: The "Sign-Marginalized" Map

The author introduces a new method called Signed-Permutation Transport.

Imagine you are trying to match the rooms again, but this time you have a special magnifying glass. Instead of just asking, "Is Room A like Room B?" you ask, "Is Room A like Room B, or is it the exact opposite of Room B?"

Look at the Magnitude: First, you check the strength of the connection, ignoring whether it's positive or negative. This finds the correct room, even if the arrow is flipped.
Check the Sign: Once you've found the matching room, you check the arrow's direction. If it's flipped, you flip it back before moving your tool.

The paper proves mathematically that if you ignore the sign flips, you will fail about 50% of the time (because half the arrows might be flipped). By accounting for the flips, you can recover the correct alignment almost 100% of the time.

Real-World Tests in the Paper

The author didn't just do math; they tested this on real AI models:

The "Refusal" Test: They took a tool that makes an AI refuse to answer harmful questions.
- Old Method: The tool failed. The AI started accepting harmful requests instead of refusing them (because the sign was flipped).
- New Method: The tool worked perfectly, preserving 95.8% of its original power.
The "Dictionary" Test: They tried to move a dictionary of features (SAE) from one model to another.
- Old Method: The dictionary was useless (reconstruction error was huge).
- New Method: The dictionary worked almost perfectly.
The "Resume" Test: They paused training an AI, changed the labels (gauge), and tried to resume training.
- Old Method: The AI forgot where it was and took a completely different path.
- New Method: The AI continued exactly where it left off, as if nothing happened.

The Big Takeaway

The paper concludes that for modern AI models, you cannot just ask, "Which neuron is this?" without first specifying the "gauge" (the rulebook for how the labels and signs are arranged).

If you want to move tools, dictionaries, or training states between these models, you must use the new "Signed-Permutation" rules. If you don't, you are essentially trying to drive a car with the steering wheel flipped 180 degrees: you think you're going straight, but you're actually going backward.

In short: Modern AI models have a hidden "sign flip" symmetry. To move things between them, you must fix the signs, not just the names. The paper provides the map and the compass to do exactly that.

Technical Summary: Signed-Permutation Coordinate Transport for RMSNorm Transformers

Problem Statement

Modern Large Language Model (LLM) workflows increasingly rely on moving coordinate-indexed artifacts—such as steering vectors, sparse autoencoder (SAE) dictionaries, top-k neuron sets, and LoRA adapters—across different model checkpoints. However, the question "which neuron is this?" is ill-posed without a fixed gauge (coordinate system). While standard alignment tools (e.g., Git Re-Basin, weight matching) assume a permutation-only symmetry group ( $S_d$ ), this assumption is incorrect for models utilizing RMSNorm (Root Mean Square Layer Normalization).

Unlike LayerNorm, which includes mean-centering and restricts symmetry to permutations (up to a global sign flip), RMSNorm lacks mean-centering. Consequently, its residual stream admits independent sign flips for every coordinate in addition to permutations. The native discrete gauge for RMSNorm is the hyperoctahedral group $B_d = S_d \ltimes \{\pm 1\}^d$ (signed permutations), not $S_d$ . Ignoring the sign component ( $S_d$ -only alignment) leads to a structural failure mode where coordinate-indexed artifacts are misaligned, causing significant performance degradation or complete functional collapse when transferred between models.

Methodology

1. Theoretical Characterization of Native Gauges

The paper establishes that for transformers with generic per-channel gain $\gamma$ :

LayerNorm: The maximal coordinate-preserving subgroup of $O(d)$ is $S_d \times C_2$ (permutations plus a global sign flip).
RMSNorm: The maximal coordinate-preserving subgroup is the full signed-permutation group $B_d$ .
Implication: Continuous orthogonal rotations (used in methods like SliceGPT or QuaRot) are valid for function preservation but destroy sparse coordinate identity. Only $B_d$ preserves the identity of sparse objects (e.g., specific neurons or SAE features).

2. Sign-Marginalized Gauge Recovery

To recover the correct $B_d$ gauge from paired activations, the paper proposes a modification to the standard Hungarian matching algorithm:

Cost Function: Instead of maximizing the signed correlation $\sum C_{i, \pi(i)}$ , the algorithm maximizes the absolute correlation $\sum |C_{i, \pi(i)}|$ .
Sign Recovery: Once the optimal permutation $\pi^*$ is found via the absolute cost matrix, the signs are recovered as $s_i = \text{sign}(C_{i, \pi^*(i)})$ .
Theoretical Guarantee: The paper proves that in the decorrelated-coordinate population case, standard signed-correlation matching has a structural accuracy ceiling of approximately 50% (the fraction of positive signs) because it actively avoids the correct anti-correlated entries. Sign-marginalized matching removes this obstruction, converging to 100% accuracy.

3. Parallel Transport via Trajectory Composition

Rather than matching only two endpoints, the paper introduces parallel transport along a training trajectory:

Local $B_d$ gauges are estimated between consecutive checkpoints using sign-marginalized matching.
These local gauges are composed to form a global transport map between distant checkpoints (or across runs sharing a base).
This approach avoids the "holonomy" issues of direct endpoint matching and leverages the high assignment margins of consecutive checkpoints.

4. Gauge-Covariant Transport Rules

The paper derives explicit transformation rules for various coordinate-indexed artifacts under $B_d$ (Table 1):

Vectors (Steering, Probes, SAE Encoders): Transform as $v \mapsto vG$ .
Matrices (SAE Decoders, LoRA WRITE): Transform as $W \mapsto G^\top W$ .
Optimizer States (AdamW): The first moment $m_W$ transforms with the weights; the second moment $v_W$ (being a square) is sign-blind and transforms via permutation only.

Key Results

1. Recovery Accuracy

Sign-Marginalized vs. Signed: In experiments on 7B/8B RMSNorm models, naive signed-correlation matching recovered only ~49.5% of the permutation and ~50% of signs. Sign-marginalized matching achieved 100% recovery.
Transport vs. Endpoint Matching: On a Qwen2.5-1.5B fine-tuning benchmark (1500 steps), transport via base composition recovered 91.1% of coordinates, whereas direct endpoint matching recovered only 60.3%. The improvement was not due to referencing the base alone (endpoint-via-base was 11.4%) but from composing many high-margin local gauges.

2. Artifact Transfer Performance

Using $B_d$ -correct transport significantly outperformed $S_d$ -only transport across multiple artifact classes:

LoRA Adapters: On Qwen2.5-1.5B (SST-2), $S_d$ transport collapsed adapter accuracy to the unadapted base level (0.523 vs. 0.525), while $B_d$ transport preserved 99.7% of the gain (0.944).
SAE Reconstruction: For TinyLlama, $S_d$ transport resulted in a reconstruction NMSE of 1.08 (failure), whereas $B_d$ transport achieved 0.004 (near-perfect).
Steering Vectors: Qwen sentiment steering preserved 95.8% of its effect under $B_d$ vs. 17.2% under $S_d$ . Refusal steering vectors reversed sign under $S_d$ , effectively flipping the behavior.

3. Stateful Training and Merging

Optimizer State: Resuming training with $S_d$ -only transported AdamW states (permutation correct, signs forced to +1) caused the training trajectory to diverge from the reference, despite starting from a functionally identical checkpoint. $B_d$ transport preserved the trajectory to numerical tolerance.
Model Merging: In coordinate-preserving merges, $S_d$ -only alignment introduced destructive sign mismatches, creating high "peak-above-chord" barriers. Signed alignment eliminated these barriers (reducing them to 0.00 in gauge-scramble tests).

Significance and Claims

The paper argues that for RMSNorm-based architectures (including widely used models like Llama 2, Llama 3, and Qwen), coordinate identity is only well-defined up to signed permutation.

Reproducibility: Coordinate-indexed interpretability claims (e.g., "neuron X is a refusal neuron") are not reproducible across runs unless a specific $B_d$ gauge is fixed and applied. The paper proposes "gauge sweeps" as a negative control to audit such claims.
Tool Transfer: The failure of $S_d$ -only methods is not a minor error but a structural collapse. Signed transport is required to successfully move steering vectors, SAE dictionaries, and optimizer states.
Limitations: The paper notes that while $B_d$ preserves coordinate identity, continuous orthogonal methods remain useful for dense fusion or compression where coordinate identity is not the primary object. Additionally, FFN hidden units possess their own independent permutation symmetries ( $S_{d_{ffn}}$ ) that are distinct from the residual stream gauge.

In summary, the work provides the theoretical foundation and practical algorithms to correctly align RMSNorm models, demonstrating that ignoring the sign component of the gauge renders coordinate-based workflows fundamentally broken.