RMSNorm Transformer 的符號排列座標傳輸與規範固定研究

这篇论文聚焦于大语言模型内部表示对齐中的一个深层理论盲区：如何在不同的模型检查点之间准确传输坐标索引对象。随着模型工作流日益复杂，研究者需要跨检查点移动引导向量、稀疏自编码器特征、Top-k 神经元集合以及归因列表等关键数据。然而，这种传输只有在固定模型的残差流规范后才能良好定义。作者通过严谨的数学推导发现，不同归一化架构对应的规范群结构截然不同：LayerNorm 的残差流图表具有排列群 $S_d$ 的规范对称性（允许全局符号翻转），而 RMSNorm 图表由于引入了通用的每通道增益，其规范对称性扩展为符号排列群 $B_d = S_d \ltimes \{\pm 1\}^d$。这意味着，对于广泛使用的 RMSNorm 模型，仅进行排列对齐是不完整的，忽略了符号维度的对称性，从而导致后续对齐工具失效。这一发现揭示了当前许多模型编辑和可解释性方法在理论上的局限性，即它们往往错误地假设了模型的规范结构，从而在处理 RMSNorm 模型时产生系统性偏差。

为了解决这一规范不完整问题，作者提出了一种基于符号边际化的匈牙利匹配算法。该方法的核心在于不再简单地将坐标视为无序集合进行排列匹配，而是显式地处理符号对称性。作者证明了，如果直接使用原始符号相关匹配，在坐标去相关的条件下，其排列准确率存在一个结构性的上限，该上限仅由真实规范中正号的比例决定，无法突破。通过引入符号边际化，算法有效地移除了这一结构性限制，从而能够更准确地恢复真实的规范变换。在技术实现上，论文将坐标保持的传输而非函数级别的合并作为主要研究对象。通过沿着相同基线的微调轨迹组合保存检查点的局部 $B_d$ 规范，作者构建了一种能够精确追踪坐标变化的传输机制。

这种机制不仅考虑了坐标的排列顺序，还精确校正了每个坐标通道的符号翻转，确保了在模型微调过程中，内部表示的语义一致性得以维持，为后续的模型编辑和干预提供了坚实的理论基础。在实验设置与关键结果方面，作者在多个基准任务上验证了所提方法的有效性。首先，在坐标恢复实验中，通过在相同基线的微调轨迹上进行 1500 步的模拟，结果显示，使用 $B_d$ 规范进行坐标保持传输能够恢复 91.1% 的跨运行坐标，而传统的端点匹配方法仅能恢复 60.3%。这一显著增益并非仅仅源于通过基线进行路由，而是直接得益于对符号对称性的正确处理。其次，在具体的应用任务中，$B_d$ 规范的优势进一步凸显。在 TinyLlama 的稀疏自编码器（SAE）重建任务中，$B_d$ 规范下的归一化均方误差（NMSE）仅为 0.004，而仅使用排列对齐的 $S_d$ 规范下误差高达 1.08。

在 Qwen 模型的情感引导任务中，$B_d$ 规范保留了 95.8% 的引导效果，而 $S_d$ 规范下这一比例骤降至 17.2%。更令人担忧的是，在拒绝引导任务中，仅使用 $S_d$ 规范会导致引导符号反转，完全破坏干预效果。此外，作者还发现，这种协方差结构同样适用于状态训练，符号传输的 AdamW 状态能够保持恢复后的轨迹，而仅排列对齐的状态则会偏离原本功能相同的检查点轨迹。从行业意义与潜在影响来看，这项研究对开源社区和工业落地具有深远影响。它指出，当前许多基于排列对齐的工具在处理 RMSNorm 模型时存在根本性缺陷，可能导致模型编辑、可解释性分析等关键任务的结果不可靠甚至完全错误。通过引入 $B_d$ 规范，研究者可以构建更鲁棒的模型干预工具，确保在微调、合并和引导过程中内部表示的一致性。此外，规范扫描审计结果表明，许多基于索引的可解释性声明只有在相对于显式规范时才可复现，这呼吁社区在报告可解释性结果时必须明确规范假设。对于工业界而言，理解并应用符号排列传输有助于优化模型合并策略，提升微调效率，并减少因规范不一致导致的性能下降。未来研究可进一步探索如何在大规模模型中高效计算和传输 $B_d$ 规范，以及该框架在其他架构中的应用潜力，从而推动大模型内部表示对齐理论的完善与实践的标准化。