GeoRA: Geometry-Aware Low-Rank Adaptation for RLVR

1. Il Problema

L'articolo affronta le sfide specifiche dell'adattamento di modelli linguistici di grandi dimensioni (LLM) tramite Reinforcement Learning con Ricompense Verificabili (RLVR), un paradigma chiave per migliorare le capacità di ragionamento (es. OpenAI-o1, DeepSeek-R1).

Sebbene l'RLVR sia efficace, presenta dinamiche di ottimizzazione distinte rispetto al Supervised Fine-Tuning (SFT):

• Dinamiche Geometriche Diverse: Mentre l'SFT tende a modificare le direzioni principali dei pesi (componenti a grande magnitudine), l'RLVR favorisce aggiornamenti in sottospazi "off-the-principals", ovvero direzioni ortogonali o a bassa magnitudine rispetto alle caratteristiche pre-addestrate, per preservare la struttura geometrica originale e evitare il collasso del comportamento.
• Limiti dei Metodi Esistenti:
  • I metodi di Low-Rank Adaptation (LoRA) orientati all'SFT (come PiSSA) inizializzano gli adattatori sulle componenti principali (valori singolari maggiori) della matrice dei pesi. Questo crea un mismatch geometrico con l'RLVR, che preferisce aggiornare regioni diverse, portando a instabilità o prestazioni subottimali.
  • I metodi di Sparse Fine-Tuning (adattamenti su parametri sparsi non strutturati) sono teoricamente più allineati all'RLVR, ma soffrono di colli di bottiglia di efficienza sull'hardware moderno (GPU/TPU) a causa della scarsa supporto per la sparsità non strutturata, rendendoli lenti in pratica.

2. Metodologia: GeoRA

Gli autori propongono GeoRA (Geometry-Aware Low-Rank Adaptation), un framework di adattamento a basso rango progettato specificamente per la geometria dell'RLVR. La metodologia si articola in tre fasi principali:

A. Costruzione del Sottospazio Vincolato dalla Geometria

Invece di lavorare sulla matrice dei pesi originale $W$, GeoRA costruisce una matrice vincolata $W_{Geo}$ identificando le regioni più stabili e plastiche per l'aggiornamento RL. Questo avviene tramite una strategia di mascheramento che combina due priors geometrici:

1. Prior Spettrale ($M_{Spec}$): Seleziona le componenti a bassa curvatura (bassa magnitudine) all'interno di una approssimazione a basso rango dei pesi, sopprimendo le direzioni ad alta energia che potrebbero destabilizzare il modello.
2. Prior Euclideo ($M_{Euc}$): Seleziona i pesi originali a bassa magnitudine, sfruttando la plasticità dei parametri vicini allo zero.
   La matrice $W_{Geo}$ è l'unione di queste due regioni stabili.

B. Inizializzazione Strutturata tramite SVD

Sulla matrice $W_{Geo}$ viene eseguita una Scomposizione ai Valori Singolari (SVD).

• Gli adattatori a basso rango ($A_{Geo}$ e $B_{Geo}$) vengono inizializzati per approssimare le direzioni principali all'interno di questo sottospazio geometrico vincolato, non su tutto lo spettro dei pesi originali.
• Questo garantisce che l'adattamento parta già allineato con la geometria preferita dall'RLVR.

C. Ancoraggio Strutturale (Residual Leash)

Per preservare le capacità pre-addestrate e garantire la stabilità:

• Viene calcolata una Matrice Residua ($W_{res}$) sottraendo l'adattatore inizializzato dai pesi originali: $W_{res} = W - \frac{\alpha}{r} B_{Geo}A_{Geo}$.
• Durante l'addestramento, $W_{res}$ viene congelato (frozen).
• Il forward pass diventa: $h = W_{res}x + \frac{\alpha}{r} B_{Geo}A_{Geo}x$.
  Questo meccanismo agisce come un "ancoraggio" che impedisce all'ottimizzatore di erodere le rappresentazioni pre-addestrate, permettendo aggiornamenti solo lungo le direzioni geometricamente allineate.

3. Contributi Chiave

1. Primo Framework Geometry-Aware per RLVR: GeoRA è il primo metodo PEFT (Parameter-Efficient Fine-Tuning) che esplicitamente allinea l'inizializzazione e la struttura degli adattatori alla geometria specifica degli aggiornamenti RLVR, superando il mismatch dei metodi basati sull'SFT.
2. Scoperta della Struttura Compressibile: Gli autori dimostrano che il sottospazio di aggiornamento efficace per l'RLVR è anisotropo e ammette una struttura a basso rango compressibile, non isotropo come ipotizzato da alcuni approcci sparsi casuali.
3. Efficienza Hardware: A differenza dei metodi sparsi non strutturati, GeoRA utilizza aggiornamenti densi a basso rango, garantendo velocità di calcolo elevate sulle GPU moderne senza sacrificare l'efficienza dei parametri.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli Qwen e Llama (da 1.5B a 32B parametri) in domini di ragionamento matematico, medico e di coding.

• Prestazioni In-Distribution (Matematica): GeoRA supera costantemente i baselines forti (LoRA, PiSSA, MiLoRA, SparseFT) su benchmark complessi come AIME24/25 e OlymMATH. Ad esempio, su Qwen3-8B, GeoRA ha raggiunto il 23.75% su AIME24 contro il 22.50% di PiSSA e il 19.58% di LoRA.
• Generalizzazione Out-of-Distribution (OOD): GeoRA mostra una capacità superiore di preservare le conoscenze generali su task non correlati (Coding, Scienze, Istruzioni), riducendo il "dimenticamento" (catastrophic forgetting) rispetto ad altri metodi.
• Stabilità dell'Addestramento:
  • Robustezza agli Iperparametri: GeoRA mantiene alte prestazioni su un ampio range di learning rate, mentre PiSSA e MiLoRA collassano con learning rate più alti.
  • Dinamica KL: Durante l'RLVR, GeoRA mantiene una divergenza KL bassa e controllata, evitando il collasso della politica che si osserva in altri metodi con aggiornamenti aggressivi.
• Efficienza: Rispetto al Full Fine-Tuning, GeoRA riduce i parametri addestrabili del 99.5% e il consumo di VRAM del 28.5%, superando anche i metodi sparsi in termini di velocità di iterazione grazie all'uso di operazioni dense.

5. Significato e Impatto

Il lavoro di GeoRA è significativo perché colma il divario teorico e pratico tra l'efficienza dei parametri e le specifiche esigenze geometriche dell'RLVR.

• Teorico: Dimostra che l'RLVR non è un processo di ottimizzazione isotropo, ma favorisce sottospazi strutturati e compressibili che possono essere catturati tramite SVD su regioni specifiche dei pesi.
• Pratico: Fornisce una soluzione pronta all'uso che combina la stabilità dell'SFT (tramite l'ancoraggio residuo) con l'efficacia dell'RLVR, permettendo di addestrare modelli di ragionamento complessi su hardware standard senza i costi proibitivi del Full Fine-Tuning o le inefficienze della sparsità non strutturata.

In sintesi, GeoRA rappresenta un passo avanti fondamentale per scalare l'addestramento RLVR su modelli di grandi dimensioni, offrendo un compromesso ottimale tra stabilità, generalizzazione ed efficienza computazionale.