Riepilogo Tecnico: Asimmetria Geometrica nella Specializzazione MoE

Enunciato del Problema

Mentre le architetture Mixture-of-Experts (MoE) sono diventate un meccanismo standard per scalare i modelli Transformer consentendo grandi quantità di parametri con costi computazionali fissi, la struttura geometrica interna della specializzazione degli esperti rimane poco compresa. Le analisi esistenti si concentrano prevalentemente sulle statistiche di instradamento, sul bilanciamento del carico o sulle prestazioni a livello di sistema. Esiste un divario critico nella comprensione di come gli esperti differiscano nello spazio delle funzioni (le trasformazioni che eseguono) rispetto allo spazio delle rappresentazioni (come i token instradati sono organizzati negli stati nascosti).

L'intuizione corrente suggerisce che i livelli MoE partizionino la distribuzione di input in cluster disgiunti, dove ogni esperto agisce come un approssimatore locale. Tuttavia, evidenze empiriche indicano che gli esperti operano spesso su regimi di token sovrapposti, sollevando la domanda: Quale struttura geometrica induce l'instradamento all'interno di un Transformer MoE preaddestrato? Nello specifico, gli esperti implementano funzioni distinte e come si relazionano le loro rappresentazioni tra loro?

Metodologia

Gli autori introducono un framework unificato Jacobiano-PCA-Grassmann per analizzare i livelli MoE sia nello spazio delle funzioni che in quello delle rappresentazioni. Questo framework viene applicato a modelli preaddestrati (Mistral, Qwen) e a modelli Transformer controllati a 3 livelli addestrati da zero.

1. Geometria Funzionale: Jacobiani Locali agli Esperti

Per caratterizzare la funzione dei singoli esperti indipendentemente dalle dinamiche di instradamento, gli autori analizzano il jacobiano locale all'esperto ($J_e(x) = \partial f_e(x) / \partial x$).

• Metrica: Calcolano il jacobiano medio ($\bar{J}_e$) per ogni esperto su un dataset e lo vettorializzano.
• Allineamento: L'allineamento funzionale a coppie è misurato tramite il prodotto interno di Frobenius normalizzato (similarità del coseno) tra questi jacobiani vettorializzati.
• Interpretazione: Una bassa similarità del coseno indica decorrelazione funzionale, il che significa che gli esperti implementano trasformazioni localmente distinte piuttosto che varianti ridondanti della stessa funzione.

2. Geometria delle Rappresentazioni: PCA Instradata e Distanze Grassmanniane

Per caratterizzare la geometria delle rappresentazioni instradate, gli autori raccolgono gli stati nascosti ponderati dalle pesi di instradamento ($h_e(x) = g_e(x)h(x)$).

• PCA: Viene eseguita l'Analisi delle Componenti Principali (PCA) standard sugli stati nascosti instradati per identificare le direzioni principali e gli spettri di varianza.
• Distanza del Sottospazio: La separazione geometrica tra gli esperti è quantificata utilizzando la Distanza Geodetica Grassmanniana tra i sottospazi PCA top-$n$ di diversi esperti.
• Interpretazione: Questa metrica cattura la separazione multidirezionale tra interi sottospazi. Una distanza vicina allo zero implica ruoli funzionali ridondanti, mentre una distanza che si avvicina al massimo teorico ($\frac{\pi}{2}\sqrt{n}$) implica ortogonalità.

3. Esperimenti di Instradamento Controllati

Per isolare l'effetto dei meccanismi di instradamento dalle dinamiche di addestramento, gli autori confrontano due regimi su un modello controllato a 3 livelli con parametri degli esperti fissi:

• Instradamento Top-k: Selezione sparsa e discreta.
• Instradamento Fully-Soft: Miscela densa e continua.

Risultati Chiave

1. Asimmetria Strutturale: Decorrelazione Funzionale vs. Sovrapposizione delle Rappresentazioni

Nei Transformer MoE preaddestrati (Mistral-8x7B e Qwen1.5-MoE-A2.7B), gli autori osservano un'asimmetria strutturale coerente:

