PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Prateek Kumar Sikdar

arXiv:2606.01838v1cs.CLcs.AIcs.LG

TL;DR

Agentic systems apply the same full transformer computation to predictable tool calls and complex planning steps, motivating adaptive inference. LayerRoute combines input-conditioned hard-gated skipping with joint LoRA fine-tuning and gate regularisation. It achieves a 12.91% skip differential while improving quality over the base model and using 1.10M trainable parameters.

  • Problem

    Current inference systems route both predictable tool calls and complex planning steps through all transformer layers at full compute cost.

  • Method

    LayerRoute combines input-conditioned hard-gated layer skipping with joint LoRA fine-tuning of attention projections and gate regularisation.

  • Results

    12.91% skip differential separates tool-call and planning compute, while LayerRoute achieves lower perplexity than the full-layer baseline and improves over the base model.

  • Takeaways & Limitations

    LayerRoute demonstrates that agentic step type can be learned implicitly from the language-modelling objective without explicit task-type labels or a separate classification head.

  • Takeaways & Limitations

    Planning steps skip only 2.34% of computation, so LayerRoute provides little benefit for planning-heavy workloads.

Abstract

from arXiv · show

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.

1 Introduction

LayerRoute targets the compute mismatch between predictable tool calls and complex planning steps by learning input-conditioned layer skipping. It combines hard-gated routing, LoRA adaptation, regularisation, and biased initialisation, achieving a 12.91% skip differential.

  • Motivation: Agentic workflows combine short, predictable tool calls with complex, open-ended planning steps.Tool calls are structured sequences, whereas planning requires multi-step reasoning over a broad output space.
  • Motivation: Current inference routes both step types through every transformer layer, despite deep layers contributing marginally for tool calls.The motivation is to avoid identical full-compute treatment when early residual-stream confidence is already high.
  • Approach: LayerRoute learns input-conditioned layer skipping with a parameter-efficient adapter.The backbone remains the target of adaptive routing rather than receiving unrestricted full-model updates.
  • Approach: Hard-gated skip connections use straight-through estimation, keeping train-time and inference-time routing decisions aligned.The gates are binary during the forward pass while gradients are estimated through the continuous router output.
  • Approach: Joint LoRA fine-tuning of attention projections lets routing and adapted weights co-evolve under the language-modelling objective.Gate regularisation prevents all gates from remaining open, while biased middle-layer initialisation enables differentiation from the first step.
  • Results: 12.91% skip differential separates tool-call and planning compute, and the system was trained in under 7 minutes.This result is reported on Qwen2.5-0.5B-Instruct as empirical validation of the approach.

2 Related Work

LayerRoute builds on input-dependent layer routing and parameter-efficient adaptation while targeting agentic step types rather than generic language processing. It reduces active-layer computation as a complement to kernel and memory optimisations.

  • Early exit and layer skipping: Prior work includes early exit and input-dependent layer skipping based on intermediate confidence or routing.DeeBERT and PABEE enable early exits, while SkipBERT skips entire layers and Mixture of Depths routes per token.
  • Early exit and layer skipping: LayerRoute differs from Mixture of Depths by using joint LoRA fine-tuning rather than training from scratch and targeting agentic step types.Its routing mechanism is therefore paired with a parameter-efficient adaptation strategy and an agentic setting.
  • Parameter-efficient fine-tuning: LoRA decomposes weight updates into low-rank matrices, allowing attention projections and routers to adapt together with few trainable parameters.The routing policy co-evolves with the adapted weights during fine-tuning.
  • Efficient inference for LLMs: LayerRoute is complementary to speculative decoding, PagedAttention, and FlashAttention because it reduces active layer count rather than optimising kernels or memory.The methods address different inference-efficiency mechanisms.
  • Agentic LLM systems: LayerRoute targets inference efficiency, whereas AgenticQwen targets agentic task capability through agentic fine-tuning.The two approaches are presented as complementary in objective.

3 Method

