LayerRoute: Input-Conditioned Adaptive Layer Skipping via LoRA Fine-Tuning for Agentic Language Models

Technische Samenvatting: LayerRoute – Input-geconditioneerde adaptieve laag-overslag via LoRA Fine-Tuning

Probleemstelling

Moderne agentische AI-systemen orkestreren Large Language Models (LLM's) via heterogene workflows bestaande uit twee verschillende typen stappen:

1. Gestructureerde tool calls: Korte, deterministische sequenties met een lage perplexiteit (bijv. databasequeries).
2. Open-ended planning/redeneren: Lange, complexe sequenties met een hoge perplexiteit die multi-staps redenering vereisen.

Huidige inferentiesystemen passen identieke computationele middelen toe op beide typen stappen, waarbij elke input door alle transformer-lagen wordt gerouteerd. Dit is inefficiënt omdat de hoge mate van vertrouwen in de vroege residu-stroomtoestanden voor tool calls suggereert dat diepe lagen marginaal bijdragen aan de uiteindelijke voorspelling. Het paper adresseert de behoefte aan een inferentiesysteem dat de rekenkracht dynamisch kan toewijzen op basis van de complexiteit van de input zonder de modelkwaliteit op te offeren.

Methodologie

De auteurs stellen LayerRoute voor, een lichtgewicht adapter ontworpen voor het Qwen2.5-0.5B-Instruct model (24 lagen, verborgen grootte 896). De architectuur breidt de bevroren backbone uit met twee primaire componenten per transformer-blok:

1. Per-laag Routers: Elk blok $i$ bevat een lichtgewicht lineaire laag ($w_i^T \bar{h}_i + b_i$) die een zachte sigmoid-waarde $\sigma_i$ produceert. Deze wordt via een drempelwaarde van 0,5 omgezet in een harde binaire gate $g_i \in \{0, 1\}$.

   • Forward Pass: Als $g_i=1$, voert het blok normaal uit (met actieve LoRA-adapters). Als $g_i=0$, passeert de verborgen toestand ongewijzigd, wat nul berekening voor dat blok betekent.
   • Training: De niet-differentieerbare harde gate wordt afgehandeld via de Straight-Through Estimator (STE), waardoor gradiënten door de sigmoid kunnen stromen tijdens backpropagation terwijl harde beslissingen worden behouden tijdens inferentie.

2. LoRA Adapters: Low-rank adapters (Rank 8) worden toegepast op de Query, Key, Value en Output projecties ($W_Q, W_K, W_V, W_O$) van alle 24 blokken. De backbone-gewichten blijven bevroren.

Trainingsstrategie & Innovaties:

• Gezamenlijke Optimalisatie: De router-gewichten en LoRA-matrices worden gezamenlijk getraind met een gecombineerde verliesfunctie: $L = L_{LM} + \lambda \cdot \frac{1}{L}\sum \sigma(s_i)$.
• Gate Regularisatie: Een regularisatieterm ($\lambda=1.0$) straft hoge zachte gate-waarden af, wat voorkomt dat de router convergeert naar een triviale oplossing waarbij alle gates open blijven.
• Gebiasde Initialisatie: Om symmetrie te doorbreken en het "kip-en-ei"-probleem te vermijden (waarbij gates niet kunnen differentiëren totdat ze overslaan), initialiseren de auteurs de routers van de middelste lagen (lagen 8–16) met een bias van $-1.0$ (startend als gesloten/overgeslagen) en de buitenste lagen met $+1.0$ (startend als open). Dit dwingt onmiddellijke differentiatie af op basis van de taalmodellering-loss.

Belangrijkste Bijdragen

Het paper beschrijft vijf specifieke bijdragen:

1. Per-laag hard-gated skip-verbindingen gebruikmakend van STE om de mismatch tussen training en inferentie te elimineren.
2. Gezamenlijke LoRA fine-tuning van aandachtsprojecties naast het leren van de router, waardoor de routing-policy kan co-evolueren met de aangepaste gewichten.
3. Gate regularisatie om het ontdekken van overslagebare blokken af te dwingen in plaats van het behouden van volledige rekenkracht.
4. Gebiasde initialisatie van middelste-laag routers om gradiënt-gebaseerde differentiatie vanaf de eerste trainingsstap mogelijk te maken.
5. Empirische validatie op Qwen2.5-0.5B-Instruct die een significante skip-differentiaal laat zien tussen tool calls en planning-stappen, getraind in minder dan 7 minuten.

Experimentele Resultaten

Het model werd getraind op een gemengde agentische dataset (Hermes, Glaive, GSM8K, Turing) gedurende 3.000 stappen (ongeveer 6,4 minuten op een A100 40GB) met gebruik van slechts 1,10M trainbare parameters (0,22% van de 494M backbone).

• Skip Differential: LayerRoute bereikte een skip differential van 12,91%. Tool calls oversloegen 15,25% van de FLOPs, terwijl planning-stappen slechts 2,34% oversloegen.
• Kwaliteitsverbetering: In tegenstelling tot typische efficiëntie-trade-offs, vertoonde het gated model een verbeterde perplexiteit ten opzichte van de full-layer baseline.
  • Tool calls: Perplexity delta van −1,29.
  • Planning: Perplexity delta van −1,30.
  • Deze verbetering wordt toegeschreven aan de LoRA-adaptatie die de kwaliteitskosten van het overslaan compenseert.
• Ablatieonderzoeken:
  • LayerRoute-NoReg: Deverteerde naar bijna nul overslag (0,1%), wat bevestigt dat regularisatie noodzakelijk is.
  • LayerRoute-UniformInit: Keerde de skip differential om (planning sloeg meer over dan tool calls), wat het cruciale belang van gebiasde initialisatie bewijst.
  • LayerRoute-BCE: Getraind met expliciete labels in plaats van gezamenlijke optimalisatie, wat resulteerde in aanzienlijk lagere skip differentials (0,3%).

Gate Structuur: De geleerde gates convergeerden naar een stabiele twee-clusterstructuur: lagen 0–7 en 17–23 bleven betrouwbaar open, terwijl lagen 8–16 (de middelste cluster) betrouwbaar werden overgeslagen voor tool calls maar behouden bleven voor planning.

Betekenis en Claims

Het paper claimt dat het type agentische stap impliciet geleerd kan worden vanuit de taalmodellering-objective zonder dat daarvoor expliciete taak-type labels of een aparte classificatiekop nodig zijn.

• Efficiëntie: Het bereikt een reductie van 15,2% in FLOPs voor tool calls terwijl de kwaliteit behouden blijft (en zelfs licht verbetert).
• Adaptiviteit: Het systeem onderscheidt succesvol tussen korte, gestructureerde sequenties en lange, intensief redeneer-sequenties, en past rekenkracht alleen toe waar nodig.
• Parameter-efficiëntie: Het gehele adaptatieproces vereist minder dan 0,22% van de backbone-parameters en kan binnen minuten worden voltooid.

Erkende Beperkingen:
De auteurs merken op dat de middelste laag-cluster (8–16) een monolithisch skip-blok vormt zonder interne differentiatie, wat de potentiële FLOPs-besparing beperkt. Daarnaast blijft de skip-rate voor planning-stappen laag (2,34%), wat minimale voordelen biedt voor planning-zware workloads. De evaluatie is momenteel beperkt tot de 0,5B modelgrootte.