Technische Samenvatting: WeightCLIP

Probleemstelling

Weight space learning (WSL) behandelt neurale netwerkparameters (NN-parameters) als een datamodaliteit voor representatieleer, wat taken mogelijk maakt zoals het voorspellen van modeleigenschappen op basis van gewichten of het synthetiseren van nieuwe modellen. Hoewel bestaande benaderingen veelbelovend zijn gebleken, lijden ze aan een kritieke beperking: de geleerde gewichtsruimte-representaties missen een expliciet semantisch referentiekader. Bijgevolg zijn latente ruimtes die puur getraind zijn voor reconstructie of eigenschapsvoorspelling moeilijk te navigeren; nabijheid in de latente ruimte komt niet noodzakelijkerwijs overeen met variatie op datasetniveau, en "navigatie" brengt niet intuïtief een relatie tot met specifieke datadistributies. Dit beperkt het vermogen om zero-shot transfer, controleerbare gewichtsgeneratie of effectieve verfijning uit te voeren zonder dure bemonstering in de hoogdimensionale gewichtsruimte.

De auteurs stellen dat getrainde modellen een impliciete "bijschrift" (caption) dragen: de dataset waarop ze zijn getraind. Het kernprobleem dat wordt aangepakt, is hoe deze datasetinformatie kan worden ingezet om de gewichtsruimte-representatie te hervormen, waardoor een gedeelde latente ruimte ontstaat waar datasets dienen als natuurlijke supervisiesignalen, analoog aan hoe tekst visuele taalmodellen stuurt.

Methodologie: WeightCLIP

WeightCLIP stelt een methode voor om datasetrepresentaties af te stemmen op modelgewicht-representaties met behulp van een contrastieve leerdoelstelling. Het framework bestaat uit drie primaire componenten:

1. Weight-Space Autoencoder: Volgens eerder werk (Schürholt et al., 2024) worden neurale netwerken gerepresenteerd als sequenties van "gewichtstokens" (platgeslagen subsets van laagparameters). Een encoder $g_\phi$ comprimeert deze tokens naar latente representaties, en een decoder $h_\psi$ reconstrueert de gewichten.
2. Dataset Encoder: Een dataset encoder $d_\theta$ (gebaseerd op de DeepSets-architectuur) brengt een kleine subset afbeeldingen uit een dataset in naar een dataset-embedding $e_D$.
3. Contrastieve Afstemming: De kerninnovatie is een token-niveau contrastieve verliesfunctie die de latente representaties van de modelgewichten afstemt op de dataset-embeddings. In tegenstelling tot methoden die een aparte projectieruimte leren, hervormt WeightCLIP direct de latente ruimte van de autoencoder. De verliesfunctie bevat een reconstructieterm ($L_{recon}$) en een bidirectionele contrastieve afstemmingsterm ($L_{align}$) die de gelijkenis tussen de token-representaties van een model en zijn trainingsdataset-embedding maximaliseert, terwijl de gelijkenis met andere datasets wordt geminimaliseerd.

Downstream Componenten

Zodra de latente ruimte is afgestemd, introduceert WeightCLIP twee mechanismen voor downstream taken:

• Dataset-naar-Model Mapping: Om modellen te genereren vanuit een nieuwe dataset-prompt, vertaalt een lichtgewicht mapper de dataset-embedding $e_D$ naar een sequentie van model-tokenrepresentaties $Z$. De auteurs stellen twee mappers voor:
  • Lineaire Mapper: Een eenvoudige affine transformatie die de tokensequentie regresseert vanuit de dataset-embedding.
  • Memory-Bank Mapper: Een methode die voorspellingen stuurt richting het manifold van echte trainingsmodellen door tokens te synthetiseren via een convexe combinatie van opgeslagen trainings-tokenrepresentaties, gewogen door een zachte buurtverdeling.
• Latente Verfijning: Om mismatches (bijv. verschillen in de classifier head) na generatie aan te pakken, stellen de auteurs voor om het model direct in de latente ruimte te verfijnen in plaats van de gewichten te finetunen. Dit proces omvat:
  1. Globale Prior: Het projecteren van de initiële latente tokens op een hypersferische schil (geschat uit de trainingsdata) om te garanderen dat ze op het geldige model-manifold liggen.
  2. Lokale Optimalisatie: Het optimaliseren van de latente tokens $Z$ om de taakverlies te minimaliseren op een kleine doel-dataset, met beperkingen die de updates dicht bij het geleerde manifold houden. Dit is structureel verschillend van standaard gewichtsruimte fine-tuning.

