Resumo Técnico: Transporte de Coordenadas por Permutação Assinada para Transformers com RMSNorm

Definição do Problema

Os fluxos de trabalho modernos de Grandes Modelos de Linguagem (LLMs) dependem cada vez mais da movimentação de artefatos indexados por coordenadas — como vetores de direção (steering vectors), dicionários de Autoencoders Esparsos (SAEs), conjuntos de neurônios top-k e adaptadores LoRA — entre diferentes checkpoints de modelos. No entanto, a pergunta "qual neurônio é este?" é mal formulada sem um calibre (gauge) fixo (sistema de coordenadas). Embora ferramentas de alinhamento padrão (ex: Git Re-Basin, correspondência de pesos) assumam um grupo de simetria de permutação apenas ($S_d$), essa suposição é incorreta para modelos que utilizam RMSNorm (Normalização de Raiz Quadrática Média).

Diferente da LayerNorm, que inclui centralização de média e restringe a simetria a permutações (até uma inversão de sinal global), a RMSNorm carece de centralização de média. Consequentemente, seu fluxo residual admite inversões de sinal independentes para cada coordenada, além de permutações. O calibre discreto nativo para RMSNorm é o grupo hiperoctaédrico $B_d = S_d \ltimes \{\pm 1\}^d$ (permutações assinadas), não o $S_d$. Ignorar o componente de sinal (alinhamento apenas de $S_d$) leva a um modo de falha estrutural onde artefatos indexados por coordenadas ficam desalinhados, causando degradação significativa de desempenho ou colapso funcional completo durante a transferência entre modelos.

Metodologia

1. Caracterização Teórica dos Calibres Nativos

O artigo estabelece que para transformers com ganho por canal $\gamma$ genérico:

• LayerNorm: O subgrupo máximo de preservação de coordenadas de $O(d)$ é $S_d \times C_2$ (permutações mais uma inversão de sinal global).
• RMSNorm: O subgrupo máximo de preservação de coordenadas é o grupo completo de permutação assinada $B_d$.
• Implicação: Rotações ortogonais contínuas (usadas em métodos como SliceGPT ou QuaRot) são válidas para preservação de função, mas destroem a identidade de coordenadas esparsas. Apenas $B_d$ preserva a identidade de objetos esparsos (ex: neurônios específicos ou características de SAE).

2. Recuperação de Calibre com Marginalização de Sinal

Para recuperar o calibre $B_d$ correto a partir de ativações pareadas, o artigo propõe uma modificação ao algoritmo de correspondência de Hungarian padrão:

• Função de Custo: Em vez de maximizar a correlação assinada $\sum C_{i, \pi(i)}$, o algoritmo maximiza a correlação absoluta $\sum |C_{i, \pi(i)}|$.
• Recuperação de Sinal: Uma vez encontrada a permutação ótima $\pi^*$ via matriz de custo absoluto, os sinais são recuperados como $s_i = \text{sign}(C_{i, \pi^*(i)})$.
• Garantia Teórica: O artigo prova que, no caso de uma população de coordenadas decorrelacionadas, a correspondência de correlação assinada padrão possui um teto de precisão estrutural de aproximadamente 50% (a fração de sinais positivos) porque ela evita ativamente as entradas anti-correlacionadas corretas. A correspondência com marginalização de sinal remove essa obstrução, convergindo para 100% de precisão.

3. Transporte Paralelo via Composição de Trajetória

Em vez de alinhar apenas dois pontos finais, o artigo introduz o transporte paralelo ao longo de uma trajetória de treinamento:

• Calibres $B_d$ locais são estimados entre checkpoints consecutivos usando correspondência com marginalização de sinal.
• Esses calibres locais são compostos para formar um mapa de transporte global entre checkpoints distantes (ou entre execuções que compartilham uma base).
• Esta abordagem evita os problemas de "holonomia" do alinhamento direto de pontos finais e aproveita as altas margens de atribuição de checkpoints consecutivos.

4. Regras de Transporte Covariantes ao Calibre

O artigo deriva regras de transformação explícitas para vários artefatos indexados por coordenadas sob $B_d$ (Tabela 1):

