Technische Samenvatting: CoFrGeNet – Kettingbreuk-architecturen voor taalgeneratie

Probleemstelling

Transformers zijn de dominante architectuur geworden voor taalgeneratie, waarbij ze vertrouwen op Multi-head Attention en Feed-Forward Networks (FFN's) met grote verborgen lagen. Deze componenten introduceren echter aanzienlijke rekenkundige overhead, vooral wat betreft het aantal parameters en de trainingsduur. Hoewel er diverse benaderingen zijn voorgesteld om de efficiëntie te verbeteren (bijvoorbeeld sparse attention, linear attention), blijft er behoefte aan fundamenteel verschillende functieklassen die deze kerncomponenten kunnen vervangen terwijl ze prestaties behouden of verbeteren.

Vorig werk introduceerde CoFrNets (Continued Fraction Networks) voor supervised learning, waarbij gebruik wordt gemaakt van een ladderachtige structuur die is geïnspireerd op kettingbreuken. Het aanpassen van deze architectuur voor generatieve modellering stelt echter drie specifieke uitdagingen:

1. Meerdimensionale output: Generatieve modellen vereisen het produceren van meerdimensionale output, in tegenstelling tot de scalaire output die vaak wordt gezien in supervised settings.
2. Rijkere functielering: De architectuur moet complexe representaties leren die geschikt zijn voor sequentiemodellering.
3. Causale beperkingen: Het model moet parameters leren die alleen afhankelijk zijn van eerdere tokens om coherente generatie te waarborgen.
4. Rekenkundige inefficiëntie: De standaard $1/x$-niet-lineariteit in kettingbreuken vereist $d$ delingen voor een ladder met diepte $d$. Aangezien deling aanzienlijk langzamer is dan vermenigvuldiging op moderne hardware, ontstaat hierdoor een bottleneck voor zowel training als inferentie, vooral naarmate de diepte ($d$) en het aantal ladders ($L$) toenemen.

Methodologie

1. Continued Fraction Generative Networks (CoFrGeNets)

De auteurs stellen CoFrGeNets voor, een familie van architecturen die Attention en/of FFN-lagen in Transformer-blokken vervangen door componenten gebaseerd op kettingbreuken. De kernfunctieklasse is gedefinieerd als:
$$w_0x + \frac{1}{w_1x + \frac{1}{w_2x + \dots}}$$
waarbij $x$ de input is en $w_k$ leerbare parametervectoren zijn.

2. Architecturale Componenten

Het artikel introduceert specifieke ontwerpen voor het vervangen van Transformer-componenten terwijl causale beperkingen worden gehandhaafd:

• Causale Attention-vervanging (CAttnU en CAttnM):

  • CAttnU: Gebruikt univariate CoFrNet-ladders waarbij inputs worden getransponeerd om dimensies over tokens te mengen. Om causaliteit af te dwingen (het voorkomen van lekken van toekomstige tokens), maakt de architectuur gebruik van bovendriehoekige lineaire lagen binnen de ensembles. Elementsgewijze vermenigvuldiging van outputs van twee ensembles genereert kruistermen, wat het representatievermogen versterkt zonder de causale stroom te schenden.
  • CAttnM: Gebruikt $L$ CoFrNet-ladders om inputs af te beelden op embeddings met de lengte van de sequentie, die dienen als attention-weights. Causaliteit wordt gehandhaafd door een causale softmax (convexe combinatie van alleen eerdere tokens) toe te passen voordat de waardenmatrix wordt gewogen.

• FFN-vervanging (Cffn):

  • Omdat FFN's feature-mixing vereisen in plaats van token-mixing, maakt de architectuur gebruik van ensembles van p-variate ladders (waarbij $p$ de embedding-dimensie is) zonder de input te transponeren.
  • Een lineaire laag aan het einde projecteert de output terug naar de doel-dimensie.
  • Cruciaal is dat de input voor deze ladders een gegateerde, niet-uitgebreide representatie is ($\alpha=1$), waardoor de parameterexplosie wordt vermeden die typisch is voor standaard Transformer FFN's, die vaak een uitbreidingsfactor $\alpha \gg 1$ gebruiken.

3. Efficiënte Implementatie via Continuants

Om de rekenkundige inefficiëntie van herhaalde delingen aan te pakken, leiden de auteurs een aangepaste implementatie af op basis van continuants (polynomen $K_k$).

• Forward Pass: In plaats van de breuk laag-voor-laag te berekenen met $d$ delingen, wordt de kettingbreuk berekend als de verhouding van twee polynomen: $\tilde{f}(a) = K_{d-1} / K_d$. Dit vereist het recursief berekenen van de polynomen met uitsluitend vermenigvuldigingen en optellingen, gevolgd door een enkele deling ($1/K_d$).
• Backward Pass: De auteurs leiden gesloten-vorm expressies af voor de gradiënten met behulp van de eigenschappen van continuants. De gradiënt $\frac{\partial \tilde{f}}{\partial a_k}$ wordt uitgedrukt als een verhouding van gekwadrateerde continuants.
• Optimalisatie: Door $K_d$ één keer te berekenen en $1/K_d$ te hergebruiken voor alle gradiëntberekeningen, wordt het aantal delingen gereduceerd van $d$ naar 1. Dit versnelt zowel training als inferentie aanzienlijk.

4. Aangepast Trainingschema

De auteurs stellen een dyadisch parameter-update-schema voor om training te stabiliseren. In plaats van alle parameters gelijktijdig te updaten:

• Lineaire componenten worden vanaf het begin bijgewerkt.
• Parameters op hogere dieptes worden aanvankelijk bevroren.
• Updates voor parameters op diepte $i$ beginnen pas na $t/2^i$ iteraties (waarbij $t$ het totale aantal iteraties is).
  Deze incrementele aanpak voorkomt instabiliteit tijdens de vroege fasen van training.

Belangrijkste Bijdragen

1. Nieuwe Architecturen: Introductie van CoFrGeNet-componenten als plug-in vervangingen voor Attention- en FFN-lagen in Transformers, wat een nieuwe functieklaas biedt voor generatieve modellering.
2. Efficiënte Gradiëntformulering: Afleiding van aangepaste gradiëntformules met behulp van continuants die het aantal delingsoperaties reduceren van $d$ naar 1 per ladder, wat de hardware-efficiëntie drastisch verbetert.
3. Trainingsstabiliteit: Ontwikkeling van een aangepast dyadisch trainingsschema dat de training van diepe kettingbreuk-architecturen stabiliseert.
4. Empirische Validatie: Pre-training en evaluatie op grote schaalmodellen (GPT2-xl en Llama3) die concurrerende prestaties tonen met een verminderd aantal parameters.

Experimentele Resultaten

Opzet

• Modellen: Experimenten werden uitgevoerd op GPT2-xl (1.5B), voorgetraind op OpenWebText en GneissWeb, en Llama3 (3.2B), voorgetraind op de docling data mix (2T tokens).
• Basislijnen: Vergelijkingen omvatten de originele Transformers, Dense Synthesizer en Sparse Attention.
• Varianten: CoFrGeNet-F (FFN vervangen), CoFrGeNet-A (Attention vervangen) en CoFrGeNet (beide vervangen).

Prestaties

• Downstream Taken (GLUE): CoFrGeNet-modellen, ondanks dat ze 2/3 tot 1/2 het aantal parameters hebben van de originele modellen, behaalden concurrerende of superieure nauwkeurigheid op classificatie-, Q&A- en redeneertaken. Bijvoorbeeld, CoFrGeNet-F (985M) presteerde beter dan de 1.5B GPT2-xl op de meeste GLUE-benchmarks.
• Perplexiteit: Op taalmodellering-benchmarks (PTB, Wikitext, Lambada) behaalden CoFrGeNet-varianten over het algemeen een lagere perplexiteit dan de basislijn-Transformers en andere efficiënte attention-basislijnen.
• Efficiëntie:
  • Trainingsduur: De op continuants gebaseerde implementatie verkleinde de trainingsduur aanzienlijk in vergelijking met een "letterlijke" implementatie (CoFrGeNetB) die geen continuants gebruikte. Voor GPT2-xl daalde de trainingsduur van ~203 uur (B) naar ~178 uur (CoFrGeNet).
  • Inferentietijd: De inferentiesnelheid verbeterde met bijna een orde van grootte voor de op continuants gebaseerde versie in vergelijking met de letterlijke implementatie.
  • Doorvoer: Op Llama3 toonden CoFrGeNet-varianten een hogere tokens/day-doorvoer en een verminderde totale trainingsduur (bijvoorbeeld CoFrGeNet-F getraind in 6,6 dagen versus 8,5 dagen voor de basislijn).

Betekenis en Claims

Het artikel stelt dat CoFrGeNets een levensvatbaar, parameter-efficiënt alternatief vertegenwoordigen voor standaard Transformer-componenten. Door gebruik te maken van de wiskundige eigenschappen van kettingbreuken en hun implementatie te optimaliseren via continuants, tonen de auteurs aan dat:

• Hoogpresterende taalmodellen kunnen worden gebouwd met aanzienlijk minder parameters.
• De rekenkundige bottleneck van delingsoperaties in kettingbreuken effectief kan worden opgelost, waardoor de architectuur praktisch wordt voor grootschalige industriële workflows.
• De aanpak "plug-and-play" is, met minimale wijzigingen in bestaande trainings- en inferentiepijplijnen.

De auteurs concluderen dat dit werk de grondslag legt voor door kettingbreuken geïnspireerde generatieve architecturen, wat potentieel AI kan democratiseren door entiteiten met minder middelen in staat te stellen hoogwaardige modellen voor te trainen. Zij merken echter op dat deze modellen geen inherente veiligheidsmechanismen bezitten en vatbaar blijven voor hallucinaties en adversariale aanvallen, wat suggereert dat toekomstig onderzoek zich moet richten op het aanpakken van deze uitdagingen binnen de specifieke functionele vorm.