Technische Samenvatting: Prototype Language Models (PRISM)

Probleemstelling

Huidige Large Language Models (LLMs) worden geconfronteerd met een fundamentele uitdaging in Training Data Attribution (TDA): het bepalen welke specifieke trainingsvoorbeelden de output van een model aansturen. In standaard dichte architecturen beïnvloedt trainingsdata voorspellingen indirect via geaccumuleerde effecten op miljarden parameters binnen dichte, verstrengelde paden. Dit maakt de invloed van een enkel voorbeeld verspreid en ontraceerbaar zonder dure, benaderende en vaak statistisch misgespecificeerde post-hoc analyse.

Bestaande gradiëntgebaseerde TDA-methoden (bijv. Influence Functions) kampen met twee cumulerende problemen op schaal:

1. Computationele Inefficiëntie: Exacte berekening van de inverse Hessian-Vector Product (iHVP) is onhaalbaar voor miljarden-parameter modellen, wat noodzakelijke benaderingen vereist die trouw inruilen voor schaalbaarheid.
2. Statistische Misspecificatie: Deze methoden gaan vaak uit van deterministische, convexe doelstellingen die geoptimaliseerd zijn tot convergentie, condities die niet gelden voor stochastische LLM-training. Bijgevolg kunnen attributieschattingen niet overeenkomen met leave-one-out hertraining en zijn ze gevoelig voor hyperparameters van de trainingspipeline (diepte, batchgrootte, datavolgorde).

Het artikel betoogt dat TDA niet louter moet worden behandeld als een post-hoc schattingsprobleem, maar als een architecturaal ontwerpprobleem.

Methodologie: PRISM Architectuur

De auteurs introduceren Prototypes for Interpretable Sequence Modeling (PRISM), een taalmodelarchitectuur die voorspellingen vormt via een ijle, niet-negatieve mengeling van geleerde prototypes, in plaats van een dichte output head.

Kernarchitectuur

PRISM behoudt een standaard Transformer decoder backbone, maar vervangt het laatste dichte pad door een ijle prototype reconstructie:

1. Prototype Bank: Een verzameling van $K$ geleerde prototype vectoren $P = [p_1, \dots, p_K] \in \mathbb{R}^{d \times K}$ wordt onderhouden.
2. IJle Activatie: Voor een gegeven verborgen toestand $z_t$ berekent het model de cosinus-gelijkenis met alle prototypes, past ReLU toe om niet-negativiteit af te dwingen, en selecteert de top-$k$ actieve prototypes ($k \ll K$).
3. Reconstructie: De verborgen toestand wordt gedecomposeerd in een prototype reconstructie $\hat{z}_t$ en een residu $r_t$:
   $$z_t = \hat{z}_t + r_t = \sum_{i \in \mathcal{K}_t} a_{t,i} p_i + r_t$$
   waarbij $\mathcal{K}_t$ de verzameling actieve prototypes is en $a_{t,i}$ hun activaties zijn.
4. Logit Formatie: De volgende-token logits worden additief gevormd:
   $$\ell_t = W z_t = W r_t + \sum_{i \in \mathcal{K}_t} a_{t,i} (W p_i)$$
   Hier is $W p_i$ de vaste "token-logit signatuur" voor elk prototype.

Trainingsdoelstellingen

PRISM wordt end-to-end getraind met vier verliestermen om de taalmodellering-getrouwheid te balanceren met prototype-interpreteerbaarheid:

• Cross-Entropy ($L_{CE}$): Zorgt voor de standaard volgende-token voorspellingskwaliteit.
• Reconstructie Verlies ($L_{REC}$): Bestraft het residu, wat prototypes aanmoedigt om het merendeel van de verborgen toestand te verklaren.
• Clustering Verliezen ($L_{R1}, L_{R2}$): Deze symmetrische doelstellingen verankeren prototypes aan coherente buurten van trainingsdata. $L_{R1}$ trekt prototypes naar hun dichtstbijzijnde token-representaties in de batch, terwijl $L_{R2}$ token-toestanden naar hun dichtstbijzijnde prototypes trekt. Dit zorgt ervoor dat prototypes op de data-manifold liggen in plaats van te driften naar willekeurige ijle dictionary-richtingen.

Theoretisch Inzicht: Curvatuur Lokalisatie

Een belangrijke theoretische bijdrage is de analyse van de loss Hessian in de prototype subspace. De auteurs tonen aan dat de clustering-doelstellingen de curvatuur lokaliseren:

• De Hessian splitst zich in een globale co-usage term (die vangt hoe prototypes samen worden gebruikt) en lokale blok-diagonaal termen (die de geometrie rond elk prototype vangen).
• Sterkere clustering vermindert de cross-prototype koppeling, waardoor de Hessian beter geconditioneerd wordt.
• Deze structuur maakt schaalbare TDA mogelijk: omdat de curvatuur gelokaliseerd is, kunnen invloedsscores snel worden berekend met behulp van gecachte ijle statistieken, waardoor de noodzaak voor volledige parameter-ruimte gradiëntproducten wordt vermeden.