Belangrijkste Bijdragen

Het artikel levert drie primaire bijdragen:

1. Dataset-Afgestemde Latente Ruimte: De auteurs laten zien dat het gebruik van datasets als referentiepunten via contrastieve afstemming de gewichtsruimte-latente ruimte organiseert. Dit maakt intuïtieve navigatie mogelijk waarbij modellen getraind op vergelijkbare datasets samen clusteren.
2. Dataset-Gestuurde Gewichtsgeneratie: Ze introduceren methoden om een data-embedding (data-prompt) direct in kaart te brengen naar een sequentie van model-embeddings, die gedecodeerd kunnen worden naar neurale netwerkgewichten. Dit maakt het genereren van modellen mogelijk die specifiek zijn afgestemd op een doel-dataset, inclusclusief out-of-distribution (OOD) datasets.
3. Latent Manifold Verfijning: Ze stellen een verfijningsprocedure voor die de taakverlies optimaliseert via de decoder, terwijl de latenten worden beperkt tot een hypersferische schil en een lokale buurt. Dit maakt effectieve test-time adaptatie mogelijk die vaak beter presteert dan standaard fine-tuning van de gegenereerde gewichten.

Experimentele Resultaten

De auteurs evalueren WeightCLIP op twee model-zoos: een grote collectie van 10.000 CNN-checkpoints en 1.000 ResNet-18 checkpoints, getraind op diverse beeldclassificatie-datasets.

• Organisatie van de Latente Ruimte: t-SNE visualisaties en k-NN classificatie tonen aan dat afstemming de modelrepresentaties significant clustert op basis van hun trainingsdataset (wat de classificatienauwkeurigheid verbetert van 77% naar 98% vergeleken met niet-afgestemde baselines). Interpolatie tussen dataset-prompts levert vloeiende trade-offs op in modelprestaties, wat de navigeerbaarheid bevestigt.
• Retrieval: WeightCLIP verbetert de dataset-naar-model retrieval aanzienlijk. In in-distribution settings bereikt het 95% Recall@1, waarmee het zowel niet-afgestemde training als de TANS-baseline overtreft. In OOD-settings bereiken de opgehaalde modellen een hogere zero-shot en gefinetunede nauwkeurigheid dan de baselines.
• Generatie: Wanneer dataset-prompts worden gebruikt om nieuwe modellen te genereren, presteert WeightCLIP beter dan training vanaf nul, TANS en hypernetwerk-baselines (Text2Model). Zelfs bij initialisatie (0 epochs fine-tuning) vertonen de gegenereerde modellen een sterke afstemming met de doeltaak. De Memory-Bank Mapper biedt over het algemeen de beste initiële prestaties, terwijl de Lineaire Mapper competitief blijft.
• Verfijning: Latente verfijning bereikt prestaties die vergelijkbaar met of superieur aan standaard gewichtsruimte fine-tuning zijn, met hetzelfde computationele budget. Ablatie-studies bevestigen dat zowel de hypersferische beperking als de initiële afstemming noodzakelijk zijn voor deze winst.

Betekenis en Claims

Het artikel claimt dat WeightCLIP een sleutelbeperking van bestaande WSL-methoden aanpakt door expliciet datasetinformatie te gebruiken om de gewichtsruimte-representatie te hervormen. Door datasets als een semantisch referentiekader te behandelen, maakt de methode het volgende mogelijk:

• Gestructureerde Navigatie: De latente ruimte wordt querybaar en navigeerbaar op basis van dataset-semantiek.
• Efficiënte Model-Synthese: Het maakt het mogelijk om sterke model-initialisaties voor nieuwe datasets te genereren zonder dure bemonstering of training vanaf nul.
• Verbeterde Adaptatie: De voorgestelde latente verfijning biedt een structureel anders en vaak effectiever alternatief voor standaard fine-tuning.

De auteurs positioneren dit werk als een stap naar het meer tijd- en rekenefficiënt maken van neurale netwerktraining voor bekende architecturen op nieuwe datasets. Zij erkennen beperkingen en merken op dat de huidige reikwijdte beperkt is tot computer vision-datasets en homogene modelarchitecturen, en dat de methode vertrouwt op trainingsaugmentaties om gewichtsruimte-symmetrieën af te handelen in plaats van ze expliciet af te dwingen.