Technische Samenvatting: Het heroverwegen van de Vormconventie van Transformer Taalmodellen

Probleemstelling

Dense Transformer-taalmodellen hebben historisch gezien vastgehouden aan een rigide architecturale conventie: het Feed-Forward Network (FFN) maakt gebruik van een "smal–breed–smal" Multilayer Perceptron (MLP) vorm. In dit ontwerp wordt de modeldimensie ($d_{model}$) uitgebreid naar een bredere intermediaire dimensie ($d_h$) met een expansieratio die doorgaans tussen de 2 en 4 ligt, voordat deze weer wordt geprojecteerd. Deze structuur domineert de parameterallocatie en beslaat vaak het grootste deel van de parameters en floating-point operaties van een model.

De auteurs stellen de vraag of deze specifieke vorm essentieel is of slechts een gemakkelijke instantiatie. Recent theoretisch werk suggereert dat residuele "breed–smal–breed" (hourglass) MLPs als optimale functiebenaderingen kunnen dienen, wat potentieel superieure expressiviteit biedt met minder parameters. Bovendien heeft empirisch bewijs in generatieve visiemodellen en Mixture-of-Experts (MoE) Transformers aangetoond dat hourglass-structuren de standaardontwerpen kunnen overtreffen. Het kernprobleem dat wordt geadresseerd is of de conventionele smal–breed–smal FFN een noodzakelijke beperking is of dat een diepere, smallere hourglass FFN een efficiëntere allocatie van capaciteit tussen aandacht- en FFN-modules kan bieden.

Methodologie

Het paper stelt een Transformer-variant voor die het conventionele FFN vervangt door een Hourglass FFN. De methodologie omvat:

1. Architecturale Herontwerp:

   • De Hourglass FFN vervangt de enkele brede expansielaag door een stapel van $K$ hourglass-vormige sub-MLP's die verbonden zijn via residuele paden.
   • Elke sub-MLP ($MLP_i$) volgt een compressie-expansiestructuur: het projecteert de input van $d_{model}$ naar een bottleneck-dimensie $d_h$ (waarbij $d_h < d_{model}$), past een niet-lineaire activatie toe (SwiGLU), en projecteert terug naar $d_{model}$.
   • De output van de laag is het resultaat van het passeren van de input door $K$ dergelijke residuele sub-blokken.

2. Strategie voor Parameterherverdeling:

   • Door de intermediaire dimensie $d_h$ te verkleinen (het FFN "smaller" maken), besparen de auteurs parameters die gewoonlijk aan de breedte van de FFN worden toegewezen.
   • Deze bespaarde parameters worden geheralloceerd naar:
     • Het vergroten van de modeldimensie ($d_{model}$) en de aandachtscapaciteit.
     • Het vergroten van de diepte van de FFN-stapel ($K$).
     • Het vergroten van het aantal Transformer-lagen ($L$).

3. Experimentele Opzet:

   • Baselines: Vergelijkingen worden gemaakt met standaard LLaMA-stijl Transformers en de state-of-the-art OLMo-2 architectuur.
   • Schalen: Experimenten worden uitgevoerd over modelgroottes variërend van 113M tot 1B parameters.
   • Restricties: Experimenten worden uitgevoerd onder gematchte totale parameterbudgetten (binnen 0,001% verschil) en vaste aandachtstijden om de effecten van de FFN-vorm te isoleren.
   • Zoekruimte: De auteurs voeren een architecturale zoektocht uit waarbij $K$ (diepte van de hourglass stapel), $L$ (aantal Transformer-lagen), $d_{model}$ en de ratio $d_h/d_{model}$ variëren.

Belangrijkste Bijdragen

Het paper levert vier primaire bijdragen:

