Technische Samenvatting: Bergson: Een Open Source Bibliotheek voor Data-attributie

Probleemstelling

Data-attributie beoogt modelgedrag te verklaren door de causale effecten van individuele trainingsitems op specifieke modeloutputs te bepalen. Hoewel dit vakgebied toepassingen biedt voor debugging, datasetcuratie en het begrijpen van memorisatie of redeneren, brengt het opschalen van deze technieken naar moderne Large Language Models (LLMs) aanzienlijke engineering-uitdagingen met zich mee. De naïeve toepassing van data-attributiemethoden resulteert vaak in onrealistische geheugeneisen, met name bij het werken met grootschalige pre-training datasets of modellen die multi-node ondersteuning vereisen. Bovendien missen veel state-of-the-art methoden, zoals die met een hoge Leave-One-Out (LOO) correlatie, open-source implementaties, wat reproduceerbaarheid en vooruitgang belemmert.

Methodologie en Bibliotheekarchitectuur

Bergson wordt geïntroduceerd als een open-source bibliotheek die is ontworpen om snellere vooruitgang in data-attributie mogelijk te maken door een composabele, schaalbare architectuur te bieden. De bibliotheek ondersteunt een breed scala aan technieken, waaronder differentiatie door volledige trainings-trajecten, benaderde unrolled differentiatie en influence functions.

Kerncomponenten

• Composabele Pipeline: Bergson maakt gebruik van een data-pipeline architectuur (Figuur 1) waarbij componenten zoals Hessian-approximaties, normaliseerders, index-builders en attributie-scorers interoperabel zijn. Workflows kunnen worden gedefinieerd via een Command Line Interface (CLI) of programmatisch met YAML-configuratiebestanden.
• Gradient Opslag en Indexering: De bibliotheek ondersteunt on-disk gradient stores die geïndexeerd zijn per module (inclusief gesplitste attention heads) en ondersteunt per-token of per-sequence collectie. Om grote indexen te verwerken, integreert Bergson met FAISS voor approximate nearest-neighbor (ANN) lookups, terwijl het ook exacte top-k lookups ondersteunt.
• Gedistribueerde Training: Bergson implementeert gerepliceerde en gesharte parallellisme (DDP, FSDP) om attributie op te schalen naar modellen met tot 405 miljard parameters. Het ondersteunt multi-node operaties waarbij FSDP-groepen onafhankelijk dataset-shards kunnen verwerken.
• PEFT Integratie: Native integratie met de PEFT-bibliotheek maakt data-attributie beperkt tot adapter-parameters (bijv. LoRA) mogelijk, wat de per-example gradient footprint met ordes van grootte vermindert.
• Niet-differentieerbare Objectives: Voor taken die geëvalueerd worden via LLM-as-a-judge, ondersteunt Bergson attributie via Group Relative Policy Optimization (GRPO) loss, waardoor attributie van gedragingen mogelijk is die niet worden gevangen door standaard cross-entropy losses.

Geïmplementeerde Algoritmen

Bergson biedt de eerste open-source implementaties voor drie leidende methoden:

1. MAGIC: Een algoritme voor exacte differentiatie door unrolled training loops met behulp van checkpointing. Het bereikt $O(N \log N)$ computation en $O(\log N)$ diskruimte (of een nieuwe $O(N)$ compute/$O(\sqrt{N})$ diskruimte trade-off).
2. SOURCE: Een methode voor benaderde unrolled differentiatie die het trainings-traject verdeelt in segmenten, waarbij Hessians en gradients als constant binnen elk segment worden behandeld.
3. TrackStar: Een schaalbare benaderingsmethode voor influence functions die per-example gradients projecteert in een laag-dimensionale ruimte met behulp van random projecties en autocorrelatie Hessian-approximaties gebruikt.

De bibliotheek ondersteunt ook standaard influence function varianten (K-FAC, EK-FAC, TK-FAC/Shampoo) en gradient normalisatie strategieën (Adam, Adafactor, unit normalisatie).

