Research Brief

Correctly aligning internal coordinate representations in RMSNorm Transformers requires accounting for signed permutations (a 'signed-permutation gauge'), which is critical for robust interpretability, steerability, and stable fine-tuning.

This paper addresses a fundamental problem in understanding and manipulating large language models (LLMs): how internal components like 'steering vectors' or 'neuron sets' can be reliably tracked and transferred across different versions or checkpoints of a model. The authors identify that the method for aligning these components (the 'gauge') depends on the model's architectural normalization. Specifically, while LayerNorm models have a simpler 'permutation gauge', RMSNorm models, which are common in modern LLMs, require a more complex 'signed-permutation gauge' that accounts for both reordering and sign flips of internal coordinates. They demonstrate that ignoring these sign flips leads to significant failure in transferring tools like sparse autoencoders, sentiment steering, and even AdamW optimizer states. They introduce a method called 'sign-marginalized Hungarian matching' to correctly handle this signed gauge, dramatically improving the preservation of coordinate-indexed functionalities across model checkpoints and fine-tuning trajectories, thus enabling more reliable interpretability and controllable AI.

• Reliable LLM Interpretability and Steerability: Enabling consistent transfer of steering vectors, attribution lists, and neuron sets, making LLMs more controllable and their internal workings more understandable across different model versions.
• Robust Model Merging and Fine-tuning: Ensuring that merging different LLMs or continuing fine-tuning from a checkpoint preserves the intended internal structure and learned functionalities, rather than breaking them due to misaligned coordinates.
• Stable Optimizer State Transfer: Preserving the training trajectory when resuming training from checkpoints by correctly transporting AdamW optimizer states, crucial for efficient and consistent model development.
• Reproducible AI Research: Establishing a framework for index-level interpretability claims to be reproducible, making research findings more robust and verifiable across different experimental setups.