LayerRoute freezes Qwen2.5-0.5B-Instruct and adds per-layer routers with LoRA adapters, using hard input-conditioned gates to skip blocks. Joint language-model and gate-regularised training encourages genuinely skippable layers while preserving hard decisions during inference.

  • Architecture: LayerRoute augments a frozen 24-layer Qwen2.5-0.5B-Instruct backbone with per-layer routers and LoRA adapters.The routers operate at each transformer block, while LoRA is added to attention projections.
  • Per-layer routing: Each router uses the mean-pooled hidden state entering its block to produce a sigmoid gate score.Each router has 897 parameters, and the 24-router collection totals 21,528 parameters.
  • Gated skip connections: A hard gate either executes the block with LoRA active or passes the hidden state through unchanged with zero block compute.The forward transformation is therefore input-conditioned at the block level.
  • Straight-through estimation: The straight-through estimator uses hard binary decisions in the forward pass while allowing gradients to flow through the continuous sigmoid values.This keeps training and inference decisions aligned.
  • Training objective: Joint optimisation combines autoregressive language-model loss with gate regularisation weighted by λ = 1.0.The regulariser penalises uniformly high soft gate values, pressuring routers to identify skippable blocks.
  • Initialisation: Layer-dependent bias initialisation starts middle layers near skipping and early or late layers near execution to break gate symmetry.Layers 8–16 begin near σ(b_i) ≈ 0.27, while layers 0–7 and 17–23 begin near σ(b_i) ≈ 0.73.

4 Experiments

Experiments train LayerRoute on mixed tool-call and planning/reasoning data and evaluate held-out examples with skip- and compute-oriented metrics. The setup also compares the full model with two LayerRoute ablations and measures perplexity and LoRA impact.

  • Training data: The mixed agentic training set combines tool-call data from Hermes and Glaive with planning/reasoning data from GSM8K and Turing.Tool-call sources contribute 1,893 and 5,000 samples, while planning sources contribute 5,000 and 50 samples.
  • Training data: 10,749 training and 1,194 validation samples are split 90/10 and padded or truncated to 512 tokens.Sequences use Qwen’s chat template.
  • Training procedure: 3,000-step AdamW training takes 381.9 seconds, approximately 6.4 minutes, on an A100 40GB.The effective batch size is 16 after gradient accumulation.
  • Evaluation: Evaluation uses 100 held-out samples: 50 Hermes tool calls and 50 GSM8K planning examples.The evaluation covers both agentic step types.
  • Metrics: Skip differential is defined as mean tool-call skip percentage minus mean planning skip percentage, while FLOPs reduction uses average active layers.Perplexity compares the gated model with a full-layer baseline, and LoRA impact compares base Qwen with the LoRA-adapted model.
  • Baselines and ablations: The comparisons include a no-skipping full model, a provenance-label BCE ablation, and a uniformly initialised joint-training ablation.LayerRoute-BCE replaces joint language-model plus gate-regularised training with binary cross-entropy labels.

5 Results

LayerRoute learns input-conditioned skipping that differentiates tool-call and planning computation, while preserving or improving perplexity through LoRA adaptation. Its routing forms stable layer clusters and yields a measurable skip differential on held-out agentic samples.

  • Main results: 12.91% skip differential separates tool-call and planning steps, with tool calls skipping more blocks than planning inputs.Figure 1 shows input-dependent gate patterns, and Table 2 defines skip differential as tool skip% − planning skip%.
  • Gate structure: Layers 0–7 and 17–23 stabilise near 0.73, while layers 8–16 stabilise near 0.27 after 3,000 steps.The resulting hard gates reliably open the outer layers and close the middle cluster.
  • Gate structure: 0.050 gate variance confirms structural differentiation rather than a flat routing policy.The reported variance distinguishes the learned structure from random or uniformly converged gates.
  • Quality analysis: PPL(gated) < PPL(full) for both step types under the shared LoRA-adapted weights.The full baseline uses all layers with the same LoRA weights trained under the gated objective.
  • Quality analysis: LoRA improves loss by −1.31 on tool calls and −0.88 on planning relative to the base Qwen model without adapters.These improvements accompany lower perplexity for the gated model than the full-layer baseline.
  • Training dynamics: The skip rate jumps to ≈38% during steps 0–50, then stabilises at 10–18% from steps 500–3000 as loss converges to ≈1.5.The intermediate phase combines rapid LM-loss reduction with declining skip rate as LoRA and routers adapt.

