Riepilogo Tecnico: KLAS - Utilizzo della Similarità per Cucire Reti Neurali per Migliori Compromessi tra Accuratezza ed Efficienza

1. Enunciato del Problema

Il dispiegamento di grandi transformer visivi (ViT) preaddestrati e reti neurali convoluzionali (CNN) affronta una sfida fondamentale: ambienti diversificati, che vanno da server cloud potenti a dispositivi edge con risorse limitate, richiedono una selezione flessibile dei modelli all'interno di budget computazionali specifici. Sebbene la Ricerca di Architettura Neurale (NAS) tenti di bilanciare accuratezza ed efficienza, la NAS tradizionale uno-a-uno è computazionalmente proibitiva, e gli approcci uno-a-molti o zero-shot sono spesso limitati a un singolo spazio di progettazione del modello, fallendo nel sfruttare il ricco pool di modelli preaddestrati esistenti.

La cucitura dei modelli (model stitching) offre un'alternativa promettente combinando blocchi da diversi modelli "ancora" preaddestrati per creare reti interpolate. Tuttavia, gli approcci esistenti di cucitura molti-a-molti, come SN-Net, si basano su euristiche naive (ad esempio, "cucitura più vicina" di ancoraggi adiacenti o accoppiamento fisso di blocchi) per selezionare le configurazioni di cucitura. Queste euristiche spesso producono compromessi subottimali tra accuratezza ed efficienza e mancano di generalizzabilità attraverso diverse famiglie di modelli. Inoltre, lo spazio di ricerca per le configurazioni di cucitura è vasto ($O(k^2n^2)$ per $k$ modelli di profondità $n$), rendendo la ricerca esaustiva non fattibile. Il problema centrale è la mancanza di un metodo principiato e automatizzato per identificare configurazioni di cucitura compatibili che massimizzino le prestazioni senza incorrere in elevati costi di addestramento.

2. Metodologia: Il Framework KLAS

Gli autori propongono KLAS (Cucitura di Ancoraggi basata sulla Divergenza KL), un framework che automatizza la selezione della cucitura sfruttando la similarità tra le rappresentazioni intermedie dei modelli preaddestrati.

Concetto Fondamentale: La Divergenza KL come Metrica di Similarità

KLAS postula che una cucitura di successo dipenda sia dalla similarità rappresentazionale (allineamento geometrico delle attivazioni) sia dalla similarità funzionale (preservazione delle informazioni rilevanti per il compito e dei confini decisionali). Gli autori sostengono che la divergenza di Kullback-Leibler (KL) sia unica nel catturare entrambi:

• Allineamento Rappresentazionale: Quantifica la distanza statistica tra le distribuzioni di attivazione, indicando se una trasformazione leggera può mappare le attivazioni sorgente su quelle target.
• Allineamento Funzionale: Valuta quanto bene le uscite di un blocco sorgente preservino i confini decisionali del modello target, garantendo che la rete cucita mantenga le prestazioni sul compito a valle con un minimo fine-tuning.

A differenza di metriche come l'Allineamento del Kernel Centrato (CKA), l'Errore Quadratico Medio (MSE) o l'Entropia Incrociata (CE), che spesso falliscono nel prevedere il successo della cucitura, la divergenza KL misura direttamente la compatibilità delle distribuzioni di probabilità generate da sonde lineari.

L'Algoritmo KLAS

Il framework opera in due fasi principali:

1. Selezione degli Ancoraggi:

   • KLAS calcola la divergenza KL tra le distribuzioni di probabilità softmax degli ultimi blocchi dei modelli preaddestrati candidati.
   • I modelli con bassa divergenza KL dell'ultimo blocco sono considerati ancoraggi compatibili, poiché condividono confini decisionali e distribuzioni di confidenza di classe simili. Questo sostituisce la "cucitura più vicina" euristica utilizzata nei lavori precedenti.

2. Selezione dei Blocchi:

   • Una volta selezionati gli ancoraggi, KLAS valuta le coppie di blocchi utilizzando un punteggio di cucitura $\Gamma(i, j)$ definito come:
     $$\Gamma(i, j) = \frac{\Theta(P^f_i, P^g_j)}{\Theta(P^g_j, P^g_{j+1})}$$
     Dove:
     • $\Theta(P^f_i, P^g_j)$ è la distanza di attivazione tra ancoraggi (divergenza KL tra il blocco sorgente $i$ e il blocco target $j$).
     • $\Theta(P^g_j, P^g_{j+1})$ è la capacità intra-blocco dell'ancoraggio (divergenza KL tra blocchi consecutivi nell'ancoraggio target).
   • Questo punteggio bilancia il carico di trasformazione sul ipotetico strato di cucitura rispetto alla capacità del blocco target di assorbire le differenze distribuzionali.
   • ProbeNet: Per calcolare efficientemente questi punteggi senza addestrare strati di cucitura completi, KLAS impiega un'architettura unificata chiamata ProbeNet. Essa attacca sonde di classificazione lineare a ogni blocco dei modelli ancoraggio e le addestra congiuntamente in un singolo passaggio forward-backward. Ciò comporta un costo one-time trascurabile (ad esempio, 0,25 giorni GPU per Swin-B) rispetto all'addestramento di sonde individuali.