Belangrijkste Bijdragen

1. Architectuur (PRISM): De eerste familie van LLM-schaal prototype taalmodellen voor volgende-token voorspelling. Het breidt case-based reasoning uit naar autoregressieve modellering door een ijle, op training gebaseerde mengeling van prototypes voor elke voorspelling te exposeren.
2. Theorie en Attributie: Demonstreert dat ijle prototype structuren en clustering doelstellingen de curvatuur in het loss landschap lokaliseren. Dit levert een gestructureerde Hessian op die training data attributie $\sim$500$\times$ sneller maakt dan dichte gradiëntgebaseerde baselines (zoals EK-FAC), terwijl het een vergelijkbaar geheugen verbruikt.
3. Empirische Prestaties en Controle:
   • Prestaties: Over modellen variërend van 130M tot 1.6B parameters getraind op maximaal 50B tokens, evenaart PRISM of blijft het binnen 2.5 procentpunten van gematchte dichte GPT baselines op downstream benchmarks.
   • Controleerbaarheid: Prototype controllers kunnen downstream accuratesse met $\sim$3 punten verbeteren door actieve prototypes te herwegen, waarbij correcties traceerbaar zijn naar specifieke trainingsbuurten.
   • Alignment: Gerichte prototype suppressie kan ongewenst gedrag (bijv. NSFW content) verwijderen tijdens inferentie zonder finetuning of meetbaar verlies in generatiekwaliteit.

Experimentele Resultaten

Schaling en Prestaties

• Schaal: PRISM werd geschaald van 130M naar 1.6B parameters op FineWeb en Nemotron-CC corpora (tot 50B tokens).
• Accuratesse: Bij 1.6B parameters bereikte PRISM getraind op 10B tokens een prestatie vergelijkbaar met dichte GPT baselines. Met 20B–50B tokens verkleinde het de kloof met dichte baselines en bereikte het een gemiddelde downstream accuratesse binnen 2.5 punten.
• Efficiëntie: De prototype laag voegt minimale overhead toe (bijv. 1.7% parameter toename bij XL schaal) en bereikt 97.3% van de trainings-throughput van dichte GPT modellen.
• Gronding: Prototype coverage metrieken ($\bar{R}_1, R_2$) bleven hoog (0.94–1.00), wat aangeeft dat prototypes gegrond bleven in echte trainingscontexten in plaats van te driften.

Attributie Efficiëntie

• Snelheid: Gecachte prototype attributie bereikte een 470$\times$ versnelling ten opzichte van EK-FAC op een 124M model onder equivalente GPU geheugenrestricties.
• Signaalbehoud: Gecachte scores behielden $\sim$80% van de top-gerangschikte trainingssequenties vergeleken met volledige prototype ruimte gradiënten, wat aantoont dat de ijle, gelokaliseerde geometrie het noodzakelijke attributie-signaal behoudt.

Controle en Sturing

• Controllers: Een ijle lineaire controller over actieve prototypes verbeterde de accuratesse op vier niet-binaire multiple-choice taken met gemiddeld +2.72 punten, wat de SAE-gebaseerde controller baseline (+0.32 punten) aanzienlijk overtreft.
• Veiligheid Alignment: Door 86 geïdentificeerde NSFW-geassocieerde prototypes (0.52% van de bank) tijdens inferentie te onderdrukken, daalde de gemiddelde NSFW score van 0.668 naar 0.154 op 50.000 generaties, terwijl de tekstkwaliteit ongewijzigd bleef. Op de moeilijkste gevallen (base NSFW score = 2), daalde de score naar 0.144.

Betekenis en Claims

Het artikel claimt dat PRISM een verschuiving laat zien in hoe interpreteerbare taalmodellen ontworpen moeten worden. In plaats van te vertrouwen op uitdagende post-hoc analyse om sporen van trainingsdata te herstellen uit dichte parameters, exposeert PRISM ijle, op training gebaseerde coördinaten direct binnen het voorspellingspad.

De auteurs stellen dat wanneer modellen vergelijkbare voorspellende prestaties behalen, architecturen waarvan de berekeningen gemakkelijker te attribueren, te auditeren en te controleren zijn, de voorkeur verdienen. PRISM bereikt dit door:

1. Training data attributie een native, computationeel goedkope eigenschap van de modelarchitectuur te maken.
2. Directe interventie (sturing, suppressie) op modelgedrag mogelijk te maken zonder finetuning.
3. Een "wetenschap van modelgeneralisatie" te bieden waarbij specifieke test-tijd gedragingen terug te leiden zijn naar specifieke trainingsbuurten.

Het werk suggereert dat "interpretable by design" modellering op schaal levensvatbaar is, en een praktisch alternatief biedt voor de verstrengeling van effecten die men vindt in standaard dichte LLMs.