• Forte Decorrelazione Funzionale: L'allineamento dei jacobiani tra esperti è costantemente vicino allo zero (media $\approx$ 0.000–0.068). Ciò indica che gli esperti non sono ridondanti; implementano trasformazioni locali distinte.
• Sovrapposizione Parziale delle Rappresentazioni: Nonostante la decorrelazione funzionale, le rappresentazioni instradate occupano sottospazi distinti ma parzialmente sovrapposti. Le distanze Grassmanniane sono sostanziali (indicando separazione) ma rimangono significativamente al di sotto del massimo teorico per l'ortogonalità completa.
• Conclusione: Decorrelazione funzionale e sovrapposizione delle rappresentazioni coesistono. Gli esperti agiscono come operatori localmente decorrelati su regimi sovrapposti su una varietà condivisa, piuttosto che partizionare lo spazio di input in cluster disgiunti.

2. Sparsità dell'Instradamento come Manopola di Controllo Geometrico

Gli esperimenti controllati rivelano che la nitidezza dell'instradamento plasma direttamente questa geometria:

• Instradamento Top-k: Induce una separazione funzionale più netta (bassa allineamento dei jacobiani) e distanze Grassmanniane più ampie (maggiore divergenza dei sottospazi).
• Instradamento Fully-Soft: Produce strutture di esperti intrecciate con alto allineamento dei jacobiani (indicando un collasso verso mappature ridondanti) e distanze Grassmanniane collassate (alta sovrapposizione dei sottospazi).
• Implicazione: La sparsità dell'instradamento è un fattore chiave nell'imporre la specializzazione.

3. Confronto con Livelli Densi

Gli esperti MoE mostrano un decadimento degli autovalori più ripido nei loro spettri di varianza rispetto ai livelli feedforward densi. Le prime componenti principali catturano una percentuale di varianza significativamente più alta negli esperti MoE, suggerendo che l'instradamento concentra i token in sottospazi a dimensionalità inferiore e più focalizzati.

4. Coerenza Cross-Modello

I pattern geometrici osservati (basso allineamento dei jacobiani + alta ma sub-massimale separazione Grassmanniana) sono coerenti tra diversi modelli (Mistral, Qwen), livelli (iniziali, intermedi, finali) e dataset, suggerendo che questa sia una proprietà robusta delle architetture MoE piuttosto che un artefatto di esecuzioni di addestramento specifiche.

Significato e Affermazioni

Il documento afferma di fornire un framework diagnostico generale per lo studio del calcolo condizionale nelle moderne architetture Transformer. I suoi contributi principali sono:

1. Strumento Diagnostico Unificato: Offre un metodo sistematico per analizzare la specializzazione degli esperti oltre le semplici statistiche di instradamento o i pattern di attivazione, caratterizzando congiuntamente lo spazio delle funzioni e lo spazio delle rappresentazioni.
2. Interpretazione Geometrica di MoE: I risultati supportano una visione dei livelli MoE come implementazione di una partizionamento morbido di una varietà di rappresentazione condivisa. Gli esperti sono funzionalmente distinti ma operano su sottovarietà sovrapposte.
3. Insight Meccanicistico: Lo studio identifica la nitidezza dell'instradamento (sparsità) come un meccanismo di controllo diretto per il grado di specializzazione, collegando le scelte di progettazione architetturale a strutture geometriche misurabili.
4. Implicazioni per la Progettazione del Modello: Gli autori suggeriscono che le strategie per l'unione o la potatura degli esperti basate esclusivamente sulla similarità delle rappresentazioni (ad esempio, sovrapposizione delle attivazioni) potrebbero essere errate. Poiché gli esperti possono avere rappresentazioni sovrapposte pur eseguendo funzioni distinte, potrebbero essere necessarie metriche funzionali (come l'allineamento dei jacobiani) per identificare correttamente la vera ridondanza.

Il documento mantiene un ambito modesto, riconoscendo che la sua analisi è locale (jacobiani del primo ordine) e lineare (PCA), e non cattura il comportamento non lineare globale o l'evoluzione della geometria durante l'addestramento. Tuttavia, stabilisce una prospettiva geometrica fondamentale su come i modelli MoE organizzano il calcolo.