Technische Samenvatting: TriRoute

Probleemstelling

Het implementeren van Large Language Models (LLMs) brengt aanzienlijke inferentiekosten met zich mee, maar de huidige toewijzing van middelen is statistisch gezien verspillend. Dichte transformers passen uniforme floating-point operaties (FLOPs) en Key/Value (KV) cachegeheugen toe op elke token, ongeacht de semantische belangrijkheid. Een syntactisch triviaal functiewoord verbruikt evenveel middelen als een zeldzame entiteit die de betekenis van een paragraaf bepaalt.

Bestaande technieken voor conditionele berekening proberen deze uniformiteit te doorbreken, maar opereren geïsoleerd op afzonderlijke assen:

1. Mixture-of-Experts (MoE): Maakt de Feed-Forward Network (FFN) ijl door tokens naar een subset van experts te routeren.
2. Mixture-of-Depths (MoD): Slaat volledige transformer-blokken (attention + FFN) over voor tokens die als voorspelbaar worden beschouwd.
3. KV-Cache Kwantisatie: Comprimeert de attention-geheugens door keys en values op te slaan met een lagere bitbreedte.

Het artikel stelt dat deze beslissingen sterk gekoppeld zijn. Zo kan een zeldzame entiteit-token bijvoorbeeld volledige attention-resolutie en een hoge-precisie cache-opslag vereisen om getrouw te kunnen worden opgehaald door toekomstige queries, terwijl de FFN-transformatie ervan voorspelbaar en overgeslagen kan worden. Omgekeeld kan een functiewoord in alle drie de dimensies overgeslagen worden. Het behandelen van deze als onafhankelijke optimalisatieproblemen met aparte hyperparameters faalt in het vastleggen van deze correlaties, wat leidt tot suboptimale toewijzing van middelen en potentieel verlies aan "tail-case" robuustheid (prestaties op zeldzame entiteiten, code en rekenkunde).

Methodologie: TriRoute

De auteurs stellen TriRoute voor, een verenigde architectuur die drie geïsoleerde mechanismen vervangt door één enkele, lichtgewicht geleerde controller. Voor elke token $t$ in elke laag $\ell$ genereert de controller een gecoördineerd beleid $\pi^{(\ell)}_t = (a_t, e_t, b_t)$ over drie heterogene assen:

1. Attention Resolutie ($a_t$): Kiest tussen skip (0 kosten), lokaal (windowed attention), of volledig (dense causal attention).
2. Expert Selectie ($e_t$): Routeert de token naar een ijle set FFN-experts of een null-expert ($f_0 \equiv 0$). De null-expert stelt het model in staat om de FFN over te slaan terwijl de volledige attention behouden blijft, een regime dat MoD niet kan uitdrukken.
3. KV-Cache Bitbreedte ($b_t$): Bepaalt de precisie (2, 4, 8 of 16 bits) waarmee de eigen Key/Value-paren van de token worden geschreven naar de cache voor toekomstige queries.

Technische Componenten

• Gedeelde Controller: Een twee-laagse MLP-trunk verwerkt een genormaliseerde residuele staat en goedkope causale zij-features (positie, afstand, predictieve entropie) om logits te genereren voor drie afzonderlijke koppen (Attention, Experts, Bits). Dit ontwerp maakt een gedeeld begrip van "token-belangrijkheid" mogelijk om alle assen te informeren, terwijl aparte output-subspaces behouden blijven.
• Heterogene Relaxatie: Training vereist differentiatie door discrete beslissingen. De auteurs gebruiken:
  • Gumbel-Softmax met Straight-Through (ST) Estimators voor de categorische attention- en bitbreedte-beslissingen.
  • Top-k Softmax Gating voor expert-selectie.
  • Cross-Axis Gradient Balancing: Omdat het overslaan van attention een veel grotere impact heeft op de loss dan het verminderen van cache-bits, introduceren de auteurs een per-as gradiënt-rescaling mechanisme (met behulp van een EMA van gradiënt-normen) om te voorkomen dat de attention-as de gedeelde trunk domineert.
• Koppeling-bewuste Balans: Om een "collapse cascade" te voorkomen (waarbij een instorting in één as leidt tot instorting in andere assen door inputs met lage variantie), gebruikt de methode:
  • Per-as Whitening: Detached normalisatie van trunk-features vóór elke kop.
  • Entropy Floor: Een penalty die alleen wordt geactiveerd wanneer de entropie van het gebruik van opties onder een drempelwaarde daalt, wat diversiteit waarborgt zonder de budgetrestrictie te bevechten.
