Sintesi Tecnica: Trasporto di Coordinate tramite Permutazione con Segno per Transformer con RMSNorm

Definizione del Problema

I moderni flussi di lavoro per i Large Language Model (LLM) si affidano sempre più al movimento di artefatti indicizzati per coordinate — come vettori di steering, dizionari di Sparse Autoencoders (SAE), set di neuroni top-k e adattatori LoRA — attraverso diversi checkpoint del modello. Tuttavia, la domanda "quale neurone è questo?" è mal posta senza un gauge (sistema di coordinate) fisso. Sebbene gli strumenti di allineamento standard (ad esempio, Git Re-Basin, weight matching) assumano un gruppo di simmetria di sola permutazione ($S_d$), tale assunzione è errata per i modelli che utilizzano RMSNorm (Root Mean Square Layer Normalization).

A differenza della LayerNorm, che include la centratura della media e limita la simmetria alle sole permutazioni, la RMSNorm manca della centratura della media. Di conseguenza, il suo stream residuo ammette flip di segno indipendenti per ogni coordinata, oltre alle permutazioni. Il gauge discreto nativo per RMSNorm è il gruppo delle permutazioni con segno (iperottaedrico) $B_d = S_d \ltimes \{\pm 1\}^d$, non $S_d$. Ignorare la componente del segno (allineamento basato solo su $S_d$) porta a un fallimento strutturale in cui gli artefatti indicizzati per coordinate risultano disallineati, causando un significativo degrado delle prestazioni o un completo collasso funzionale durante il trasferimento tra modelli.

Metodologia

1. Caratterizzazione Teorica dei Gauge Nativi

Il documento stabilisce che per i transformer con guadagno per canale $\gamma$ generico:

• LayerNorm: Il sottogruppo massimale di $O(d)$ che preserva le coordinate è $S_d \times C_2$ (permutazioni più un flip di segno globale).
• RMSNorm: Il sottogruppo massimale che preserva le coordinate è il gruppo completo delle permutazioni con segno $B_d$.
• Implicazione: Le rotazioni ortogonali continue (usate in metodi come SliceGPT o QuaRot) sono valide per la preservazione della funzione ma distruggono l'identità sparsa della coordinata. Solo $B_d$ preserva l'identità degli oggetti sparsi (ad esempio, specifici neuroni o feature SAE).

2. Recupero del Gauge tramite Marginalizzazione del Segno

Per recuperare il corretto gauge $B_d$ da coppie di attivazioni, il documento propone una modifica all'algoritmo standard di matching di Hungarian:

• Funzione di Costo: Invece di massimizzare la correlazione firmata $\sum C_{i, \pi(i)}$, l'algoritmo massimizza la correlazione assoluta $\sum |C_{i, \pi(i)}|$.
• Recupero del Segno: Una volta trovata la permutazione ottimale $\pi^*$ tramite la matrice di costo assoluto, i segni vengono recuperati come $s_i = \text{sign}(C_{i, \pi^*(i)})$.
• Garanzia Teorica: Il documento dimostra che nel caso di una popolazione di coordinate decorrelate, il matching standard basato su correlazione firmata ha un tetto di accuratezza strutturale di circa il 50% (la frazione di segni positivi) perché evita attivamente le voci anti-correlate corrette. Il matching con marginalizzazione del segno rimuove questa ostruzione, convergendo verso il 100% di accuratezza.

3. Trasporto Parallelo tramite Composizione di Traiettorie

Invece di accoppiare solo due endpoint, il documento introduce il trasporto parallelo lungo una traiettoria di addestramento:

• I gauge locali $B_d$ vengono stimati tra checkpoint consecutivi utilizzando il matching con marginalizzazione del segno.
• Questi gauge locali vengono composti per formare una mappa di trasporto globale tra checkpoint distanti (o tra diverse esecuzioni che condividono una base).
• Questo approccio evita i problemi di "olonomia" del matching diretto tra endpoint e sfrutta gli alti margini di assegnazione dei checkpoint consecutivi.

4. Regole di Trasporto Covarianti rispetto al Gauge

Il documento deriva esplicite regole di trasformazione per vari artefatti indicizzati per coordinate sotto $B_d$ (Tabella 1):