1. Het Uitdagen van de Standaard: Het daagt de de facto "smal–breed–smal" FFN-vorm uit en stelt een efficiënte "breed–smal–breed" Hourglass FFN voor die expressiviteit behoudt door middel van gestapelde, residueel verbonden sub-MLP's.
2. Parameterherallocatie: Het demonstreert een strategie waarbij het herdirigeren van parameters van FFN-breedte naar aandachtbreedte of interne FFN-diepte ($K > 1$) superieure prestaties oplevert, met name bij kleinere schalen (113M).
3. Architecturale Inzichten: Het identificeert een U-vormige breedte-diepte trade-off, vindt optimale ratio's voor $d_{model}/L$ en $d_h/d_{model}$, en stelt vast dat diepere Hourglass FFN-structuren robuust zijn.
4. Validatie van Schaalbaarheid: Het valideert de Hourglass FFN van 113M tot 1B parameters, waarbij wordt aangetoond dat het een competitief alternatief is voor standaard baselines, met specifieke configuraties die conventionele modellen overtreffen op kleinere schalen en deze evenaren op grotere schalen.

Resultaten

De empirische validatie levert de volgende bevindingen op:

• Prestaties bij Vaste Aandacht: Op de 113M-schaal met vaste aandachtparameters bereiken Hourglass FFN's een lagere validatieverlies en perplexiteit (PPL) dan conventionele FFN's. Bijvoorbeeld, bij $K=8$ daalde de PPL van 36,441 (conventioneel) naar 36,179.
• Optimale Configuraties:
  • Diepte vs. Breedte: Een ratio van $d_h/d_{model} \approx 0,4$ met $K=4$ lagen leverde het laagste validatieverlies op.
  • Parameterverschuiving: Bij 113M resulteerde het verschuiven van parameters om het aandachtsmechanisme dominant te maken (bijv. 66M aandacht vs. 47M FFN) in de beste prestaties (PPL 35,101), wat suggereert dat contextuele verwerking waardevoller is dan statische feitelijke extractie voor kleine modellen.
  • Breedte-Diepte Trade-off: Er werd een duidelijke U-vormige curve waargenomen voor de $d_{model}/L$ ratio, met een optimaal 'sweet spot' tussen 100 en 250.
• Schaalbaarheid:
  • Kleine tot Middelgrote Schalen (113M – 906M): Hourglass-varianten presteerden consistent beter dan conventionele baselines in termen van validatieverlies, perplexiteit en downstream redeneertaken (bijv. Arc Easy, HellaSwag, PIQA).
  • Grote Schaal (1B): Op de 1B-parameter schaal bereikte de Hourglass-variant een prestatie die vergelijkbaar is met de sterke conventionele baseline (Val PPL 20,082 vs. 20,002) en overtrof de OLMo-2 configuratie (20,466).
• Ablatie-studies: Het verhogen van de hourglass-diepte $K$ verbeterde consequent de prestaties, hoewel met afnemende meeropbrengsten boven $K=6$. Het verlagen van $d_h/d_{model}$ onder 0,4 leidde tot significante prestatieverslechtering, wat wijst op een ondergrens voor FFN-compressie.

Betekenis en Claims

Het paper claimt dat de langdurige "smal–breed–smal" MLP-vorm in Transformers niet essentieel is en mogelijk redundant is. Door een Hourglass FFN te adopteren, demonstreren de auteurs dat:

1. Efficiëntie: Modellen een competitieve of superieure prestatie kunnen bereiken met een efficiëntere parameterverdeling.
2. Resourceherallocatie: De Hourglass-structuur maakt een fundamentele verschuiving mogelijk waarbij het aandachtsmechanisme de dominante consument van parameters kan worden, wat het traditionele FFN-gedomineerde budget uitdaagt.
3. Herwaardering van Ontwerp: De bevindingen stimuleren een heroverweging van de balans tussen aandacht en FFN, en suggereren dat voor efficiënte en expressieve moderne taalmodellen, diepere, smallere FFN's gecombineerd met bredere aandachtmechanismen een superieur ontwerpplan kunnen zijn.

De auteurs blijven bescheiden over de schaalbaarheid naar modellen met tientallen of honderden miljarden parameters, waarbij zij opmerken dat hun bevindingen zijn geëxtrapoleerd van de 113M-schaal en dat de "vloek van diepte" (signaalpropagatieproblemen) verdere investigatie vereist in diepere architecturen. Ze merken ook op dat de interactie tussen Hourglass FFN's en geavanceerde aandachtmechanismen (bijv. Group-Query Attention) nog onverkend is.