Riepilogo Tecnico: CLHA – Aggregazione dell'Essenziale Cross-Layer per la Quantizzazione di Transformer Mixture-of-Experts con Condivisione dei Pesi (Weight-Shared)

1. Definizione del Problema

Recenti architetture Transformer Mixture-of-Experts (MoE) efficienti, come quelle in Megrez2, impiegano la Cross-Layer Weight Sharing (CLWS). In questo design, i tensori dei pesi degli esperti routati adiacenti (ad esempio, $\ell$ e $\ell+1$) riutilizzano lo stesso peso, dimezzando l'impronta dei parametri pur mantenendo lo stesso numero di parametri attivati per token.

Tuttavia, gli esistenti metodi di Post-Training Quantization (PTQ), inclusi gli approcci di tipo GPTQ, trattano questi pesi condivisi come se appartenessero a un singolo strato. Calibrano la quantizzazione utilizzando le statistiche di attivazione di un solo strato "primario". Ciò crea un disallineamento fondamentale:

• La rappresentazione del residuo evolve tra gli strati, il che significa che le distribuzioni dei token che attivano l'esperto condiviso allo strato $\ell_a$ differiscono da quelle dello strato $\ell_b$.
• I metodi esistenti ignorano le statistiche di attivazione dello strato secondario, portando a una minimizzazione subottimale dell'errore di ricostruzione, particolarmente nei regimi a basso bit (ad esempio, 2-bit).

2. Metodologia

2.1 Fondamento Teorico: Aggregazione dell'Essenziale Cross-Layer

Il documento formalizza l'obiettivo per la quantizzazione di un peso $W$ condiviso tra gli strati $\ell_1, \dots, \ell_n$. L'obiettivo corretto è minimizzare l'errore totale di ricostruzione attraverso tutti gli strati utilizzando il peso:
$$L(W_q) = \sum_{i} \alpha_i \text{tr}[(W_q - W)H_{\ell_i}(W_q - W)^\top]$$
dove $H_{\ell_i}$ è l'Essenziale (covarianza delle attivazioni) dello strato $\ell_i$, e $\alpha_i$ rappresenta la frequenza di routing (proporzione di token) per quello strato.

La Proposizione 1 stabilisce che minimizzare questo errore combinato è matematicamente equivalente ad applicare GPTQ con un Essenziale combinato:
$$H_{comb} = \sum_{i} \alpha_i H_{\ell_i}$$
Ciò contrasta con la Per-Layer Independent Calibration (PLIC), che minimizza l'errore utilizzando solo l'Essenziale di un singolo strato ($H_{\ell_a}$).

2.2 Algoritmo CLHA

Gli autori propongono CLHA (Cross-Layer Hessian Aggregation), un algoritmo di calibrazione a tre passaggi:

1. Raccolta (Collection): Per ogni sequenza di calibrazione, accumula gli Essenziali per strato e il conteggio dei token per ogni esperto attivo.
2. Aggregazione (Aggregation): Per ogni coppia di esperti condivisi $(\ell_a, \ell_b)$, calcola un Essenziale combinato. Il documento valuta due strategie di ponderazione:
   • Routing-Frequency Weighting ($\alpha = n_a / (n_a + n_b)$): Pesa gli Essenziali in base al numero effettivo di token ricevuti da ciascuno strato.
   • Equal Weighting ($\alpha = 0.5$): Agisce come un regolarizzatore, utile quando i dati di calibrazione sono limitati o le stime dell'Essenziale dello strato secondario sono rumorose.
3. Quantizzazione (Quantization): Applica gli aggiornamenti sequenziali per colonna standard di GPTQ utilizzando l'Essenziale aggregato $H_{comb}$.

2.3 CLHA-MP (Mixed Precision)

Riconoscendo che gli esperti condivisi sempre-attivi ricevono tutti i $N$ token (a differenza degli esperti routati che ne ricevono una frazione $K/E$), il documento nota che la loro sensibilità alla quantizzazione è $E/K$ volte superiore. CLHA-MP assegna $b+1$ bit a questi esperti condivisi mantenendo gli esperti routati a $b$ bit, comportando un overhead di archiviazione minimo (10–20% dei parametri degli esperti).

2.4 Correzione dell'Implementazione: Ambiguità degli Hook

