Riepilogo Tecnico: kNN-MoE

Enunciato del Problema

Le architetture Mixture-of-Experts (MoE) scalano i Large Language Models (LLM) in modo efficiente utilizzando un "router" parametrico per instradare i token verso un sottoinsieme sparso di esperti. Tuttavia, i router standard sono classificatori leggeri addestrati una sola volta e poi congelati. Sebbene efficaci sui dati in distribuzione, mancano di flessibilità durante l'inferenza quando si affrontano spostamenti di distribuzione (ad esempio, domini specializzati o task fuori distribuzione). In tali scenari, il router deve estrapolare da prior fissi, portando spesso a una selezione subottimale degli esperti e a un'alta perplessità. I metodi esistenti di adattamento al momento dell'inferenza, come C3PO, tentano di affinare i logit di routing online ma comportano un significativo sovraccarico computazionale (circa 5 volte la baseline). Altri approcci richiedono aggiornamenti permanenti dei parametri tramite fine-tuning supervisionato (SFT), che è computazionalmente intensivo e meno efficiente per l'inferenza.

Metodologia: kNN-MoE

Gli autori propongono kNN-MoE, un framework di routing potenziato dal recupero che affina le decisioni di assegnazione degli esperti senza modificare i parametri del modello. Il metodo opera in due fasi distinte:

1. Costruzione della Memoria (Offline)

Il sistema costruisce un archivio di memoria non parametrico $\mathcal{M}$ per ogni layer MoE, costituito da coppie chiave-valore $(k_t, v_t)$:

• Chiave ($k_t$): L'input del router (stato nascosto) $x_t$ da un dataset di riferimento $D_{ref}$.
• Valore ($v_t$): Un'assegnazione "localmente ottimale" degli esperti derivata ottimizzando i logit di routing per massimizzare la probabilità del token ground-truth $y_t$.
  • Invece di memorizzare l'output del router congelato, il metodo risolve un problema di ottimizzazione per ogni token in $D_{ref}$: trovare un vettore di logit $r^*_t$ che minimizzi la log-verosimiglianza negativa del token target, dati i parametri della rete congelata.
  • Questa ottimizzazione viene eseguita tramite pochi passi di discesa del gradiente (tipicamente $S=1$) partendo dai logit parametrici originali.
  • I logit ottimizzati risultanti $r^{(S)}_t$ servono come valore $v_t$ nella memoria.

2. Miscelazione Adattiva Consapevole della Confidenza (Inferenza)

Durante l'inferenza, per un dato input del router $x$:

• Recupero: Il sistema recupera i $K$ vicini più prossimi $\mathcal{N}(x)$ dalla memoria $\mathcal{M}$ basandosi sulla somiglianza tra $x$ e le chiavi memorizzate.
• Proposta Non Parametrica: Viene calcolata una proposta di logit basata sulla memoria $r_{mem}(x)$ aggregando i logit ottimali recuperati, ponderati in base alla loro somiglianza con l'input.
• Miscelazione Guidata dalla Confidenza: I logit di routing finali sono un'interpolazione lineare tra i logit del router parametrico originale $r(x)$ e la proposta basata sulla memoria $r_{mem}(x)$:
  $$r_{final}(x) = (1 - \lambda(x))r(x) + \lambda(x)r_{mem}(x)$$
  Il coefficiente di miscelazione $\lambda(x)$ è definito come la somiglianza media dei vicini recuperati. Questo agisce come un punteggio di confidenza:
  • Alta Somiglianza ($\lambda \approx 1$): Il modello si fida dell'assegnazione localmente ottimale recuperata.
  • Bassa Somiglianza ($\lambda \approx 0$): Il modello ricade sul router congelato originale, evitando il rumore da casi irrilevanti.
• Output: I logit finali vengono passati attraverso un Top-K Softmax per generare l'assegnazione degli esperti sparsi, mantenendo i vincoli di sparsità originali del modello.

