Riemannian Gradient Descent for Low-Rank Architectures

Sintesi Tecnica: Discesa del Gradiente Riemanniana per Architetture a Basso Rango

Enunciato del Problema
Nel deep learning moderno, i parametri del modello sono spesso matrici vincolate a un basso rango per ridurre il costo computazionale e il numero di parametri (ad esempio, nei meccanismi di multihead attention dove $W = AB^T$). L'addestramento standard aggiorna i fattori $A$ e $B$ indipendentemente utilizzando gradienti euclidei, ignorando la struttura geometrica dell'insieme delle matrici di rango-$r$. Questo articolo investiga se trattare la matrice a rango vincolato $W$ come un punto su una sottovarietà liscia e applicare l'ottimizzazione riemanniana possa portare a una migliore convergenza o a soluzioni migliori. La sfida centrale è che la fattorizzazione $W=AB^T$ non è unica (invariante sotto $A \to AS, B \to BS^{-T}$) e l'insieme delle matrici di rango-$r$ è una varietà, non uno spazio vettoriale. Gli autori mirano a limitare le modifiche solo all'ottimizzatore, evitando cambiamenti all'architettura del modello, mantenendo al contempo una complessità per iterazione comparabile ai metodi standard ($O(mr^2 + nr^2)$).

Metodologia
Gli autori propongono un framework di discesa del gradiente riemanniano adattato per parametri fattorizzati a basso rango. L'iterazione generale prevede:

1. Recuperare il gradiente euclideo ambiente $G_W$ dai gradienti dei fattori $G_A$ e $G_B$.
2. Proiettare questo gradiente sullo spazio tangente della specifica varietà per ottenere il gradiente riemanniano $G_W^R$.
3. Aggiornare il momentum e normalizzare la dimensione del passo (step size).
4. Muoversi lungo la geodetica (usando una ritrazione) verso l'iterata successiva.
5. Trasportare parallelamente il vettore del momentum nel nuovo spazio tangente.

Il documento istanzia questo template attraverso dieci specifici design algoritmici categorizzati per geometria e vincoli:

1. Geometrie per Matrici a Rango Fisso ($R_{r}^{m \times n}$):

   • Geometria Incorporata (Embedded): Utilizza la metrica di Frobenius ambiente. La ritrazione viene eseguita tramite SVD troncata di rango-$r$ di $W + \Xi$. Il trasporto parallelo utilizza la proiezione metrica.
   • Geometria del Quoziente: Tratta la varietà come un quoziente dello spazio prodotto dei fattori $(A, B)$. La metrica è definita "sopra" sui fattori, inducendo una metrica sul quoziente. La ritrazione è additiva sui fattori; il trasporto coinvolge la proiezione orizzontale.

2. Geometrie per Isometrie Parziali di Rango-$r$ ($P_{r}^{m \times n}$):

   • Qui, $W$ è vincolata in modo che i suoi valori singolari non nulli siano tutti pari a 1 (equivalente a $A$ e $B$ essere ortonormali nelle colonne).
   • Geometria Incorporata: Restringe la metrica di Frobenius. La ritrazione usa la SVD troncata con specifici vincoli di ortogonalità.
   • Geometria del Quoziente: Utilizza il prodotto delle varietà di Stiefel ($St_{m,r} \times St_{n,r}$) come spazio "sopra". Le ritrazioni usano decomposizioni QR o Polari.
   • Geometria Canonica: Utilizza la metrica canonica sui fattori di Stiefel, inducendo una metrica di prodotto sul quoziente.

3. Condivisione dei Pesi (Griglie di Varietà):

   • Estende quanto sopra ad architetture in cui i fattori sono condivisi tra blocchi di righe e blocchi di colonne (ad esempio, Group-Query Attention).
   • Per il Rango Fisso, gli autori dimostrano che impilare i fattori in matrici più grandi permette l'uso di algoritmi singleton, a patto che la preservazione del rango sia verificata per blocchi.
   • Per le Isometrie Parziali, l'impilamento semplice fallisce a causa dei vincoli di ortogonalità. Gli autori sviluppano algoritmi più complessi, incluso un approccio "ibrido" per il caso incorporato dove si elevano alla geometria del quoziente per eseguire ritrazioni e trasporti, per poi proiettare nuovamente verso il basso.

