Technische Samenvatting: TaskPGM – Het leren van taakmengsels uit taakaffiniteiten

Probleemstelling

Supervised Fine-Tuning (SFT) voor Large Language Models (LLMs) leunt zwaar op de verdeling van het trainingsbudget over een heterogene set instructietaken. Huidige praktijken vertrouwen vaak op eenvoudige heuristieken, zoals uniforme sampling of op datasetgrootte gebaseerde proportionele weging. Deze methoden houden geen rekening met taakinteracties, wat leidt tot twee primaire inefficiënties:

1. Redundantie: Het overmatig wegen van taken die bijna identieke gedragsupdates induceren, waardoor het computationele budget wordt verspild.
2. Gemiste Transfer: Het onderbenutten van taken die complementaire informatie of cross-domain transfervoordelen bieden.

Bestaande geautomatiseerde benaderingen (bijv. op invloed gebaseerde retrieval zoals LESS, of submodulaire selectie zoals SMART) optimaliseren vaak voor discrete deelverzamelingen in plaats van continue mengverhoudingen, vertrouwen op semantische gelijkenis die het downstream modelgedrag negeert, of vereisen dure iteratieve training. De kernuitdaging is het automatisch bepalen van een continu taakmengsel dat tegelijkertijd representatief (hoge utiliteit) en divers (lage redundantie) is om de downstream prestaties onder een vast budget te maximaliseren.

Methodologie: TaskPGM

De auteurs stellen TaskPGM (Task Mixtures via Probabilistic Graphical Modelling) voor, een framework dat SFT-datamenging beschouwt als een energie-minimalisatieprobleem over een dichte Markov Random Field (MRF).

1. Formulering van het Probabilistisch Grafisch Model

Taken worden gemodelleerd als knopen in een MRF. Het doel is om een optimale samplingdistributie $p^* \in \Delta_n$ (waarbij $\Delta_n$ de waarschijnlijkheidssimplex is) te vinden die een energiefunctie $E(p)$ maximaliseert. De energiefunctie bestaat uit twee componenten:

• Unary Potentials ($\Psi_{un}$): Coderen de utiliteit of representativiteit van individuele taken. Dit wordt afgeleid van de totale gelijkenis massa van een taak ten opzichte van de gehele collectie.
• Pairwise Potentials ($\Psi_{pair}$): Coderen redundantie en diversiteit. Deze termen bestraffen de co-occurrence van zeer gelijkaardige taken in het mengsel.

Het optimalisatie-objectief wordt geformuleerd als:
$$\max_{p \in \Delta_n} E(p) = \Psi_{un}^\top p - \frac{1}{2} p^\top \Psi_{pair} p$$
Dit is equivalent aan het minimaliseren van een convexe kwadratische functie (onder specifieke condities) onder simplex-restricties.

2. Gedragsmatige Taakaffiniteitsmetrieken

Een belangrijke innovatie is de methode voor het berekenen van paarwijze taakaffiniteiten ($S_{ij}$). In plaats van te vertrouwen op semantische gelijkenis (embeddings) of metadata, berekent TaskPGM affiniteiten op basis van gedragsafwijkingen afgeleid van de predictieve distributies van single-task fine-tuned modellen.

• Pointwise Mutual Information (PMI): Meet de wederzijdse informatie tussen de predictieve distributies van twee taken wanneer ze op elkaars datasets worden geëvalueerd.
• Jensen–Shannon Divergence (JSD): Een symmetrische maat voor de discrepantie tussen de predictieve distributies van twee modellen op cross-task voorbeelden.

Deze metrieken vangen hoe taken het modelgedrag op out-of-task voorbeelden vormen, wat direct de cross-task transfer en redundantie adresseert.

3. Optimalisatie en Solvers

• Continu Mengsel: De auteurs leiden een gesloten vorm optimum af voor de continue mengverhoudingen met behulp van Karush-Kuhn-Tucker (KKT) condities.
  • Als de paarwijze gelijkenis-matrix niet Positief Semi-Definiet (PSD) is, passen zij spectrum shifting toe (het toevoegen van een constante aan de diagonaal) om convexiteit te garanderen, wat een regularisatieterm introduceert die de oplossing richting uniformiteit stuurt als de niet-PSD structuur ernstig is.
  • Wanneer de inverse slecht geconditioneerd is of de niet-negativiteitsrestricties worden geschonden, gebruiken zij Projected Gradient Descent (PGD) op de waarschijnlijkheidssimplex.
