Technische Samenvatting: TraceNAS

Probleemdefinitie

De implementatie van Large Language Models (LLMs) wordt beperkt door enorme parameteraantallen en computationele vereisten. Hoewel gestructureerde pruning een weg biedt naar efficiënte implementatie door architecturale eenheden (bijv. attention heads, transformer blocks) te verwijderen, kampen bestaande methoden met aanzienlijke beperkingen:

1. Lokale versus Globale Afhankelijkheden: De meeste pruning-methoden evalueren de belangrijkheid van lagen, heads of kanalen in isolatie met behulp van lokale heuristieken (bijv. activatie-gewogen magnituden of Hessian-gebaseerde reconstructie). Deze benaderingen negeren complexe globale structurele afhankelijkheden en de propagatie van pruning-geïnduceerde fouten door het model.
2. Computationele Kosten van Training-Aware Search: Methoden die proberen globale afhankelijkheden te vangen, vertrouwen vaak op training-bewuste Neural Architecture Search (NAS). Deze benaderingen vereisen iteratieve gradiëntupdates en continue stromen van ongeziene tokens om overfitting te voorkomen, waardoor het zoekproces even computationeel duur is als de post-pruning recovery-training zelf.
3. Tekortkomingen van Bestaande Zero-Shot Proxies: Huidige training-vrije proxies zijn ontworpen om trainbare, willekeurig geïnitialiseerde architecturen te identificeren door te focussen op trainbaarheid (bijv. gradiëntstabiliteit, gladheid van het loss-landschap). Ze falen in het evalueren van "functionele overerving" — het vermogen van een gepruind model om prestaties te herstellen door het pre-trained loss-landschap te behouden.

Methodologie

TraceNAS wordt voorgesteld als een training-vrij NAS-framework dat ontworpen is om de diepte en breedte van LLMs gezamenlijk te optimaliseren zonder zoek-tijd training.

1. Zoekruimte en Codering

Het framework behandelt de pre-trained LL als een super-netwerk. De zoekruimte $G$ is gedefinieerd als een gezamenlijke distributie over:

• Diepte ($d$): Een discrete maskering over $L$ transformer blocks, wat niet-uniforme dieptereductie mogelijk maakt terwijl de continuïteit van de residual stream behouden blijft.
• Breedte ($\kappa$): Parameterretentieratio's voor attention en MLP sub-blocks, wat heterogene sparsiteit over het model mogelijk maakt.

2. In-Place Architecturale Realisatie

Om de geheugenoverhead van het instantiëren van duizenden kandidaat-modellen te vermijden, gebruikt TraceNAS een in-place masking strategie:

• Breedte Maskering: Gebruikt een activatie-gewogen product (vergelijkbaar met Wanda) om binaire maskers voor gewichtskanalen te genereren. Deze $O(d^2)$ heuristiek identificeert functioneel kritieke kanalen terwijl outlier-features behouden blijven.
• Diepte Maskering: Slaat volledige blokken over als ze gedeactiveerd zijn, waarbij de activaties via residual connections worden gerouteerd.
• Executie: Het framework wijzigt de gewichtsmatrices van het basismodel tijdelijk om gradiënten te berekenen op een minimale kalibratieset, en herstelt vervolgens de originele gewichten, waardoor het basismodel behouden blijft voor opeenvolgende iteraties.

3. De Zero-Shot Proxy: Gradient Trace Correlation

De kerninnovatie is een schaal-invariante proxy $\Phi$ die de functionele uitlijning meet tussen een gepruind kandidaat-model ($M_{sub}$) en de pre-trained basis ($M_{base}$).

• Gradient Traces: Gradiënten worden berekend op een kalibratieset. Om schaalbaarheid te garanderen, worden deze geprojecteerd in een Low-Rank (LoRA) subspace, waarbij gebruik wordt gemaakt van de intrinsieke lage dimensionaliteit van LLMs om de geheugenoverhead met ordes van grootte te verminderen.
• Uitlijningsmetriek: De proxy berekent de Pearson Correlatiecoëfficiënt tussen de gestandaardiseerde gradient traces van de gepruinde en de basismodellen voor elke laag. Dit zorgt voor schaal-invariantie door te focussen op directionele uitlijning in plaats van magnitude-verschuivingen.
• Sparsity-Gewogen Aggregatie: De uiteindelijke score $\Phi$ is een gewogen som van laag-specifieke correlaties, waarbij de gewichten overeenkomen met de parameterretentieratio's. Dit verankert de globale score in gebieden met een hoge capaciteit, om te voorkomen dat ruis van agressief gepruinde sub-blokken de ranking vertekent.

