Technische Samenvatting: TAS-LoRA

Probleemstelling

Transformer Architecture Search (TAS) heeft als doel om automatisch optimale Vision Transformer (ViT)-architecturen te ontdekken, waardoor de handmatige inspanning die voor het ontwerp nodig is, wordt verminderd. De huidige state-of-the-art TAS-methoden vertrouwen voornamelijk op one-shot-benaderingen, waarbij een "supernet" wordt getraind dat alle mogelijke subnets bevat. In deze gewichtsverstrengelde supernets delen alle subnets een gemeenschappelijke set gewichten.

Het artikel identificeert een kritieke beperking in dit paradigma: het feature collapse-probleem. Omdat het supernet wordt geoptimaliseerd om een globale verliesfunctie over alle subnets te minimaliseren, neigen de gedeelde gewichten ertoe om generieke kenmerken te leren die adequaat werken voor diverse architecturen, maar falen in het vastleggen van subnet-specifieke kenmerken. Bijgevolg kunnen individuele subnets hun unieke architecturale eigenschappen (bijvoorbeeld specifieke aantallen attention heads of MLP-ratio's) niet volledig benutten, wat leidt tot suboptimale prestaties in vergelijking met subnets die onafhankelijk vanaf nul zijn getraind. Empirische analyse in het artikel bevestigt dat subnets binnen gewichtsverstrengelde supernets bijna identieke feature-representaties vertonen (hoge cosinusgelijkenis), terwijl onafhankelijk getrainde subnets onderscheidende representaties ontwikkelen.

Methodologie: TAS-LoRA

Om feature collapse aan te pakken terwijl de computationele efficiëntie behouden blijft, stellen de auteurs TAS-LoRA voor (Transformer Architecture Search met Mixture-of-LoRA Experts). De methode integreert Low-Rank Adaptation (LoRA) in het trainingsproces van het supernet om gespecialiseerde feature-learning voor elke subnet mogelijk te maken, zonder de verbodsdure kosten van het toewijzen van unieke parameters aan elke mogelijke subnetconfiguratie.

1. Mixture-of-LoRA-Experts (MoLE)

In plaats van een uniek LoRA-module toe te wijzen aan elke subnet (wat computationeel onuitvoerbaar is vanwege de enorme zoekruimte), hanteert TAS-LoRA een Mixture-of-LoRA-Experts (MoLE)-strategie.

• Architectuur: LoRA-experts worden toegepast op de lineaire lagen van het supernet (inclusief QKV-projectie en MLP-lagen).
• Dynamische Routing: Een lichtgewicht router selecteert en combineert dynamisch een beperkte set LoRA-experts op basis van de specifieke architectuur van de bemonsterde subnet.
• Inferentie-efficiëntie: Tijdens de inferentie worden de geselecteerde LoRA-parameters samengevoegd met de oorspronkelijke gewichten van het supernet. Aangezien de subnetconfiguratie tijdens de inferentie vaststaat, draait de router slechts één keer en introduceren de samengevoegde gewichten geen extra computationele overhead.

2. Routerontwerp

De router is ontworpen om de sequentiële afhankelijkheden van transformer-blokken vast te leggen:

• Input: Deze neemt zowel block-niveau attributen (bijvoorbeeld aantal attention heads, MLP-ratio) als subnet-niveau attributen (bijvoorbeeld embedding-dimensie, diepte) als input.
• Mechanisme: Deze attributen worden gecodeerd door een leerbare embedding-laag en verwerkt door een LSTM om sequentiële afhankelijkheden over blokken heen vast te leggen. De LSTM geeft verborgen representaties uit, die via een volledig verbonden "router head" worden geleid om een kansverdeling te genereren over de beschikbare LoRA-experts voor elke laag.

3. Groepsgebaseerde Routerinitialisatie

Een belangrijke uitdaging die wordt geïdentificeerd, is dat de router tijdens de vroege training de neiging heeft om vergelijkbare gewichten toe te wijzen aan alle experts, wat leidt tot redundantie en specialisatie verhindert. Om dit te mitigeren, introduceren de auteurs een groepsgebaseerde routerinitialisatie-techniek:

• Groepering: Transformer-blokken over verschillende subnets worden gegroepeerd op basis van architecturale gelijkenissen (specifiek aantal attention heads, MLP-ratio en embedding-dimensie).
• Initialisatie: Het aantal LoRA-experts wordt ingesteld om overeen te komen met het aantal groepen. De router wordt zo geïnitieerd dat elke expert een bias heeft naar een specifieke groep. Dit wordt bereikt door specifieke bias-termen in de router head in te stellen op een positieve waarde $\beta$ voor het corresponderende groep-expert-paar, terwijl de anderen op nul worden gezet.
• Effect: Dit dwingt de router om in een vroeg stadium van de training verschillende experts toe te wijzen aan verschillende architecturale groepen, wat divers feature-learning stimuleert. Naarmate de training vordert, werkt de router deze toewijzingen dynamisch bij op basis van de eigenschappen van de subnet.

Belangrijkste Bijdragen

1. Identificatie van Feature Collapse in TAS: Het artikel identificeert en analyseert formeel het feature collapse-probleem in gewichtsverstrengelde supernets, en toont aan dat gedeelde gewichten subnets beletten om gespecialiseerde representaties te leren.
2. TAS-LoRA Framework: Een nieuwe methode die LoRA combineert met een dynamisch routingmechanisme (MoLE) om subnet-specifieke kenmerken efficiënt te leren. Het stelt subnets in staat om gedeelde kenmerken te verfijnen op basis van hun architecturale eigenschappen zonder het volledige supernet vanaf nul opnieuw te trainen.
3. Groepsgebaseerde Routerinitialisatie: Een techniek om expert-redundantie in de vroege stadia van de training te voorkomen, waardoor wordt gegarandeerd dat LoRA-experts onderscheidende feature-representaties leren.
4. State-of-the-Art Prestaties: Uitgebreide experimenten tonen aan dat TAS-LoRA aanzienlijk beter presteert dan bestaande TAS-methoden (bijvoorbeeld AutoFormer, DYNAS, FocusFormer) over diverse zoekruimten en datasets.

Experimentele Resultaten

De auteurs hebben TAS-LoRA geëvalueerd op het ImageNet-dataset en verschillende transfer learning-benchmarks (CIFAR-10/100, Flowers, CARS, INAT-19).

• ImageNet-prestaties: TAS-LoRA verbeterde consistent de top-1-accuraatheid van supernets die met AutoFormer en DYNAS waren getraind, over Tiny, Small en Base zoekruimten.
  • Bijvoorbeeld, in de AutoFormer-T-ruimte behaalde TAS-LoRA 75,7% nauwkeurigheid, wat beter presteerde dan de originele AutoFormer (74,7%) en andere SOTA-methoden zoals PreNAS-Ti (77,1%, maar met minder flexibiliteit).
  • In de AutoFormer-B-ruimte behaalde TAS-LoRA-B 82,6% nauwkeurigheid met 54M parameters, wat AutoFormer-B (82,4%) overtrof en gelijkstond aan of grotere modellen zoals GLiT-B (82,3% met 96M parameters) overtrof.
• Transfer Learning: TAS-LoRA toonde superieure generalisatie, met concurrerende of state-of-the-art resultaten op CIFAR-10/100, Flowers, CARS en INAT-19.
• Ablatiestudies:
  • Initialisatie: Groepsgebaseerde initialisatie presteerde consistent beter dan willekeurige initialisatie, wat de rol bevestigt in het verminderen van expert-redundantie.
  • Feature-gelijkenis: Analyse van de cosinusgelijkenis toonde aan dat TAS-LoRA met groepsgebaseerde initialisatie aanzienlijk meer diverse feature-representaties tussen experts produceerde in vergelijking met willekeurige initialisatie of standaard gewichtsverstrengelde methoden.
  • Vergelijking met Fine-tuning: TAS-LoRA presteerde beter dan een volledig gefine-tuned supernet (waarbij alle gewichten worden bijgewerkt), ondanks dat het slechts 50 epochs op de MoLE-lagen werd getraind. Dit geeft aan dat de winst voortkomt uit de subnet-specifieke feature-learningcapaciteit van MoLE, en niet louter uit verlengde training.

Betekenis en Claims

Het artikel beweert dat TAS-LoRA een aanzienlijke vooruitgang in Transformer Architecture Search vertegenwoordigt door het feature collapse-probleem dat inherent is aan gewichtsverstrengelde supernets effectief op te lossen.

• Efficiëntie versus Prestaties: In tegenstelling tot methoden die het trainen van meerdere supernets vereisen (bijvoorbeeld MoS) of het opnieuw trainen vanaf nul wanneer hardwarebeperkingen veranderen (bijvoorbeeld PreNAS), traint TAS-LoRA een enkel supernet met subnet-specifieke adaptatielagen. Dit maakt flexibele aanpassing aan nieuwe hardwarebeperkingen (bijvoorbeeld FLOPs, latentie) mogelijk zonder het volledige model opnieuw te trainen.
• Schaalbaarheid: Door gebruik te maken van een beperkte set experts en een dynamische router, schaalt de methode naar supernets met enorme aantallen mogelijke subnetconfiguraties zonder een lineaire toename in parameters of computationele kosten.
• Generaliseerbaarheid: De aanpak wordt gepresenteerd als een algemene oplossing die toepasbaar is op elk gewichtsverstrengeld supernet, waardoor de prestaties verbeteren ongeacht de onderliggende trainingsstrategie (bijvoorbeeld AutoFormer of DYNAS).

De auteurs concluderen dat TAS-LoRA feature collapse succesvol mitigeert, waardoor subnets gespecialiseerde representaties kunnen leren en state-of-the-art prestaties behalen in zowel architectuuroptimalisatie als transfer learning-taken.