LayerRoute: Eingabebedingtes adaptives Layer-Überspringen durch LoRA-Feintuning für agentische Sprachmodelle

Zusammenfassung

Agentische Sprachmodellsysteme wechseln zwischen zwei strukturell unterschiedlichen Schritttypen: strukturierte Werkzeugaufrufe (kurz, deterministisch, niedrige Perplexität) und offene Planungs-/Denkschritte (lang, komplex, hohe Perplexität). Trotz dieser Heterogenität wenden aktuelle Inferenzsysteme auf jeden Schritt identische Rechenleistung an. Wir stellen LayerRoute vor, einen leichten Adapter, der lernt, Transformer-Blöcke selektiv pro Eingabe zu überspringen. LayerRoute erweitert jeden der 24 Transformer-Blöcke in Qwen2.5-0.5B-Instruct um: (1) einen pro Schicht arbeitenden Router (~897 Parameter, Linear(896,1)), der über den Straight-Through-Schätzer ein hartes binäres Gatter ausgibt, und (2) LoRA-Adapter (Rang 8, ~1,08 Mio. Parameter) auf den Q/K/V/O-Aufmerksamkeitsprojektionen. Die Grundgerüstgewichte bleiben eingefroren. Ein einziger durchgängiger Trainingsdurchlauf auf agentischen Daten (Hermes, Glaive, GSM8K, Turing) mit einem Regularisierungsterm für das Gatter zwingt das System zu entdecken, welche Blöcke pro Eingabetyp überspringbar sind. Nach 3.000 Schritten (6,4 Minuten auf einer A100 40GB) erreicht LayerRoute eine Sprungdifferenz von 12,91 %: Werkzeugaufrufe überspringen 15,25 % der FLOPs, während Planungsschritte nur 2,34 % überspringen, wobei nur 1,10 Mio. trainierbare Parameter (0,22 % des 494 Mio. Grundgerüsts) verwendet werden. Die Qualität verbessert sich gegenüber dem Basismodell aufgrund der LoRA-Adaption, mit einem Perplexitätsdelta von -1,29 bei Werkzeugaufrufen und -1,30 bei Planung.

English

Agentic language model systems alternate between two structurally distinct step types: structured tool calls (short, deterministic, low perplexity) and open-ended planning/reasoning steps (long, complex, high perplexity). Despite this heterogeneity, current inference systems apply identical compute to every step. We introduce LayerRoute, a lightweight adapter that learns to selectively skip transformer blocks on a per-input basis. LayerRoute augments each of the 24 transformer blocks in Qwen2.5-0.5B-Instruct with: (1) a per-layer router (~897 parameters, Linear(896,1)) that outputs a hard binary gate via the straight-through estimator, and (2) LoRA adapters (rank 8, ~1.08M parameters) on the Q/K/V/O attention projections. The backbone weights remain frozen. A single end-to-end training pass on agentic data (Hermes, Glaive, GSM8K, Turing) with a gate regularisation term forces the system to discover which blocks are skippable per input type. After 3,000 steps (6.4 minutes on an A100 40GB), LayerRoute achieves a 12.91% skip differential: tool calls skip 15.25% of FLOPs while planning steps skip only 2.34%, using only 1.10M trainable parameters (0.22% of the 494M backbone). Quality improves over the base model due to LoRA adaptation, with perplexity delta of -1.29 on tool calls and -1.30 on planning.