• Differentiabel Kostenmodel & Lagrangian Controle: Eén enkele Lagrangian-restrictie dwingt een gemiddeld doelbudget $C^\star$ af dat zowel FLOPs als geheugen beslaat. De duale variabele (Lagrange-multiplier) wordt online bijgewerkt, waardoor gebruikers de volledige Pareto-frontier kunnen verkennen met een enkele scalaire knop in plaats van drie afzonderlijke sparsity-hyperparameters af te stemmen.

Belangrijkste Bijdragen

1. Verenigde Formulering: TriRoute is de eerste architectuur die attention-resolutie, expert-selectie en KV-cache precisie gezamenlijk onder één controller en budget leert.
2. Stabiel Joint Training Recept: Het papier identificeert en mitigeert de "cross-axis routing collapse cascade" door middel van per-as whitening, entropy floors en gradiënt-balancering, wat end-to-end training van heterogene discrete beslissingen haalbaar maakt.
3. Design Space Exploratie: De auteurs tonen aan dat een gedeelde trunk met afzonderlijke koppen de optimale granulariteit is, wat beter presteert dan volledig gescheiden routers of één volledig gedeelde kop. Ze laten ook zien dat per-kop attention-routing gecombineerd met token-niveau FFN/bit-routing een "sweet spot" is.
4. Preserveren van Tail-Cases: In tegen tegenover gemiddelde-perplexiteit-geoptimaliseerde adaptieve modellen, die vaak regressie vertonen op zeldzame inputs, behoudt TriRoute robuustheid op zeldzame entiteiten, code en rekenkunde door middelen precies daar in te zetten waar ze nodig zijn.

Resultaten

Geëvalueerd op decoder-only modellen variërend van 160M tot 1.3B parameters, getraind op compute-optimale token-aantallen:

• Pareto Dominantie: TriRoute volgt een kosten-kwaliteit-frontier die de best afzonderlijk afgestemde combinatie van MoD + MoE + KV-kwantisatie strikt domineert. Bij een gelijkwaardig budget van ~55% FLOPs en 40% KV-geheugen, herstelt TriRoute 96–99% van de downstream nauwkeurigheid van het dense model.
• Prestatiewinst: Vergeleken met de onafhankelijke combinatie verbetert TriRoute de perplexiteit met 0.3–0.4 en de downstream nauwkeurigheid met 0.7–1.0 punten over alle schalen.
• Tail Robustness: Op zeldzame entiteit-, code- en reken-buckets vertoonde de onafhankelijke combinatie significante degradatie (bijv. +2.1 perplexiteit op zeldzame entiteiten). TriRoute bleef binnen 0.4 perplexiteit van het dense model, wat aantoont dat gezamenlijke coördinatie "stille degradatie" van hoogwaardige inputs voorkomt.
• Interpretabiliteit: Analyse laat zien dat de controller interpreteerbare patronen leert:
  • Read/Write Fidelity: Zeldzame subwords en benoemde entiteiten ontvangen volledige attention en een hoge-precisie cache.
  • Ontkoppeling: Zeldzame entiteiten slaan vaak de FFN over (null expert) terwijl ze volledige attention en een hoge-precisie cache behouden, een gedrag dat onmogelijk is voor blok-niveau MoD.
  • Diepte Specialisatie: Vroege lagen behouden hoge attention over de hele linie, terwijl latere lagen selectiever worden en attention overslaan voor niet-kritieke tokens.

Betekenis en Claims

Het artikel claimt dat een gedeelde geleerde controller een principiële manier is om een vast inferentiebudget daar te besteden waar het het meest nodig is. Door de drie assen van conditionele berekening te verenigen, vermijdt TriRoute de inefficiënties van handmatig afgestemde, geïsoleerde mechanismen.

De auteurs benadrukken dat hun resultaten suggereren dat de "juiste" hoeveelheid compute niet een enkele scalar per token is, maar een vector van gekoppelde keuzes over heterogene middelen. De verenigde aanpak verbetert niet alleen de Pareto-frontier, maar verandert fundamenteel welke inputs de besparingen betalen, waardoor wordt gegarandeerd dat zeldzame en kritieke tokens niet worden onderbediend.

Het artikel sluit af met een bescheiden noot over praktische implementatie: hoewel de theoretische FLOP-besparingen aanzienlijk zijn, worden de gerealiseerde wall-clock snelheidswinsten momenteel beperkt door systeem-niveau hiaten (ragged attention, mixed-precision cache management en router overhead). De auteurs beschouwen het dichten van deze FLOP-naar-latency kloof als de primaire volgende stap voor praktische adoptie.