3. Selezione della Configurazione:

   • Le configurazioni di cucitura candidate vengono filtrate in base a una soglia regolabile $\tau$ rispetto al punteggio minimo all'interno dei bucket FLOPs. Ciò garantisce una fronte di Pareto densa di modelli dispiegabili attraverso lo spettro accuratezza-efficienza.

3. Contributi Chiave

1. Identificazione dei Limiti delle Euristiche: Il documento dimostra che le euristiche naive (come quelle in SN-Net) non sono universalmente ottimali, spesso mancando configurazioni di "cucitura lontana" superiori (ad esempio, collegare direttamente Swin-Ti a Swin-B) che le euristiche ignorano.
2. Metrica di Similarità Principiata: L'introduzione della divergenza KL come metrica per la selezione della cucitura che cattura sia l'allineamento rappresentazionale sia quello funzionale senza richiedere l'istanziazione o l'addestramento degli strati di cucitura.
3. Framework KLAS: Un framework generalizzabile e a basso costo che automatizza la selezione di ancoraggi e blocchi attraverso diverse famiglie di modelli (ViT, CNN e LLM).
4. Validazione Empirica: Esperimenti completi che mostrano miglioramenti coerenti rispetto alle baseline su molteplici dataset (ImageNet-1K, CIFAR-100, ADE20K, TruthfulQA) e architetture.

4. Risultati Sperimentali

Gli autori hanno valutato KLAS rispetto a SN-Net e altre baseline in condizioni di fine-tuning identiche (50 epoche per la visione, 10 epoche per gli LLM).

• Compromessi Accuratezza-Efficienza:

  • ImageNet-1K: KLAS raggiunge fino al 1,21% in più di accuratezza Top-1 allo stesso costo computazionale o una riduzione di 1,33× dei FLOPs a parità di accuratezza rispetto a SN-Net.
  • Punteggi AUC: Sulla famiglia Swin, KLAS ha ottenuto un'Area sotto la Curva (AUC) di 0,8950, superando significativamente SN-Net (0,8345) e altre metriche di similarità come CKA (0,8124).
  • Cross-Architettura: KLAS ha cucito con successo ResNet-18 con Swin-B, superando le baseline di inferenza a cascata del 2-3% di accuratezza Top-1 a FLOPs simili.

• Predizione Densa (Segmentazione Semantica):

  • Su ADE20K utilizzando Mask2Former, le cuciture selezionate da KLAS hanno raggiunto fino a +0,9% mIoU rispetto alle baseline SN-Net a budget FLOPs comparabili.

• Grandi Modelli Linguistici (LLM):

  • Utilizzando ancoraggi Llama da 1B e 3B sul benchmark TruthfulQA, i modelli cuciti con KLAS hanno superato ESTA (un metodo di cucitura LLM basato su euristiche) fino a 0,017 ROUGE-1 e 0,033 ROUGE-2, nonostante avessero meno parametri in alcune configurazioni.

• Efficacia della Metrica:

  • La divergenza KL ha recuperato con successo l'88,9% delle configurazioni di cucitura Swin Ti-B vitali, mentre CKA ne ha recuperato solo il 5,5% e le euristiche SN-Net il 61,1%.

5. Significato e Affermazioni

Il documento afferma che KLAS affronta il divario critico nella cucitura dei modelli sostituendo le euristiche ad hoc con un approccio principiato e guidato dalla similarità. Il suo significato risiede in:

• Generalizzabilità: Funziona attraverso diverse famiglie architetturali (Transformer, CNN, Ibridi) e modalità (Visione, Linguaggio), superando i limiti dei metodi ristretti a singole famiglie.
• Efficienza dei Costi: Genera una fronte di Pareto continua di modelli dispiegabili con un sovraccarico computazionale aggiuntivo trascurabile (solo addestramento delle sonde e fine-tuning dello strato di cucitura), evitando i costi proibitivi della NAS tradizionale.
• Ottimizzazione Sistematica: Catturando esplicitamente le differenze distribuzionali tramite la divergenza KL, KLAS permette la scoperta di configurazioni di cucitura non intuitive e ad alte prestazioni (ad esempio, "cucitura lontana") che i metodi euristici mancano.

Gli autori concludono che KLAS fornisce una soluzione robusta e automatizzata per ottimizzare il compromesso tra accuratezza ed efficienza negli scenari di dispiegamento, sfruttando efficacemente l'ecosistema esistente di modelli preaddestrati senza richiedere l'addestramento da zero.