• Vetores (Direcionamento, Probes, Codificadores SAE): Transformam-se como $v \mapsto vG$.
• Matrizes (Decodificadores SAE, Escrita LoRA): Transformam-se como $W \mapsto G^\top W$.
• Estados de Otimizador (AdamW): O primeiro momento $m_W$ transforma-se com os pesos; o segundo momento $v_W$ (sendo um quadrado) é cego ao sinal e transforma-se via permutação apenas.

Principais Resultados

1. Precisão de Recuperação

• Marginalização de Sinal vs. Assinado: Em experimentos em modelos RMSNorm de 7B/8B, a correspondência de correlação assinada ingênua recuperou apenas ~49,5% da permutação e ~50% dos sinais. A correspondência com marginalização de sinal alcançou 100% de recuperação.
• Transporte vs. Correspondência de Ponto Final: Em um benchmark de ajuste fino do Qwen2.5-1.5B (1500 passos), o transporte via composição de base recuperou 91,1% das coordenadas, enquanto a correspondência direta de ponto final recuperou apenas 60,3%. A melhoria não se deveu apenas à referência à base (ponto final via base foi de 11,4%), mas da composição de muitos calibres locais de alta margem.

2. Desempenho de Transferência de Artefatos

O transporte com $B_d$ correto superou significativamente o transporte apenas de $S_d$ em múltiplas classes de artefatos:

• Adaptadores LoRA: No Qwen2.5-1.5B (SST-2), o transporte $S_d$ colapsou a precisão do adaptador para o nível da base não adaptada (0,523 vs. 0,525), enquanto o transporte $B_d$ preservou 99,7% do ganho (0,944).
• Reconstrução de SAE: Para o TinyLlama, o transporte $S_d$ resultou em um NMSE de reconstrução de 1,08 (falha), enquanto o transporte $B_d$ alcançou 0,004 (quase perfeito).
• Vetores de Direcionamento (Steering): O direcionamento de sentimento do Qwen preservou 95,8% de seu efeito sob $B_d$ contra 17,2% sob $S_d$. Vetores de direcionamento de recusa inverteram o sinal sob $S_d$, efetivamente invertendo o comportamento.

3. Treinamento de Estado e Fusão

• Estado do Otimizador: Retomar o treinamento com estados AdamW transportados apenas por $S_d$ (permutação correta, sinais forçados para +1) fez com que a trajetória de treinamento divergisse da referência, apesar de partir de um checkpoint funcionalmente idêntico. O transporte $B_d$ preservou a trajetória dentro da tolerância numérica.
• Fusão de Modelos (Model Merging): Em fusões que preservam coordenadas, o alinhamento apenas de $S_d$ introduziu incompatibilidades de sinal destrutivas, criando barreiras de "pico acima da corda" elevadas. O alinhamento assinado eliminou essas barreiras (reduzindo-as a 0,00 em testes de desordem de calibre).

Significância e Alegações

O artigo argumenta que para arquiteturas baseadas em RMSNorm (incluindo modelos amplamente utilizados como Llama 2, Llama 3 e Qwen), a identidade de coordenada só é bem definida até uma permutação assinada.

• Reprodutibilidade: Alegações de interpretabilidade indexadas por coordenadas (ex: "o neurônio X é um neurônio de recusa") não são reprodutíveis entre execuções, a menos que um calibre $B_d$ específico seja fixado e aplicado. O artigo propõe "varreduras de calibre" (gauge sweeps) como um controle negativo para auditar tais alegações.
• Transferência de Ferramentas: A falha dos métodos apenas de $S_d$ não é um erro menor, mas um colapso estrutural. O transporte assinado é necessário para mover com sucesso vetores de direção, dicionários de SAE e estados de otimizador.
• Limitações: O artigo observa que, embora $B_d$ preserve a identidade de coordenada, métodos ortogonais contínuos permanecem úteis para fusão densa ou compressão, onde a identidade de coordenada não é o objeto primário. Além disso, as unidades ocultas de FFN possuem suas próprias simetrias de permutação independentes ($S_{d_{ffn}}$) que são distintas do calibre do fluxo residual.

Em resumo, o trabalho fornece a fundação teórica e os algoritmos práticos para alinhar corretamente modelos RMSNorm, demonstrando que ignorar o componente de sinal do calibre torna os fluxos de trabalho baseados em coordenadas fundamentalmente quebrados.