• Vettori (Steering, Probe, Encoder SAE): Si trasformano come $v \mapsto vG$.
• Matrici (Decoder SAE, LoRA WRITE): Si trasformano come $W \mapsto W G^\top$.
• Stati dell'Ottimizzatore (AdamW): Il primo momento $m_W$ si trasforma con i pesi; il secondo momento $v_W$ (essendo un quadrato) è cieco al segno e si trasforma tramite sola permutazione.

Risultati Chiave

1. Accuratezza di Recupero

• Marginalizzazione del Segno vs. Firmato: In esperimenti su modelli RMSNorm da 7B/8B, il matching ingenuo basato su correlazione firmata ha recuperato solo il ~49,5% delle permutazioni e il ~50% dei segni. Il matching con marginalizzazione del segno ha raggiunto il 100% del recupero.
• Trasporto vs. Endpoint Matching: Su un benchmark di fine-tuning di Qwen2.5-1.5B (1500 step), il trasporto tramite composizione della base ha recuperato il 91,1% delle coordinate, mentre il matching diretto tra endpoint ha recuperato solo il 60,3%. Il miglioramento non è dovuto al semplice riferimento alla base (l'endpoint-via-base era l'11,4%), ma dalla composizione di molti gauge locali ad alto margine.

2. Prestazioni di Trasferimento degli Artefatti

Il trasporto corretto rispetto a $B_d$ ha superato significativamente il trasporto basato solo su $S_d$ in diverse classi di artefatti:

• Adattatori LoRA: Su Qwen2.5-1B (SST-2), il trasporto $S_d$ ha fatto collassare l'accuratezza dell'adattatore al livello del base non adattato (0,523 vs 0,525), mentre il trasporto $B_d$ ha preservato il 99,7% del guadagno (0,944).
• Ricostruzione SAE: Per TinyLlama, il trasporto $S_d$ ha prodotto un NMSE di ricostruzione di 1,08 (fallimento), mentre il trasporto $B_d$ ha ottenuto 0,004 (quasi perfetto).
• Vettori di Steering: Lo steering del sentiment di Qwen ha preservato il 95,8% del suo effetto sotto $B_d$ rispetto al 17,2% sotto $S_d$. I vettori di steering del rifiuto hanno invertito il segno sotto $S_d$, invertendo efficacemente il comportamento.

3. Addestramento Stateful e Merging

• Stato dell'Ottimizzatore: Riprendere l'addestramento con stati AdamW trasportati solo tramite $S_d$ (permutazione corretta, segni forzati a +1) ha causato la divergenza della traiettoria di addestramento rispetto al riferimento, nonostante la partenza da un checkpoint funzionalmente identico. Il trasporto $B_d$ ha preservato la traiettoria entro la tolleranza numerica.
• Model Merging: Nei merge che preservano le coordinate, l'allineamento basato solo su $S_d$ ha introdotto mismatch di segno distruttivi, creando alte barriere "peak-above-chord". L'allineamento con segno ha eliminato queste barriere (riducendole a 0,00 nei test di scumble del gauge).

Significato e Rivendicazioni

Il documento sostiene che per le architetture basate su RMSNorm (inclusi modelli ampiamente utilizzati come Llama 2, Llama 3 e Qwen), l'identità della coordinata è ben definita solo fino alla permutazione con segno.

• Riproducibilità: Le affermazioni di interpretabilità indicizzate per coordinate (ad esempio, "il neurone X è un neurone di rifiuto") non sono riproducibili tra diverse esecuzioni a meno che non venga fissato e applicato un gauge $B_d$ specifico. Il documento propone i "gauge sweeps" come controllo negativo per auditare tali affermazioni.
• Trasferimento di Strumenti: Il fallimento dei metodi basati solo su $S_d$ non è un errore minore, ma un collasso strutturale. Il trasporto con segno è necessario per spostare con successo vettori di steering, dizionari SAE e stati dell'ottimizzatore.
• Limitazioni: Il documento nota che, sebbene $B_d$ preservi l'identità della coordinata, i metodi ortogonali continui rimangono utili per la fusione densa o la compressione dove l'identità della coordinata non è l'oggetto primario. Inoltre, le unità nascoste FFN possiedono le proprie simmetrie di permutazione indipendenti ($S_{d_{ffn}}$) che sono distinte dal gauge dello stream residuo.

In sintesi, il lavoro fornisce la base teorica e gli algoritmi pratici per allineare correttamente i modelli RMSNorm, dimostrando che ignorare la componente del segno del gauge rende i flussi di lavoro basati sulle coordinate fondamentalmente compromessi.