GeoRA: Geometry-Aware Low-Rank Adaptation for RLVR

Probleemstelling

Reinforcement Learning met Verifieerbare Beloningen (RLVR) is een cruciale paradigma voor het verbeteren van redeneervermogens in grote taalmodellen (zoals OpenAI-o1 en DeepSeek-R1). In tegenstelling tot Supervised Fine-Tuning (SFT), waarbij kennis wordt ingebracht door de belangrijkste gewichtsrichtingen aan te passen, werkt RLVR als een beperkt optimalisatieproces dat latente redeneergedragingen versterkt via beloningsgeïnduceerde sampling bias.

Bestaande parameter-efficiënte fijnafstemmingstechnieken (PEFT) hebben echter moeite met deze specifieke dynamiek:

1. Geometrisch misverstand: Methodes zoals PiSSA en LoRA zijn ontworpen voor SFT en initialiseren adapters op basis van de belangrijkste singuliere componenten (principale componenten). Dit botst met de RLVR-dynamiek, die updates juist "off the principals" (orthogonaal op de pre-getrainde structuur) prefereert om het bestaande redeneervermogen niet te destabiliseren.
2. Efficiëntieproblemen: Methodes die wel proberen de voorkeur van RLVR voor ongestructureerde sparsiteit na te bootsen, lopen vast op hardware-efficiëntie. Moderne GPU's ondersteunen ongestructureerde sparsiteit slecht, wat leidt tot geen echte snelheidswinst of zelfs extra overhead.

Er is dus behoefte aan een methode die de specifieke geometrie van RLVR-updates respecteert, de pre-getrainde structuur behoudt, en toch hardware-efficiënt is.

Methodologie: GeoRA

De auteurs stellen GeoRA (Geometry-Aware Low-Rank Adaptation) voor, een raamwerk dat specifiek is ontworpen voor RLVR. De kern van de methode bestaat uit drie fasen:

1. Constructie van een Geometrisch Beperkte Subruimte ($W_{Geo}$):
   In plaats van het volledige gewichtsmatrix $W$ te gebruiken, construeren de auteurs een masker om een subruimte te isoleren die stabiel en plastisch is voor RLVR-updates. Dit gebeurt via een unie van twee priors:

   • Spectrale Prior ($M_{Spec}$): Selecteert de componenten met de laagste magnitude binnen een rang-$r$ benadering van $W$. Dit onderdrukt hoge-energie (hoge kromming) richtingen die instabiliteit kunnen veroorzaken.
   • Euclidische Prior ($M_{Euc}$): Selecteert de oorspronkelijke gewichten met de laagste magnitude. Dit vangt parameters op die plastisch zijn (gemakkelijk aan te passen).
     De resulterende matrix is $W_{Geo} = W \odot (M_{Spec} \cup M_{Euc})$.

2. Geometrie-bewuste Initialisatie via SVD:
   Op $W_{Geo}$ wordt een Singular Value Decomposition (SVD) uitgevoerd. De adapters $A$ en $B$ worden niet willekeurig geïnitieerd, maar worden afgeleid van de top-$r$ singuliere componenten van $W_{Geo}$. Hierdoor worden de trainbare richtingen direct uitgelijnd met de meest effectieve, gestructureerde subruimte voor RLVR.

3. Residuele Anker (Structural Anchor):
   Om te garanderen dat de modeloutput bij initialisatie ongewijzigd blijft en de kerncapaciteiten behouden blijven, wordt een Residuele Matrix ($W_{res}$) berekend:
   $$W_{res} = W - \frac{\alpha}{r} B_{Geo}A_{Geo}$$
   Tijdens het trainen wordt $W_{res}$ bevroren. De forward pass wordt dan:
   $$h = W_{res}x + \frac{\alpha}{r} B_{Geo}A_{Geo}x$$
   Dit zorgt ervoor dat de optimizer alleen de geometrisch uitgelijnde laag-rang manifold kan aanpassen, terwijl $W_{res}$ fungeert als een stabiel anker dat de pre-getrainde representaties beschermt tegen erosie.

Belangrijkste Bijdragen

• GeoRA Framework: Het eerste geometrie-bewuste low-rank adaptatiekader dat expliciet is ontworpen voor RLVR, in plaats van voor SFT.
• Inzicht in RLVR-geometrie: Het paper demonstreert dat de effectieve update-subruimte voor RLVR anisotroop en comprimeerbaar is (gestructureerd low-rank), in plaats van isotroop of willekeurig spars.
• Hybride Architectuur: GeoRA combineert de voordelen van low-rank adaptatie (efficiëntie) met een bevroren residueel anker (stabiliteit), waardoor het de geometrische mismatch van bestaande methodes en de efficiëntieproblemen van ongestructureerde sparsiteit oplost.

Resultaten

GeoRA is geëvalueerd op modellen van 1,5B tot 32B parameters (Qwen en Llama families) in domeinen zoals wiskunde, geneeskunde en codering.

• In-Distribution (Wiskunde): GeoRA presteert consistent beter dan sterke baselines (LoRA, PiSSA, MiLoRA, SparseFT) op uitdagende benchmarks zoals AIME24/25, MATH-500 en OlymMATH. Het bereikt vaak resultaten die dicht bij Full Fine-Tuning liggen, maar met veel minder parameters.
• Out-of-Distribution (Generalisatie): GeoRA toont superieure generalisatie op taken buiten het traindomein (zoals HumanEval voor codering, GPQA voor wetenschap en TruthfulQA). Het vertoont minder "catastrophic forgetting" van algemene capaciteiten dan andere methodes.
• Trainingsstabiliteit: GeoRA bereikt sneller hoge prestaties en vertoont een stabielere trainingsdynamiek, zelfs bij agressieve leergraden. Andere methodes (zoals PiSSA) tonen vaker instabiliteit of een vroege plateau.
• Efficiëntie: GeoRA reduceert het aantal trainbare parameters met 99,5% ten opzichte van Full Fine-Tuning. In tegenstelling tot SparseFT, behoudt GeoRA dichte matrixbewerkingen, wat leidt tot hogere VRAM-efficiëntie en snellere trainingstijden op moderne hardware.

Betekenis en Conclusie

GeoRA sluit de kloof tussen parameter-efficiënte fijnafstemming en de unieke optimalisatiegeometrie van RLVR. Het paper toont aan dat het simpelweg toepassen van SFT-georiënteerde initiatie strategieën (zoals PiSSA) suboptimaal is voor RLVR. Door de update-subruimte te construeren op basis van de inherente geometrie van RLVR-updates en deze te combineren met een structureel anker, biedt GeoRA een stabiele, efficiënte en krachtige strategie voor het trainen van grote redeneermodellen. Dit maakt het een belangrijke stap voorwaarts voor het schalen van RLVR in praktische toepassingen.