Technische Samenvatting: Prototype Transformer (ProtoT)

Probleemstelling

State-of-the-art autoregressieve taalmodellen (LM's), zoals GPT-4 en LLaMA, behalen superieure prestaties in diverse domeinen, maar lijden aan een gebrek aan transparantie in hun redeneerprocessen. Hun interne computationele paden zijn grotendeels opaak, wat leidt tot uitdagingen in het begrijpen van hoe outputs worden gegenereerd, risico's op hallucinaties en moeilijkheden bij gerichte gedragsmodificatie. Huidige interpreteerbaarheidsmethoden zijn primair post-hoc (bijv. aandachtsvisualisatie, probing, sparse autoencoders) en opereren op architecturen die niet zijn ontworpen met interpreteerbaarheid als primaire beperking. Bovendien brengen standaard self-attention mechanismen een kwadratische computationele kost met zich mee ten opzichte van de sequentielengte, wat de schaalbaarheid beperkt.

Methodologie: De Prototype Transformer (ProtoT)

De auteurs introduceren ProtoT, een autoregressieve LM-architectuur die het standaard self-attention mechanisme met kwadratische kosten vervangt door een prototype-gebaseerd communicatiemodule met lineaire kosten. ProtoT behoudt de algemene structuur van LLaMA-3 (inclusclusief SwiGLU feed-forward netwerken en RMS pre-layernorm) maar verandert fundamenteel de "mixer" laag.

Kernarchitectuur

1. Prototype-gebaseerde Mixer: In plaats van self-attention gebruikt ProtoT $R$ leerbare parametervectoren genaamd prototypes. Deze prototypes fungeren als filters die $R$ onafhankelijke communicatiekanalen definiëren.
2. Write en Read Gates:
   • Write Gate: Aggregeert informatie van voorgaande tokens ($x_{j}, j < i$) in de $R$ kanalen. Het berekent gelijkenisscores tussen inputtokens en prototypes, waarbij een softmax over de prototypes wordt gebruikt om informatie naar specifieke kanalen te routeren. Dit creëert een "semantische routeringsbottleneck" waarbij meerdere betekenissen worden gedwongen tot een compromisrepresentatie per kanaal, wat semantische specialisatie stimuleert.
   • Read Gate: Haalt informatie op uit de $R$ kanalen op basis van de gelijkenis tussen het huidige token en de prototypes.
3. Prefix Mean (PM) & Time Discounting: De geaggregeerde informatie in elk kanaal wordt onderhouden als een prefix mean van voorgaande activaties. Een leerbare tijd-discount parameter ($\beta_k$) past een Exponential Moving Average (EMA) toe op de prefix mean, waardoor verschillende prototypes op verschillende tijdschalen kunnen opereren (kortetermijn versus langetermijn afhankelijkheden).
4. Lineaire Complexiteit: De architectuur bereikt een lineaire computationele complexiteit ($O(R)$) ten opzichte van de sequentielengte, aangezien prototypes niet met elkaar interageren (in tegen tegenstelling tot de $O(R^2)$ interactie in Perceiver-stijl latente ruimtes) en de recursie zorgt voor constante-tijd token generatie ($O(1)$) via caching.
5. Alpha Gate: Een ReZero-achtige scalaire gate wordt toegepast op de output van elke mixer module om de trainingsstabiliteit te debuggen en te waarborgen dat de mixer bijdraagt aan de voorspellingen.

Belangrijkste Bijdragen

1. Architecturale Innovatie: De introductie van ProtoT, een nieuwe LM-architectuur die self-attention vervangt door prototype-gebaseerde multi-channel communicatie, waarbij een lineaire computationele complexiteit wordt bereikt terwijl competitieve prestaties worden behouden.
2. Emergente Interpreteerbaarheid: De demonstratie dat prototypes automatisch ontvlezelde, benoembare concepten leren (bijv. "vrouw", "Nieuw-Zeeland", "COVID-19") tijdens de training over verschillende abstractieniveaus. Dit gebeurt zonder expliciete supervisie, gedreven door de architecturale druk om kanalen te specialiseren.
3. Gerichte Gedragsmodificatie: Bewijs dat het redeneren van het model chirurgisch kan worden aangepast door te interveniëren op specifieke prototypes (bijv. het herinitialiseren of maskeren van een "vrouwelijke" prototype). Deze interventies veranderen de waarschijnlijkheid van gerelateerde tokens aanzienlijk (bijv. het verlagen van "vrouwen" met ~16%) terwijl de algehele perplexiteit van het model nauwelijks degradeert (<1% in de meeste gevallen).
4. Robuustheidsanalyse: Het aantonen dat ProtoT robuust is tegen betekenisbehoudende lexicale perturbaties (synoniemen, typefouten) en dat deze robuustheid wordt gemedieerd door de prototype routeringsmechanismen.

Experimentele Resultaten

De auteurs evalueerden ProtoT tegen drie baselines: LLaMA (standaard self-attention), Mamba (state-space model) en DeltaNet (lineaire attention), met gebruik van een vaste backbone (6 lagen, 256 hidden size) en getraind op 250M tokens van FineWeb-Edu.

• Schaalbaarheid: ProtoT schaalt goed met modelgrootte en datavolume. In grootschalige settings (512 hidden size, 12 lagen) bereikt ProtoT een perplexiteit van 29,5, waarmee het DeltaNet (31,5) overtreft, hoewel het achterblijft bij LLaMA (25,8).
• Tekstgeneratie: In gepaarde menselijke voorkeursevaluaties (LLM-als-rechter) presteert ProtoT beter dan LLaMA en DeltaNet, en komt dicht in de buurt van Mamba.
• Downstream Taken: Op de GLUE benchmark blijft ProtoT competitief met dense baselines, waarbij het vaak Mamba en DeltaNet evenaart of overtreft, en sterke cross-domein robuustheid demonstreert.
• Doorvoersnelheid (Throughput): ProtoT schaalt beter dan LLaMA met toenemende contextlengte vanwege de lineaire compute, en overtreft LLaMA bij 32k tokens, hoewel DeltaNet de hoogste doorvoersnelheid behoudt bij zeer lange contexten.
• Interpreteerbaarheidsmetrieken: ProtoT prototypes scoren hoger op ontvlechting (disentanglement) en dekking (coverage) van hoofdthema's vergeleken met LLaMA attention heads en Sparse Autoencoder (SAE) features.
• Effectiviteit van Interventie: Het verstoren van specifieke prototypes (bijv. L9 P7 voor "vrouw") veroorzaakte significante dalingen in de doel-token waarschijnlijkheden (bijv. -16,60% voor "vrouwen") met minimale nevenschade aan de sequentiecoherentie.

Betekenis en Claims

Het artikel claimt dat ProtoT een veelbelovende stap is naar autoregressieve taalmodellen die interpreteerbaar zijn door ontwerp. Door interpreteerbaarheid direct in de architectuur te integreren via prototype kanalen, kan het model:

• Coherente, benoembare concepten automatisch vastleggen.
• Gerichte, chirurgische wijzigingen aan het gedrag van het model aanbrengen zonder het hele systeem te destabiliseren.
• Een alternatief met lineaire compute bieden voor self-attention dat robuust is tegen perturbaties.

De auteurs erkennen beperkingen, waarbij zij opmerken dat hoewel de prestaties competitief zijn met lineaire baselines, het momenteel onder de gevestigde self-attention modellen zoals LLaMA ligt qua ruwe perplexiteit. Ze waarschuwen ook dat prototype labels als benaderingen moeten worden beschouwd en niet als de grondwaarheid, aangezien sommige prototypes polysemantisch blijven. Desalniettemin suggereert het werk dat het balanceren van capaciteit met transparantie haalbaar is, wat een pad biedt naar veiligere en meer controleerbare taalmodellen.