Belangrijkste Bijdragen

1. Open-Source Implementatie: Bergson introduceert de eerste publieke implementaties van MAGIC, SOURCE en TrackStar, waarmee een gat in de tooling voor cutting-edge data-attributiemethoden wordt opgevuld.
2. Schaalbaarheid: De bibliotheek is ontworpen om miljarden-parameter modellen en massieve datasets te verwerken, met ondersteuning voor on-disk gradient stores, FAISS-integratie en multi-node gedistribueerde training.
3. Methode Selectie Gids: Het paper biedt een gids voor practitioners om methoden te selecteren op basis van bronbeperkingen, query-volume en de noodzaak voor precisie (bijv. het aanbevelen van unrolled differentiatie voor veiligheidskritische filtering en gecomprimeerde indexen voor high-query debugging).
4. Evaluatie Framework: Bergson bevat tools voor validatie met behulp van de Linear Data Modeling Score (LDS) en varianten voor top-k/bottom-k filtering, evenals ondersteuning voor leave-k-out re-training experimenten.

Experimentele Resultaten

Het paper presenteert verschillende validatie-experimenten en casestudies:

• Nauwkeurigheid Validatie: Op een GPT-2 (124M) model gefinetuned op WikiText-2, bereikte MAGIC (unrolled differentiation) de hoogste Spearman correlatie (0.979) met ground-truth leave-one-out effecten, gevolgd door EK-FAC (0.318) en TrackStar (0.164). Echter, in sorted subset leave-k-out evaluaties (relevanter voor data filtering), behaalden alle methoden correlaties van 0.7 of hoger.
• Token-Level Attributie (Biosecurity): Gebruikmakend van de WMDP Bio benchmark, hebben de auteurs een 7B model gefinetuned met LoRA en per-token unrolled differentiation gebruikt om invloedrijke tokens te identificeren. Re-training met verhoogde importance weights op de top 10% van de invloedrijke tokens verbeterde de evaluatie-nauwkeurigheid met een extra 1.5 procentpunt vergeleken met unweighted re-training.
• Fact Tracing over Stijlen: In een synthetisch experiment met biografische feiten geschreven in "Shakespeareaanse" versus "Legalese" stijlen, slaagde TrackStar erin de juiste feitelijke promotors te identificeren in 15.26% van de top-5 resultaten, vergeleken met 0.80% voor een raw gradient inner product baseline. TrackStar verminderde ook "style leak" (het matchen van de query-stijl in plaats van de feitelijke bron) aanzienlijk.
• Schaalbaarheid Performance: De tijd voor gradient collectie schaalt efficiënt met het aantal GPU's (1 tot 16 A100's) voor modellen variërend van 160M tot 12B parameters, waarbij FSDP en data-parallelisme strategieën worden gebruikt om communicatie-overhead te minimaliseren.

Betekenis en Claims

De auteurs beweren dat Bergson de drempel voor data-attributie onderzoek verlaagt door een uniforme, schaalbare infrastructuur te bieden. Door de eerste implementaties van MAGIC, SOURCE en TrackStar aan te bieden, stelt de bibliotheek practitioners in staat om te experimenteren met state-of-the-art technieken die voorheen ontoegankelijk waren of een aanzienlijke engineering-inspanning vereisten om te bouwen.

Het paper benadrukt dat effectieve data-attributie kan helpen bij het begrijpen van modelgedrag, het debuggen van onverwachte generaties en het cureren van trainingsdatasets voor veiligheid en alignment. De auteurs merken op dat hoewel data-curatie dual-use is (het kan gedrag zowel versterken als verminderen), de specificiteit van data-attributie naar verwachting de alignment en veiligheidsinspanningen naar verhoudingsgewijs meer zal ten goede komen door precieze targeting van relevante data mogelijk te maken. De bibliotheek is gepositioneerd als een tool om meer onderzoekers in staat te stellen interessant modelgedrag te bestuderen via data-attributie, met name in de context van grootschalige modellen waar dergelijke analyse voorheen onpraktisch was.