Geometry-Guided Layerwise FFN Width Allocation in Transformers

Technische Samenvatting: Geometrie-gestuurde Laagspecifieke FFN-breedte Allocatie in Transformers

Probleemstelling
In standaard Transformer-architecturen is de verborgen breedte van Feed-Forward Networks (FFNs) doorgaans constant over alle lagen, ondanks bewijs dat laagfuncties en representatiedynamiek veranderen met de diepte. Hoewel recent werk (Tapered Language Models) heeft aangetoond dat het heralloceren van capaciteit van diepere naar ondiepere lagen de prestaties verbetert, vertrouwen bestaande methoden op handmatig ontworde, monotone schema's (bijv. cosinus-tapers). Deze voorgeschreven vormen passen zich niet aan de specifieke computationele eisen van een bepaand model of datadistributie. De paper vraagt zich af of een datarijke, laagspecifieke breedte-allocatie direct kan worden afgeleid uit het geometrische gedrag van activaties tijdens een forward pass, in plaats van te vertrouwen op een vaste architecturale aanname.

Methodologie
De auteurs stellen een framework voor dat FFN-lagen behandelt als transporteurs van de wolk van token-representaties. De kernmethodologie omvat drie stadia:

1. Meting van Geometrische Arbeid:
   De paper definieert "geometrische arbeid" als de verandering die een FFN-laag induceert op een wolk van token-representaties. Om ervoor te zorgen dat de metriek de werkelijke bijdrage van de FFN reflecteert in plaats van artefacten van residuele schaling, evalueren de auteurs drie verschillende arbeid-statistieken onder drie geometrische kaders:

   • Ruwe Geometrie (Raw Geometry): Gebruikt Euclidische afstand op ongenormaliseerde toestanden. Dit is gevoelig voor schaal en dient als controle.
   • Sferische Geometrie: Gebruikt eenheids-genormaliseerde toestanden en hoekafstand.
   • Hyperbolische Geometrie: Brengt eenheids-genormaliseerde toestanden in een Poincaré-bal, wat fungeert als een niet-lineaire herweging van hoekafstanden.

   Voor elke geometrie worden drie arbeid-statistieken berekend:

   • Verschuiving (Shift): De gemiddelde verplaatsing van tokens (correspondentie-behoudend).
   • Metrische Distorsie (Metric Distortion): Gemeten via de Gromov-Wasserstein (GW) afstand, wat kwantificeert hoe paarrelaties behouden blijven of worden gewijzigd.
   • Topologische Verandering (Topological Change): Gemeten via de 2-Wasserstein afstand tussen graad-één persistente homologie-diagrammen (Vietoris–Rips filtratie), wat de creatie of destructie van lussen vastlegt.

2. Approximatie-surrogaat en Optimalisatie:
   De auteurs introduceren een surrogaatmodel om geometrische arbeid te mappen naar optimale laagbreedtes. Zij postuleren dat de approximatie-fout van een FFN-laag schaalt met de breedte volgens een intrinsieke dimensie ($d_{int}$) en een loss-sensitiviteit ($\lambda$).

   • Aannames: (1) Loss-effecten zijn lokaal additief; (2) Fout schaalt als $w^{-a}$ waarbij $a$ afhankelijk is van de intrinsieke dimensie; (3) Gemeten geometrische arbeid is een proxy voor de moeilijkheid van approximatie.
   • Optimizer: Een geconstraineerd optimalisatieprobleem wordt opgelost om de som van gewogen approximatie-fouten te minimaliseren onder een vast totaal parameterbudget. Dit levert een gesloten vorm oplossing op voor de optimale breedte $w_\ell$ per laag, waarbij rekening wordt gehouden met ondergrenzen en diepte-afhankelijke exponenten.

3. Schema Constructie:
   Een referentiemodel (uniforme breedte) wordt uitgevoerd om activatiewolken te verzamelen. Het geometrische arbeidsprofiel wordt geschat, gladgestreken en gemapt naar een breedte-schema met behulp van de afgeleide optimizer. Dit schema wordt vervolgens toegepast om een nieuw model vanaf nul te trainen.