Contributi Chiave

1. Routing Potenziato dal Recupero: Introduce un framework che riutilizza assegnazioni di esperti localmente ottimali calcolate offline da una memoria di casi passati simili per migliorare le decisioni di routing.
2. Adattamento Non Parametrico: A differenza dei metodi SFT o di adattamento online, kNN-MoE non aggiorna i parametri del modello. Sposta il carico computazionale su una fase offline di costruzione della memoria, consentendo un'inferenza efficiente.
3. Fallback Guidato dalla Confidenza: Il metodo utilizza la somiglianza media dei vicini recuperati come coefficiente di miscelazione dinamico, garantendo robustezza tornando al router congelato quando non vengono trovati casi storici rilevanti.
4. Efficienza: L'approccio richiede un singolo passo di discesa del gradiente per token durante la costruzione della memoria ed evita i pesanti costi computazionali dell'ottimizzazione online o del fine-tuning completo del modello.

Risultati Sperimentali

Gli autori hanno valutato kNN-MoE su tre modelli MoE (OLMoE, GPT-OSS e Qwen3) attraverso sei benchmark (GPQA, SuperGPQA, MMLU, USMLE, MedMCQA e MBPP).

• Prestazioni vs Baseline: kNN-MoE ha costantemente superato la baseline Zero-shot e la baseline "SFT (Router Only)" su tutti i modelli e dataset.
• Competitività con SFT: Il metodo ha raggiunto prestazioni competitive con il pieno Supervised Fine-Tuning (SFT) e, in diversi scenari con scarsità di dati (ad esempio GPQA, USMLE), ha superato l'SFT, che occasionalmente ha sofferto di overfitting su piccoli set di riferimento.
• Stabilità: A differenza del prompting a 5-shot, che ha degradato le prestazioni su certi modelli (GPT-OSS, OLMoE) a causa di problemi di lunghezza dell'input, kNN-MoE ha mostrato miglioramenti stabili.
• Studi di Ablazione:
  • Vicini ($K$): Contrariamente alle tipiche scoperte di kNN-LM, le prestazioni hanno raggiunto il picco a $K=1$. L'aggregazione di più vicini ($K>1$) ha introdotto rumore, suggerendo che i segnali di routing realmente benefici sono rari e specifici.
  • Passi di Ottimizzazione ($S$): L'aumento dei passi di discesa del gradiente oltre $S=1$ ha prodotto guadagni di accuratezza trascurabili aumentando i costi di costruzione.
  • Dimensione del Set di Riferimento: Le prestazioni sono migliorate in modo monotono con la dimensione del set di riferimento $|D_{ref}|$, con rendimenti decrescenti osservati oltre 0,5k–1k esempi.
• Efficienza: La costruzione della memoria ha richiesto 0,46 ore per 1k esempi (vs ~1,34h per SFT). La variante IVFPQ (Inverted File + Product Quantization) ha ridotto il sovraccarico di latenza nell'inferenza a ~3,67% e l'aumento dell'uso di VRAM a ~9,43% rispetto alla baseline.

Significato e Affermazioni

Il documento afferma che kNN-MoE colma il divario tra il routing parametrico statico e le esigenze di inferenza dinamica. Il suo significato principale risiede nell'offrire una via scalabile per migliorare l'adattabilità delle MoE senza il carico computazionale del fine-tuning.

Gli autori sottolineano che il metodo è particolarmente efficace quando il router parametrico è incerto (correlato ad alta perplessità), poiché il meccanismo di recupero fornisce un segnale correttivo da casi storici simili e di alta qualità. Sebbene il metodo dipenda dalla disponibilità di un set di riferimento etichettato con una distribuzione simile ai dati di test, offre un'alternativa pratica e non parametrica alle tecniche di adattamento costose, raggiungendo risultati competitivi con l'SFT con un tempo di preparazione significativamente inferiore e senza aggiornamenti dei parametri.