Il documento identifica un critico problema di ingegneria: gli standard PyTorch forward hooks attaccati a un modulo ExpertFFN condiviso scattano per qualsiasi strato che chiami quel modulo. Ciò causa la miscelazione delle attivazioni di diversi strati in un singolo accumulatore di Essenziali, corrompendo l'attribuzione per strato.

• Soluzione: Gli autori propongono un meccanismo di callback consapevole dello strato (layer-aware) passato attraverso la firma MoELayer.forward(). Questo fornisce esplicitamente gli indici degli strati alla callback, garantendo che siano accumulati separatamente gli Essenziali per $H_a$ e $H_b$, anche quando il modulo di peso sottostante è condiviso.

3. Contributi Chiave

1. Proposizione Teorica: Ha formalizzato che l'Essenziale ottimale per i pesi condivisi è la somma pesata per frequenza di routing degli Essenziali per strato.
2. Algoritmo CLHA: Un'implementazione pratica della calibrazione dell'Essenziale combinato all'interno del framework GPTQ.
3. CLHA-MP: Un'estensione a precisione mista che assegna un bit extra agli esperti condivisi altamente sensibili.
4. Correzione dell'Implementazione: Ha identificato e risolto il problema dell' "ambiguità degli hook" che corrompe silenziosamente la calibrazione nelle architetture a moduli condivisi.
5. Validazione Empirica: Ha dimostrato che i guadagni derivanti da CLHA correlano con la divergenza del routing tra gli strati.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su un CLWS-MoE da 34.5M di parametri addestrato da zero su WikiText-2.

• Performance a 2-bit: Questo è il regime critico dove il malallineamento si aggrava.
  • CLHA-rt (routing-weighted) ha ridotto la Perplessità (PPL) del 4.3% rispetto a PLIC.
  • CLHA-eq (equal-weighted) ha ridotto la PPL del 5.4%.
  • CLHA-MP ha ottenuto un miglioramento del 18.6% rispetto a PLIC.
  • A 4-bit e 3-bit, le differenze rientravano nei margini di rumore (<1%).
• Calibrazione vs. Errore di Test: Curiosamente, PLIC ha ottenuto un errore dell'Essenziale di Calibrazione (CHE) più basso sul set di addestramento perché ottimizza direttamente per quella distribuzione. Tuttavia, CLHA è generalizzato meglio sul set di test, indicando che l'Esseniale combinato cattura la vera distribuzione del peso condiviso meglio di una visione a singolo strato.
• Esperimento sulla Nitidezza del Routing (Routing Sharpness): Rendendo artificialmente più netti i logit del router (aumentando la divergenza tra gli strati), gli autori hanno osservato che:
  • La performance di PLIC è degradata del 5% all'aumentare della divergenza.
  • CLHA è rimasto stabile (<1% di degradazione).
  • Il divario di performance tra CLHA e PLIC è cresciuto da +4.3% (routing naturale) a +8.0% (alta divergenza), confermando la previsione teorica che i guadagni di CLHA scalano con la divergenza cross-layer.

5. Significato e Rivendicazioni

Il documento sostiene di fornire la prima analisi PTQ specificamente per la condivisione dei pesi tra gli strati nei modelli MoE. La sua importanza risiede nel:

• Correzione dell'Obiettivo: Dimostrare che trattare i pesi condivisi come entità a singolo strato è matematicamente subottimale e che l'aggregazione dell'Essenziale è la correzione necessaria.
• Scalabilità dei Guadagni: Mostrare che man mano che i modelli diventano più grandi e i pattern di routing più specializzati (aumentando la divergenza cross-layer), i benefici di CLHA probabilmente cresceranno, rendendolo essenziale per i futi modelli CLWS-MoE di produzione (es. Megrez2, DeepSeek, Qwen).
• Necessità Pratica: Evidenziare che l'applicazione ingenua degli esistenti strumenti PTQ alle architetture CLWS è fondamentalmente errata a causa dell'ambiguità degli hook, fornendo la correzione necessaria per qualsiasi futuro metodo di calibrazione in questo dominio.

Gli autori rimangono modesti riguardo alle prestazioni assolute, notando che il loro modello è piccolo (34.5M) e addestrato per pochi step, ma sottolineano che i miglioramenti relativi rispetto alla calibrazione standard sono robusti e teoricamente fondati.