Belangrijkste Bijdragen

1. Definitie van Geometrische Arbeid: De paper formaliseert drie complementaire forward-pass arbeid-statistieken (verschuiving, GW-distorsie, topologische verandering) onder ruwe, sferische en hyperbolische metrieken, waarbij onderscheid wordt gemaakt tussen ruwe magnitude en geschaalde structurele veranderingen.
2. Exact Budget-Optimaal Solver: Het leidt een unieke continue optimizer af voor laagwijze breedte-allocatie onder een vast parameterbudget, waarbij rekening wordt gehouden met ondergrenzen en variërende intrinsieke dimensies per laag.
3. Uitgebreide Diagnostiek: De studie evalueert schaalgevoeligheid, lokale loss-additiviteit, intrinsieke dimensie en de correlatie tussen geometrische arbeid en perturbatie-gebaseerde laag-sensitiviteit over zeven vooraf getrainde modellen (Llama 3.1/3.2, Gemma 2/3, Mistral).
4. Empirische Validatie: De paper vergelijkt geometrie-afgeleide schema's met uniforme breedte, een handmatig ontworpen cosinus-taper en een "anti-topologische" controle over drie model-schalen (128M, 256M en 440M).

Resultaten

• Schaalgevoeligheid: Ruwe geometrische arbeid (Euclidisch) volgt grotendeels de groei van de residuele norm, terwijl genormaliseerde metrieken (sferisch/hyperbolisch) een "front-loaded" profiel onthullen waarbij de vroegste lagen de hoogste geometrische arbeid vertonen.
• Sensitiviteit vs. Allocatie: Terwijl Gromov-Wasserstein arbeid een sterkere correlatie vertoont met perturbatie-gebaseerde laag-sensitiviteit, levert de topologische/hyperbolische allocatie de laagste gemiddelde validatieverlies op tijdens trainingsruns. Dit suggereert dat sensitiviteit en approximatie-moeilijkheid verschillende factoren zijn.
• Trainingsprestaties:
  • In 128M en 256M runs presteren genormaliseerde geometrie-schema's over het algemeen beter dan uniforme breedte en de anti-topologische controle, hoewel de resultaten verkennend zijn vanwege beperkte seeds.
  • In de 440M vijf-seeds experiment bereikte het topologische/hyperbolische schema een gemiddelde reductie in validatieverlies van 0.019 ten opzichte van uniforme breedte. Dit is ongeveer 6,3 keer de reductie van de TLM-stijl cosinus-taper (0.003).
  • De anti-topologische/ruwe controle presteerde slechter dan uniforme breedte, wat bevestigt dat de richting van de allocatie ertoe doet.

Betekenis en Claims
De paper claimt dat geometrische arbeid een principieel, datagestuurd alternatief biedt om globale FFN-breedte constanten te vervangen door laagafhankelijke schema's. Het stelt vast dat:

• FFN-capaciteit niet als uniform moet worden aangenomen; laag-updates veranderen van rol naarmate de diepte toeneemt.
• Vaste monotone tapers niet de enige effectieve alternatieven zijn; activatie-afgeleide profielen kunnen niet-monotone structuren vangen en grotere prestatiewinsten behalen dan handontworpen schema's onder identieke condities.
• Genormaliseerde geometrische veranderingen (met name topologische/hyperbolische) een overdraagbaar "early-heavy" profiel onthullen dat de efficiëntie van modellen verbetert.

De auteurs blijven bescheiden in hun conclusies en merken op dat de resultaten beschrijvende gemiddelde vergelijkingen zijn totdat de onzekerheid van gepaarde verschillen volledig is gerapporteerd. Ze beweren niet dat ze de universele superioriteit van een specifieke geometrie (bijv. topologie vs. GW) hebben bewezen, maar dat genormaliseerde geometrische verandering een nuttig signaal is voor capaciteitsallocatie. Dit werk positioneert activatie-afgeleide allocatie als een praktische uitbreiding van tapering, wat toekomstig werk motiveert om geometrische arbeid direct te kalibreren aan de laagwijze vraag naar approximatie.