6 Analysis

Analysis attributes LayerRoute’s routing behaviour to biased initialization and joint language-model training, while identifying limits in within-cluster differentiation and contrasting sequence-level routing with token-level alternatives.

  • Initialization and routing: Without biased initialization, every block initially runs and gate regularisation pushes gates down uniformly without producing layer-specific discrimination.All routers start at σ(+1.0) = 0.73 in the UniformInit variant.
  • Initialization and routing: Biased initialization starts layers 8–16 below threshold, exposing their contribution to LM loss and enabling routers to differentiate tool calls from planning inputs.Middle layers start at σ(−1.0) = 0.27, whereas uniformly initialized routers provide no per-layer discrimination signal.
  • Uniformity within the middle cluster: Layers 8–16 converge to indistinguishable values near 0.274, indicating binary middle-versus-rest routing rather than a continuous importance gradient.The paper identifies finer-grained differentiation within this cluster as future work.
  • Relationship to Mixture of Depths: LayerRoute routes entire sequences with one gate per block, unlike Mixture of Depths, which routes individual tokens and requires training from scratch.Sequence-level routing matches the paper’s framing of agentic step type as a sequence-level property.
  • Adaptation strategy: Joint LoRA and router training lets the routing policy co-evolve with adapted attention projections while the backbone remains suitable for post-hoc adaptation.The cited analysis contrasts LayerRoute’s training strategy with token-level routing approaches.

7 Limitations

LayerRoute’s main limitations are modest and uneven compute savings, coarse routing within the skipped middle-layer block, limited model-size validation, and imbalanced planning data.

  • Layers 8–16 form a monolithic skip block without internal differentiation, limiting potential FLOPs savings.Finer routing within this cluster is identified as a possible way to increase savings.
  • 15.2% FLOPs reduction on tool calls is modest and larger savings may risk quality degradation or require per-token routing.The stated alternatives are a higher gate regularisation weight or per-token routing, as in Mixture of Depths.
  • 2.34% skip on planning yields minimal compute savings, so planning-heavy workloads receive little benefit.
  • Evaluation covers only the 0.5B model, and scaling to 3B or 7B may change which layers are skippable.
  • Turing Open Reasoning contributes only 50 planning samples because of dataset access limitations, while a more balanced split could improve routing discrimination.

8 Conclusion

LayerRoute learns input-conditioned layer skipping for agentic language-model inference by combining hard-gated routing, LoRA fine-tuning, regularisation, and biased initialisation. It produces a stable distinction between tool-call and planning computation, with future work targeting finer routing, larger models, and complementary decoding methods.

  • LayerRoute combines hard-gated skip connections, LoRA fine-tuning, gate regularisation, and biased initialisation to learn input-conditioned layer skipping.
  • 12.91% skip differential separates tool calls, which use 15.2% fewer FLOPs, from planning, which uses 2.3% fewer while preserving quality.
  • Future work includes finer within-cluster routing, extension to larger model sizes, and integration with speculative decoding.

C Training Curves

The architecture and training figures show hard per-layer routing, a stable two-cluster gate pattern, and convergence toward lower loss with a stabilized skip rate.

  • LayerRoute architecture: Each transformer block uses a hard STE gate to either run with LoRA-adapted weights or bypass the block through a skip connection.
  • Learned gate values: Layers 0–7 and 17–23 stabilize near 0.73 and run, whereas layers 8–16 stabilize near 0.27 and skip.The reported gate variance is 0.050.
  • Training curves: LM loss converges from 3.96 to 1.53 while skip percentage declines from 37.9% and stabilizes at 10–18%.The skip percentage begins from a biased initialization.
Loading 2606.01838v1…