Contributi Chiave

• Esplorazione Sistematica: Il documento fornisce un sondaggio esaustivo e un'implementazione di dieci diverse varianti di ottimizzazione riemanniana per parametri matriciali a basso rango, coprendo sia i vincoli di rango fisso che di isometria parziale, nonché scenari di condivisione dei pesi.
• Dettagli di Implementazione: Dettaglia le specifiche ritrazioni (SVD, QR, Polar) e i trasporti vettoriali richiesti per ogni geometria, affrontando la non unicità delle fattorizzazioni e la preservazione dei vincoli della varietà (rango e ortogonalità).
• Valutazione Empirica: I metodi sono applicati alle matrici Q, K, V, O in modelli linguistici su piccola scala (6 strati, dimensione di embedding 512) utilizzando il dataset FineWeb.

Risultati Sperimentali

• Performance: Dopo la calibrazione dei tassi di apprendimento (learning rates), i metodi riemanniani non hanno superato in modo conclusivo il baseline AdamW. Nei setting di Multi-Head Attention (MHA) e Group-Query Attention (GQA), le differenze di perdita erano spesso entro la varianza dei seed casuali.
• Iperparametri: La normalizzazione della metrica ha fornito lievi miglioramenti (0.01–0.05 nats). Un momentum pesante ($\nu=7/8$) ha generalmente danneggiato le performance o fornito guadagni trascurabili rispetto al baseline.
• Stabilità: Alcune esecuzioni hanno mostrato picchi di perdita (ad esempio, intorno al passo 6660) che si sono poi recuperati, spesso correlati a dati strutturati nel set di addestramento.
• Complessità: Gli autori riconoscono che, sebbene i metodi siano matematicamente solidi, la loro complessità computazionale per iterazione è superiore ad AdamW, rendendoli meno attraenti per gli specifici esperimenti su piccola scala condotti.

Significatività e Rivendicazioni
Il documento mantiene una posizione modesta riguardo alla propria novità e al proprio impatto immediato:

• Riscoperta: Gli autori dichiarano esplicitamente che qualsiasi apparente novità è "quasi certamente una riscoperta", basandosi fortemente su testi consolidati (Absil et al., Boumal) e letteratura precedente.
• Motivazione Euristica: Il lavoro è motivato dall'euristica secondo cui un migliore account delle relazioni tra i parametri (tramite la geometria della varietà) potrebbe portare a una convergenza più rapida o a soluzioni migliori, sebbene ciò non sia stato dimostrato definitivamente negli esperimenti.
• Potenziale Futuro: Gli autori sostengono che la mancanza di superamento delle prestazioni nei modelli su piccola scala non esclude il merito per scale più ampie. Suggeriscono che l'approccio è matematicamente solido e che il "tradeoff costo-accuratezza" potrebbe spostarsi favorevolmente con modelli più grandi o architetture diverse (ad esempio, modelli tensoriali come MLA, MoLAE).
• Limitazioni: Lo studio ammette significative limitazioni, tra cui la piccola dimensione del modello, la mancanza di annealing del tasso di apprendimento e l'uso di un singolo baseline (AdamW) senza la calibrazione dei suoi iperparametri.

In conclusione, il documento funge da rigorosa dimostrazione tecnica dell'applicazione dell'ottimizzazione riemanniana ai componenti a basso rango del deep learning, fornendo una libreria di algoritmi e un baseline per la ricerca futura, pur riconoscendo che i risultati empirici attuali non giustificano ancora la sostituzione di ottimizzatori standard come AdamW.