• Discreet/Budgetteerde Variant: Voor scenario's die de selectie van een deelverzameling van taken (top-$k$) vereisen, tillen de auteurs de continue optimalisatie naar een discrete verzamelfunctie. Zij bewijzen dat deze functie monotoon en zwak submodulair is onder standaard regulariteitscondities (Restricted Strong Convexity en Smoothness), wat efficiënte greedy constructie met benaderingsgaranties mogelijk maakt.

Belangrijkste Bijdragen

1. Probabilistisch Framework voor Mengselontwerp: Een nieuwe formulering van SFT-datamenging als energie-minimalisatie op een MRF, waarbij expliciet een afweging wordt gemaakt tussen representativiteit (unary termen) en diversiteit (pairwise termen).
2. Gedragsgebaseerde Affiniteitsmetrieken: Het voorstel om taakrelaties te berekenen met behulp van divergenties in predictieve distributies (JSD/PMI) van single-task fine-tuned modellen, waarmee verder wordt gegaan dan statische semantische gelijkenis om functionele taakafhankelijkheden te vangen.
3. Theoretisch Onderbouwde Optimalisatie:
   • Afleiding van een gesloten vorm oplossing voor continue mengsels via KKT-condities.
   • Vaststelling van zwakke submodulariteit voor de discrete budgetteerde variant, wat greedy selectie met theoretische garanties ondersteunt.
4. Empirische Validatie: Demonstratie van consistente prestatiewinsten over meerdere modelfamilies (Qwen2-7B, Llama-2-7B) en benchmarks (BIG-Bench Hard, MMLU, etc.) vergeleken met heuristieken en state-of-the-art selectie-baselines.

Experimentele Resultaten

Het framework werd geëvalueerd op 309 instructietaken met behulp van Qwen2-7B en Llama-2-7B modellen met budgetten van 25K en 50K instanties.

• Prestatiewinsten: TaskPGM (gebruikmakend van zowel PMI als JSD varianten) presteerde consequent beter dan uniforme sampling, size-proportional (EPM) en random baselines.
  • Op Qwen2-7B (25K budget), verbeterde de PMI-variant de prestaties met bijna 3% ten opzichte van heuristische baselines op complexe redeneerbenchmarks.
  • Op Llama-2-7B, presteerde het door PMI afgeleide mengsel met 2% tot meer dan 4% beter dan Uniform en EPM op BBH.
• Vergelijking met SOTA: TaskPGM bleef competitief met of superieur aan LESS (een sterke op influence-function gebaseerde baseline). Opvallend genoeg bereikt TaskPGM dit met een enkele general-purpose checkpoint, terwijl LESS voor elk target een onderscheidend, target-bewust model construeert.
• Efficiëntie: TaskPGM is computationeel efficiënter dan LESS. De gelijkenisberekening is een eenmalige geamortiseerde kostenpost (~16–34 GPU-uren), waarna mengsels voor elk budget in seconden kunnen worden afgeleid via de gesloten vorm oplossing. In tegen tegenstelling hiertoe vereist LESS ~320 GPU-uren om gradient datastores voor elk target te bouwen.
• Schaalbaarheid: Het verhogen van het budget van 25K naar 50K leverde consistente winsten op voor TaskPGM, terwijl heuristische methoden (Uniform, EPM) gefragmenteerd en niet-robuust gedrag vertoonden naarmate de budgetten groeiden.

Betekenis en Claims

Het artikel claimt dat TaskPGM een principieel, interpreteerbaar en efficiënt benadering biedt voor SFT-datamenging.

• Interpreteerbaarheid: De geleerde gewichten en gelijkenis-heatmaps onthullen taakclusters, redundantie en "attractor" taken, wat inzicht biedt in waarom bepaalde bronnen helpen.
• Generalisatie: In tegen tegenstelling tot methoden die optimaliseren voor specifieke targets, produceert TaskPGM een enkel mengsel dat de prestaties over een breed spectrum van downstream taken (redeneren, taalbegrip, kennis) verbetert.
• Robuustheid: De winsten worden gedreven door functionele taalselectie op basis van modelgedrag in plaats van gunstige sampling-variantie, zoals blijkt uit de consistente verbeteringen over verschillende base-modellen en gelijkenis-metrieken.

De auteurs concluderen dat een principiële, op gelijkenis gebaseerde mixing essentieel is voor schaalbare instructie-tuning, en vaak heuristische mengsels evenaart of zelfs overtreft met aanzienlijk minder data. Zij erkennen beperkingen met betrekking tot de vaste pool van taken en de computationele kosten van affiniteitschatting, en suggereren schaalbare benaderingen als een toekomstige richting.