4. Evolutionaire Zoekstrategie

TraceNAS hanteert een evolutionaire strategie om de discrete zoekruimte te navigeren:

• Initialisatie: Een populatie wordt gegenereerd op basis van globale belangrijheidspriors.
• Evolutie: De populatie ondergaat crossover en mutatie. Diepte-evolutie gebruikt bit-flip mutaties en continue-range crossover, terwijl breedte-evolutie interpolatie en Gaussian jitter gebruikt.
• Selectie: Kandidaten worden uitsluitend gerangschikt op basis van de proxy $\Phi$, waarbij modellen worden geselecteerd met de hoogste potentie voor prestatieherstel tijdens de post-pruning continued pretraining (CPT).

Belangrijkste Bijdragen

1. Schaalbare Training-Vrije NAS: TraceNAS biedt een verenigd framework voor gezamenlijke diepte- en breedte-pruning, wat de ontdekking van niet-uniforme architecturen mogelijk maakt zonder zoek-tijd training.
2. Functionele Overervingsproxy: Introduceert een nieuwe, schaal-invariante gradiënt-gebaseerde proxy die de uitlijning van gradient traces tussen gepruinde en pre-trained modellen evalueert. Het bereikt een hoge correlatie met downstream accuratesse (Spearman $\rho = 0.94$, Kendall $\tau = 0.82$).
3. Efficiëntie: Het framework vermindert de zoekoverhead door gebruik te maken van low-rank gradiënt subspaces en in-place masking, wat de ontdekking van high-fidelity modellen op een enkele GPU mogelijk maakt.

Experimentele Resultaten

TraceNAS werd geëvalueerd op Llama-2-7B, Llama-3.1-8B, en Qwen-2.5-14B, waarbij ze respectievelijk gepruind werden naar 2.7B, 4.6B, en 8.4B parameters.

• Zoekefficiëntie: TraceNAS identificeerde optimale architecturen in 8.5 GPU-uren (op een A100), wat een 10× reductie betekent in zoektijd vergeleken met training-bewuste baselines zoals DarwinLM. Het vereiste slechts 98M zoek-tokens, een 10× reductie vergeleken met DarwinLM.
• Prestaties:
  • Llama-2-7B: TraceNAS behaalde de hoogste gemiddelde accuratesse (62.81%) over acht reasoning benchmarks, waarmee het DarwinLM (62.57%) en ShearedLLaMA (62.63%) versloeg. Opvallend genoeg daalde de prestatie van ShearedLLaMA significant (naar 60.86%) wanneer het werd beperkt tot hetzelfde token-budget dat door TraceNAS werd gebruikt.
  • Llama-3.1-8B & Qwen-2.5-14B: TraceNAS presteerde beter dan uniforme pruning baselines en training-bewuste methoden (DarwinLM, E3-Pruner) op complexe reasoning taken (bijv. ARC-Challenge, BoolQ) terwijl het aanzienlijk minder recovery tokens vereiste (20B vs. 94B voor Minitron).
• Proxy Validatie: De proxy demonstreerde een hoge stabiliteit over variërende hyperparameters (steekproefgrootte, contextlengte en LoRA-rank), wat bevestigt dat de uitlijning van de gradient trace effectief de functionele overerving vastlegt, zelfs onder beperkte settings.

Betekenis en Claims

Het artikel claimt dat TraceNAS de kritieke bottleneck van computationele kosten in gestructureerde pruning aanpakt. Door de focus te verschuiven van "trainbaarheid" (leren vanaf nul) naar "overerving" (het behoud van het pre-trained loss-landschap), stelt TraceNAS in staat om high-performing, niet-uniforme gepruinde modellen te ontdekken met minimale middelen.

De auteurs benadrukken dat hun methode:

• De prohibitieve kosten elimineert van zoek-tijd training, waardoor grootschalige modelcompressie toegankelijk wordt.
• Globale structurele afhankelijkheden vastlegt die lokale heuristieken missen, waardoor het behoud van complexe redeneer-capaciteiten wordt gegarandeerd.
• De CO2-voetafdruk vermindert van modelcompressie door de benodigde GPU-uren voor architectuurzoekopdrachten drastisch te verlagen.

Het werk positioneert TraceNAS als een schaalbare fundering voor de implementatie van efficiënte LLMs in omgevingen met beperkte middelen, terwijl de functionele integriteit van dense pre-trained modellen behouden blijft.