Apprendere sottospazi funzionali per comprimere le reti neurali

Abstract degli autori

I transformer moderni offrono capacità notevoli, ma richiedono molta memoria e potenza di calcolo. La fattorizzazione dei pesi a rango ridotto limita entrambe le esigenze mantenendo dense le matrici, che restano così efficienti sull’hardware standard. I metodi esistenti, però, scelgono il sottospazio da eliminare da ogni matrice dei pesi usando criteri locali in forma chiusa: l’energia delle attivazioni, l’errore di ricostruzione del singolo livello o un’approssimazione quadratica della funzione di perdita. Questi criteri ignorano il modo in cui gli errori si propagano nella rete: a livelli elevati di compressione, gli errori si accumulano con la profondità e le prestazioni crollano. Presentiamo Learnable Subspace Projections (LSP), un metodo che apprende invece da un’estremità all’altra quali sottospazi eliminare. A ogni livello lineare, o gruppo vincolato di livelli che leggono le stesse attivazioni, viene assegnato un proiettore ortogonale. Tutti i proiettori vengono ottimizzati congiuntamente rispetto a un obiettivo globale — la divergenza KL rispetto alla distribuzione delle uscite del modello denso oppure la funzione di perdita usata originariamente per addestrare il modello — mentre i pesi preaddestrati restano congelati. I proiettori vengono inizializzati mediante il troncamento di una SVD sbiancata e i ranghi vengono assegnati in base alla divergenza KL delle uscite indotta da ciascun proiettore per ogni parametro risparmiato. Dopo l’addestramento, i proiettori vengono incorporati in fattori standard a rango ridotto, con un fattore condiviso da ciascun gruppo vincolato. Nell’attenzione, questo consente anche al modello di memorizzare nella cache un’unica rappresentazione latente di dimensioni ridotte al posto di chiavi e valori completi. Su diversi LLM (OPT-125M/1.3B, Qwen3-4B, Llama-2-7B) e su ViT-B/16, LSP supera i metodi di riferimento e il suo vantaggio cresce all’aumentare della compressione. Con una compressione del -70%, LSP porta Llama-2-7B a una perplexity di 10,9 su WikiText-2 e a un’accuratezza media zero-shot del 42,2%, contro rispettivamente 13,3 e 36,0% per il miglior metodo di riferimento. Il modello fattorizzato esegue la decodifica fino a 1,6 volte più velocemente del modello denso con batch di piccole dimensioni; memorizzare nella cache la rappresentazione latente condivisa riduce inoltre di 13,5 volte la memoria complessiva occupata dai pesi e dalla cache KV con un contesto di 128k token, contro una riduzione massima di 6,5 volte per le fattorizzazioni di riferimento senza fattori condivisi.