Signed-Permutation Coordinate Transport for RMSNorm Transformers
Introduces signed-permutation coordinate transport for RMSNorm transformers, recovering 91.1% of coordinates.
Key Findings
Methodology
The paper introduces a sign-marginalized Hungarian matching algorithm to address the symmetry-incomplete issue in RMSNorm models. It focuses on coordinate-preserving transport by composing local Bd gauges along fine-tuning trajectories rather than function-level merging.
Key Results
- In 1500-step fine-tuning, coordinate-preserving transport recovered 91.1% of cross-run coordinates, compared to 60.3% for endpoint matching.
- TinyLlama SAE reconstruction achieved NMSE of 0.004 under Bd, versus 1.08 under Sd.
- Qwen sentiment steering preserved 95.8% of its effect under Bd, compared to 17.2% under Sd.
Significance
This research addresses the symmetry-incomplete issue in RMSNorm models, significantly improving cross-run coordinate recovery. It has important implications for model transfer and parameter sharing in academia and industry.
Technical Contribution
Technical contributions include the introduction of a sign-marginalized Hungarian matching algorithm and proving a structural permutation-accuracy ceiling in decorrelated coordinates. It also demonstrates coordinate-preserving transport via local Bd gauges.
Novelty
This is the first application of sign-marginalized matching in RMSNorm models, solving the symmetry-incomplete issue, and offers significant innovation over existing permutation-only alignment methods.
Limitations
- High probe budget requirement for natural alignment between independently trained models.
- Method relies on saved checkpoints and local Bd gauges.
Future Work
Future work could explore applying this method to larger models and datasets, and reducing probe budget requirements.
AI Executive Summary
In modern large language model (LLM) workflows, moving coordinate-indexed objects across checkpoints is increasingly common. However, existing methods face issues with symmetry-incompleteness in RMSNorm models, leading to inaccurate coordinate recovery. This paper introduces a sign-marginalized Hungarian matching algorithm, focusing on coordinate-preserving transport by composing local Bd gauges along fine-tuning trajectories, significantly improving coordinate recovery accuracy.
Experimental results show that this method recovered 91.1% of cross-run coordinates in 1500-step fine-tuning, compared to only 60.3% for traditional endpoint matching. Additionally, TinyLlama SAE reconstruction achieved NMSE of 0.004 under Bd, versus 1.08 under Sd, and Qwen sentiment steering preserved 95.8% of its effect under Bd, compared to 17.2% under Sd.
This research is significant not only in academia but also provides new insights for model transfer and parameter sharing in the industry. Future research could further explore the application of this method to larger models and datasets, and how to reduce probe budget requirements.
Deep Analysis
Background
With the evolution of deep learning, the scale and complexity of models have increased. RMSNorm, a common normalization technique, is widely used in large language models. However, existing model alignment methods face issues with symmetry-incompleteness in RMSNorm models, leading to inaccurate coordinate recovery.
Core Problem
The symmetry-incomplete issue in RMSNorm models leads to inaccurate coordinate recovery. This is because RMSNorm models allow independent sign flips for each coordinate, while existing methods only consider permutation alignment.
Innovation
The paper introduces a sign-marginalized Hungarian matching algorithm to solve the symmetry-incomplete issue in RMSNorm models. By composing local Bd gauges for coordinate-preserving transport, it significantly improves coordinate recovery accuracy.
Methodology
- �� Introduce a sign-marginalized Hungarian matching algorithm to address the symmetry-incomplete issue.
- �� Compose local Bd gauges along the same base fine-tuning trajectory for coordinate-preserving transport.
- �� Validate the method's effectiveness in RMSNorm models through experiments.
Experiments
Experiments used the Qwen2.5-1.5B model for fine-tuning, comparing the effects of sign-marginalized matching and traditional endpoint matching. Results show that sign-marginalized matching significantly improves coordinate recovery accuracy.
Results
Experimental results show that sign-marginalized matching recovered 91.1% of cross-run coordinates in 1500-step fine-tuning, compared to only 60.3% for traditional endpoint matching. Additionally, TinyLlama SAE reconstruction achieved NMSE of 0.004 under Bd, versus 1.08 under Sd.
Applications
This method can be applied to model transfer and parameter sharing in large language models, especially in scenarios requiring high-precision coordinate recovery, such as sentiment analysis and semantic understanding.
Limitations & Outlook
The method relies on saved checkpoints and local Bd gauges, and requires a high probe budget for natural alignment between independently trained models. Future research could explore how to reduce probe budget requirements.
Plain Language Accessible to non-experts
Imagine you're in a library where books are arranged in a specific order on shelves. Each book has a label indicating its position. Now, you need to move these books to another library while keeping their order intact. The problem is, the new library allows you to flip the labels during transport, just like the sign flips in RMSNorm models. To ensure the order remains unchanged, you need a method to handle these label flips. The method proposed in this paper acts like a clever librarian who can automatically adjust the labels during transport, ensuring the books' order and position remain unchanged.
ELI14 Explained like you're 14
Imagine you're playing a puzzle game. Each piece has a special pattern, and you need to place them in the right spots. Sometimes, these pieces can flip, just like the sign flips in RMSNorm models. To make sure the puzzle is completed correctly, you need a way to handle these flips. The method in this paper is like a super-smart puzzle master who can automatically adjust each piece's position and direction during the puzzle, ensuring the final picture is perfect!
Glossary
Signed Permutation
In mathematics, a signed permutation is a permutation along with a flip of signs for each element.
In RMSNorm models, signed permutation describes the flipping and rearranging of coordinates.
Hungarian Matching
An algorithm used to solve assignment problems, capable of finding optimal matches in polynomial time.
The paper uses the Hungarian matching algorithm for sign-marginalized matching.
RMSNorm
A normalization technique that does not perform mean-centering, only standard deviation normalization.
RMSNorm is widely used in large language models.
Coordinate-preserving Transport
A method aimed at maintaining coordinate consistency across different model checkpoints.
The proposed sign-marginalized matching is used to achieve coordinate-preserving transport.
Local Bd Gauge
A mathematical structure used to describe coordinate flips and permutations in RMSNorm models.
Local Bd gauges are used to achieve coordinate-preserving transport.
Open Questions Unanswered questions from this research
- 1 How to apply sign-marginalized matching to larger models and datasets remains to be explored.
- 2 Effective strategies to reduce probe budget requirements are not yet clear.
Applications
Immediate Applications
Model Transfer
This method can improve the transferability of large language models across different tasks, especially in scenarios requiring high-precision coordinate recovery.
Long-term Vision
Cross-model Parameter Sharing
By improving coordinate recovery accuracy, this method could facilitate parameter sharing between different models, advancing deep learning further.
Abstract
Modern LLM workflows move coordinate-indexed objects across checkpoints: steering vectors, sparse autoencoders, top-$k$ neuron sets, attribution lists, and merge alignments. This is only well posed after fixing the model's residual-stream gauge, which we show is architecture-dependent: LayerNorm residual charts have permutation gauge $S_d$ (up to a global sign flip), while RMSNorm charts with generic per-channel gain have signed-permutation gauge $B_d = S_d \ltimes \{\pm 1\}^d$. Permutation-only alignment is therefore symmetry-incomplete for RMSNorm models. We introduce sign-marginalized Hungarian matching and prove a sharp failure mode: with decorrelated coordinates, raw signed-correlation matching has a structural permutation-accuracy ceiling at the positive-sign fraction of the true gauge, which sign-marginalization removes. We then make coordinate-preserving transport, not function-level merging, the primary object: composing saved-checkpoint local $B_d$ gauges along same-base fine-tuning trajectories recovers 91.1% of cross-run coordinates at 1500 steps versus 60.3% for endpoint matching, and the gain is not explained by merely routing through the base. The recovered gauge transfers tools that permutation-only alignment breaks: TinyLlama SAE reconstruction has NMSE 0.004 under $B_d$ versus 1.08 under $S_d$; Qwen sentiment steering preserves 95.8% of its effect versus 17.2%; refusal steering reverses sign under $S_d$; coordinate-preserving merges behave the same way. The same covariance governs stateful training: signed transport of AdamW state preserves the resumed trajectory, while permutation-only state follows a different one from a functionally identical checkpoint. Finally, gauge-sweep audits show index-level interpretability claims are reproducible only relative to an explicit gauge.