Preference-Aligned LoRA Merging: Preserving Subspace Coverage and Addressing Directional Anisotropy

1. Het Probleem

Het samenvoegen (merging) van meerdere Low-Rank Adaptation (LoRA) modules is een veelbelovende techniek om algemene systemen te bouwen zonder kostbare multi-task training. Echter, naadloze samenvoeging is uitdagend omdat LoRA-updates voor verschillende taken vaak in verschillende subruimtes (subspaces) plaatsvinden en ongelijk bijdragen aan de prestaties.

Wanneer LoRA-modules naïef worden samengevoegd (bijvoorbeeld door lineaire combinatie van gewichten), treden twee fundamentele problemen op:

1. Verlies van Subspace Coverage: De diverse representatieve richtingen van de verschillende taken overlappen of interfereren destructief, waardoor het samengevoegde model minder in staat is om alle taken even goed te vertegenwoordigen.
2. Directionele Anisotropie: Er is een onbalans in de invloed van verschillende LoRA-richtingen. Sommige richtingen zijn extreem gevoelig voor veranderingen in de loss-functie, terwijl andere dat niet zijn. Als deze onbalans wordt genegeerd, worden minder belangrijke richtingen overgewaardeerd en kritieke richtingen ondergewaardeerd, wat leidt tot een verslechtering van de taakprestaties.

Bestaande methoden (zoals Task Arithmetic, TIES-Merging, of AdaMerging) lossen vaak slechts één van deze problemen op, of negeren de specifieke structuur van LoRA volledig, wat resulteert in suboptimale prestaties.

2. Methodologie: TARA-Merging

De auteurs stellen TARA (Task-Rank Anisotropy Alignment) voor, een raamwerk dat beide problemen aanpakt door de voorkeuren van de gebruiker te integreren en de subruimte-structuur te behouden.

Kerncomponenten van TARA:

• Subspace Coverage Behoud:
  In plaats van het samenvoegen van volledige update-matrices ($\Delta W$), behoudt TARA de interne rang-1 structuur van de LoRA-adapters. Het analyseert de effectieve rang (effective rank) van de verzameling LoRA-richtingen om te garanderen dat de samenvoeging een breed spectrum aan representatieve richtingen behoudt.

  • Variant A: Herweegt individuele rang-1 factoren binnen elke adapter om anisotropie te controleren.
  • Variant B: Construeert een gedeelde orthonormale basis over alle taken via SVD (Singular Value Decomposition) om cross-task interferentie te minimaliseren en de subruimte-coverage expliciet te behouden.

• Anisotropie en Voorkeursuitlijning:
  TARA erkent dat de gevoeligheid van de loss-functie voor bepaalde richtingen afhangt van de taakvoorkeuren.

  • Het gebruikt een smooth Tchebycheff scalarisatie met een voorkeursvector $\rho$ (die aangeeft welke taken belangrijker zijn).
  • Het optimaliseert richting-gewichten ($\phi$) door een pseudo-loss te minimaliseren gebaseerd op de voorspellende entropie (vergelijkbaar met AdaMerging), maar dan specifiek gericht op het uitlijnen van de richting-gewichten met de taakvoorkeuren.
  • Dit zorgt voor een richtingsgewijze herweging (direction-wise reweighting) die de anisotropie compenseert: kritieke richtingen voor prioritaire taken krijgen meer gewicht, terwijl minder belangrijke richtingen worden gedempt.

Formulering:
De samengevoegde gewichten worden berekend als:
$$W_{merged} = W_0 + \sum \phi_{ik} \cdot S_{ik}$$
waarbij $S_{ik}$ de basisrichtingen zijn (rang-1 factoren of gedeelde SVD-richtingen) en $\phi_{ik}$ de geoptimaliseerde gewichten zijn die de voorkeuren en anisotropie in acht nemen.

3. Belangrijkste Bijdragen

1. Identificatie van twee kritieke eigenschappen: De auteurs identificeren en kwantificeren subspace coverage en directional anisotropy als de twee belangrijkste factoren die succesvol LoRA-merging bepalen. Ze tonen aan dat het negeren van deze factoren leidt tot subruimte-inzinking (collapse) en bias.
2. Het TARA-raamwerk: Een nieuwe methode die taakvoorkeuren expliciet integreert in het merging-proces. Het combineert het behoud van de subruimte (via rang-1 of SVD-structuur) met een mechanisme voor anisotropie-uitlijning via entropie-minimalisatie.
3. Uitgebreide Empirische Validatie: De methode is getest op acht visuele benchmarks (met CLIP/ViT) en zes Natural Language Inference (NLI) taken (met LLaMA-3), evenals op Vision-Language Models (LLaVA).

4. Resultaten

TARA presteert consequent beter dan bestaande "vanilla" merging-methoden (zoals Task Arithmetic, TIES) en recente LoRA-bewuste methoden (zoals KnOTS, LoRA-LEGO, AdaMerging).

• Visuele Taken (CLIP): TARA-Variant B bereikte een gemiddelde genormaliseerde nauwkeurigheid van 76,3%, vergeleken met ~73,6% voor de beste concurrenten (Iso-CTS) en ~68% voor KnOTS.
• Taalkundige Taken (LLaMA-3): Op NLI-benchmarks behaalde TARA een gemiddelde van 80,3%, wat een aanzienlijke verbetering is ten opzichte van baselines die vaak rond de 75% blijven.
• Generalisatie naar Ongeziene Taken: TARA toont superieure generalisatie naar taken waarvoor geen LoRA-checkpoints beschikbaar waren tijdens het samenvoegen.
• Robuustheid: De methode is minder gevoelig voor veranderingen in de voorkeursvector (preference sensitivity) dan AdaMerging, wat resulteert in stabielere prestaties over verschillende taakcombinaties.
• Efficiëntie: Hoewel TARA een gradient-based optimalisatie vereist, is het computatiekosten en geheugengebruik (VRAM) vergelijkbaar met andere geavanceerde methoden en aanzienlijk lager dan het trainen van een nieuw multi-task model.

5. Betekenis en Impact

Deze paper is significant omdat het een fundamenteel inzicht biedt in waarom bestaande LoRA-merging-methoden falen: ze behandelen LoRA-updates vaak als zwarte dozen of negeren de interne structuur en de ongelijke gevoeligheid van richtingen.

• Theoretisch Inzicht: Het introduceert het concept van "directional anisotropy" in de context van model merging, wat een nieuwe richting opent voor onderzoek naar hoe we modelupdates moeten combineren.
• Praktische Toepassing: TARA biedt een robuuste oplossing voor het creëren van "general-purpose" modellen uit gespecialiseerde LoRA-adapters. Dit is cruciaal voor de schaalbaarheid van Large Language Models (LLMs) en Vision-Language Models (VLMs), waar het trainen van één groot model voor alle taken vaak onhaalbaar is.
• Toekomstige Richting: Het werk suggereert dat toekomstige merging-algoritmen niet alleen naar de grootte van de updates moeten kijken, maar ook naar de geometrie van de subruimtes en de richting-gevoeligheid om echte multi-task prestaties te bereiken.

Kortom, TARA bewijst dat het actief managen van subruimte-coverage en anisotropie de sleutel is tot het creëren van hoogwaardige, generalistische modellen via LoRA-merging.