Efficient Estimation of Kernel Surrogate Models for Task Attribution

Technische Samenvatting: Efficiënte Schatting van Kernelsurrogaatmodellen voor Taaktoewijzing

Probleemstelling
Moderne AI-systemen, met name grote taalmodellen (LLM's) en multi-task learners, worden getraind op diverse sets van taken (bijvoorbeeld vertaling, codegeneratie, redeneren). Een fundamentele uitdaging in interpretatie is taaktoewijzing: het kwantificeren van de invloed van elke individuele trainingstaak op de prestaties van het model op een specifieke doeltaak. De grondwaarheidsmetriek hiervoor is Leave-One-Out (LOO) hertraining, die de verandering in prestaties meet wanneer een specifieke taak wordt uitgesloten van de training. LOO is echter op grote schaal computationeel onuitvoerbaar, waarbij $K+1$ volledige trainingruns nodig zijn voor $K$ taken.

Bestaande efficiënte alternatieven omvatten invloedsfuncties (het schatten van het effect van infinitesimale dataperturbaties via Hessiaan-vectorproducten) en lineaire surrogaatmodellen (het leren van een lineaire mapping van taakopname-indicatoren naar prestaties). Hoewel lineaire surrogaten schaalbaar zijn, zijn ze beperkt tot het vastleggen van eerste-orde, additieve relaties. Ze falen in het modelleren van niet-lineaire interacties tussen taken, zoals synergie, antagonisme of XOR-achtige effecten, die veelvoorkomend zijn in complexe leerlandschappen.

Methodologie
Het artikel stelt KERNELSM (Kernel Surrogate Models) voor, een raamwerk dat taaktoewijzing uitbreidt om tweede-orde en hogere-orde niet-lineaire interacties tussen trainingstaken vast te leggen.

1. Theoretische Analyse van Lineaire Surrogaten:
   De auteurs stellen eerst een unifyend taakgewichtingsraamwerk op. Via een tweede-orde Taylor-ontwikkeling van de prestatiefunctie $F(s)$ rondom de uniforme gewichtsvector $s^* = [1]^K$, analyseren ze lineaire surrogaatmodellen.

   • Ze bewijzen dat lineaire surrogaatcoëfficiënten ($\hat{\beta}$) de gradiënt van het prestatielandschap ($\nabla_s F(s^*)$) benaderen, tot termen die worden veroorzaakt door de steekproefverdeling en tweede-orde interactie-effecten.
   • Cruciaal tonen ze aan dat, onder de aanname van kleine tweede-orde interacties (d.w.z. een bijna lineair landschap), lineaire surrogaatmodellen ongeveer equivalent zijn aan invloedsfuncties. Dit verklaart waarom lineaire methoden in sommige settings goed werken, maar falen wanneer taakinteracties niet-additief zijn.

2. Kernelsurrogaatmodellering:
   Om niet-additiviteit aan te pakken, vervangen de auteurs de lineaire mapping door een niet-lineaire functie die wordt geleerd via Kernel Ridge Regression (KRR) in een Reproducing Kernel Hilbert Space (RKHS).

   • Kernkeuze: Ze maken gebruik van de Radiale Basisfunctie-kern (RBF-kern) ($k(s^{(a)}, s^{(b)}) = \exp(-\gamma \|s^{(a)} - s^{(b)}\|^2)$). Deze keuze wordt gemotiveerd door de geometrie van de binaire subsetruimte $\{0, 1\}^K$, waarbij de RBF-kern op natuurlijke wijze localiteit codeert op basis van Hamming-afstand, waardoor het model complexe afhankelijkheden tussen taakcombinaties kan vastleggen.
   • Doel: De surrogaatfunctie $g_\theta$ minimaliseert het empirische verlies tussen de waargenomen prestatie $F(s^{(i)})$ en voorspellingen, geregulariseerd door de RKHS-norm.

3. Efficiënte Gradiëntgebaseerde Schatting:
   Een grote bottleneck in surrogaatmodellering is de noodzaak om $m$ aparte modellen te trainen om prestatiewaarden $F(s^{(i)})$ te verkrijgen voor $m$ gestreken subsets. De auteurs introduceren een efficiënte schattingsprocedure om herhaalde training te vermijden:

   • Eerste-orde Benadering: Ze maken gebruik van een eerste-orde Taylor-ontwikkeling van de modeloutput $f_W(x)$ rondom een vooraf getrainde initialisatie $W_0$: $f_W(x) \approx f_{W_0}(x) + \langle \nabla f_{W_0}(x), W - W_0 \rangle$.
   • Gradiëntkenmerken: In plaats van opnieuw te trainen, schatten ze de optimale parameterperturbatie $Z^*_s$ voor een subset $s$ door een geregulariseerd regressieprobleem op te lossen waarbij de kenmerken de gradiënten van het vooraf getrainde model zijn ($\nabla f_{W_0}$).
   • Dimensionaliteitsreductie: Om te gaan met hoge-dimensionale gradiënten, passen ze Gaussische willekeurige projecties toe (Johnson-Lindenstrauss-lemma) om de kenmerkruimte te comprimeren, wat snelle berekening op CPU's mogelijk maakt.
   • Nauwkeurigheid: Empirische resultaten tonen aan dat deze eerste-orde benadering minder dan 2% relatieve fout veroorzaakt over diverse datasets en modelgroottes (tot 34B parameters).

Belangrijkste Bijdragen

• Theoretische Unificatie: Het artikel legt een formele verbinding tussen lineaire surrogaatmodellen en invloedsfuncties, en toont aan dat ze equivalent zijn onder voorwaarden van verwaarloosbare tweede-orde taakinteracties.
• Kernelsurrogaatmodellen (KERNELSM): De introductie van een op kernen gebaseerde aanpak (specifiek RBF) die effectief niet-lineaire taakinteracties (synergie/antagonisme) vastlegt die lineaire modellen missen.
• Schaalbaar Schattingsalgoritme: Een gradiëntgebaseerde schattingsprocedure die de prestaties van elke taaksubset benadert met behulp van gradiënten van een enkel vooraf getraind model, waardoor de noodzaak voor herhaalde hertraining wordt geëlimineerd terwijl hoge nauwkeurigheid wordt behouden.
• Empirische Validatie: Uitgebreide evaluatie over modulaire rekenkundige redenering, in-context learning (ICL) en multi-objective reinforcement learning (RL).

Experimentele Resultaten
De auteurs evalueren KERNELSM tegenover baselines, waaronder Invloedsfuncties, TracIn, TRAK, SOURCE en Lineaire Surrogaatmodellen.

• Toewijzingsnauwkeurigheid: KERNELSM presteert consequent beter dan baselines. Bij modulaire rekenkundige taken (met XOR en kwadratische structuren) verbetert het de toewijzingsnauwkeurigheid met tot 42% ten opzichte van bestaande methoden. Over alle domeinen heen bereikt het een 25% hogere correlatie met de LOO-grondwaarheid in vergelijking met lineaire surrogaten en invloedsfuncties.
• Selectie van Downstreamtaken: Wanneer gebruikt om optimale subsets voor downstream-taken te selecteren:
  • In-Context Learning: KERNELSM resulteert in een 40% lagere loss bij demonstratieselectie in vergelijking met lineaire surrogaten.
  • Multi-Objective RL: Het verbetert de doelprijzen met 15% in de Meta-World MT10-benchmark.
• Schaalbaarheid: De methode schaalt naar grote aantallen taken (bijvoorbeeld 500 prompts) zonder prestatiedegradatie en werkt efficiënt op CPU's.
• Robuustheid: De kernbenadering blijft robuust onder verschillende optimalizers (SGD, SAM, NSO) en Hessiaan-regularisaties, terwijl lineaire surrogaten significant degraderen wanneer het verlieslandschap sterk gekromd is (hoge Hessiaan-trace).

Betekenis en Claims
Het artikel beweert dat KERNELSM een schaalbare en flexibele oplossing biedt voor taaktoewijzing in black-box systemen. Door verder te gaan dan de additieve aannames van lineaire modellen, maakt het een nauwkeurigere begrip mogelijk van hoe complexe, niet-lineaire interacties tussen trainingstaken het modelgedrag vormgeven. De auteurs benadrukken dat hun aanpak niet afhankelijk is van specifieke aannames over het trainingsproces (bovenop de eerste-orde benadering van het vooraf getrainde model) en breed toepasbaar is op opkomende paradigma's zoals in-context learning en reinforcement learning. Het werk overbrugt de kloof tussen theoretische invloedsfuncties en praktische, schaalbare datatoewijzing voor moderne grote AI-systemen.