LayerRoute: Salto adaptativo de capas condicionado por la entrada mediante ajuste fino LoRA para modelos de lenguaje agentivos

Resumen

Los sistemas de modelos de lenguaje agentivos alternan entre dos tipos de pasos estructuralmente distintos: llamadas estructuradas a herramientas (cortas, deterministas, baja perplejidad) y pasos de planificación/razonamiento abiertos (largos, complejos, alta perplejidad). A pesar de esta heterogeneidad, los sistemas de inferencia actuales aplican un cómputo idéntico a cada paso. Presentamos LayerRoute, un adaptador ligero que aprende a omitir selectivamente bloques transformadores por cada entrada. LayerRoute aumenta cada uno de los 24 bloques transformadores en Qwen2.5-0.5B-Instruct con: (1) un enrutador por capa (~897 parámetros, Linear(896,1)) que genera una compuerta binaria rígida mediante el estimador de paso directo, y (2) adaptadores LoRA (rango 8, ~1,08M de parámetros) en las proyecciones de atención Q/K/V/O. Los pesos de la red principal permanecen congelados. Una única pasada de entrenamiento de extremo a extremo en datos agentivos (Hermes, Glaive, GSM8K, Turing) con un término de regularización de compuerta obliga al sistema a descubrir qué bloques son omitibles por tipo de entrada. Después de 3.000 pasos (6,4 minutos en una A100 de 40 GB), LayerRoute logra un diferencial de omisión del 12,91%: las llamadas a herramientas omiten el 15,25% de los FLOPs, mientras que los pasos de planificación omiten solo el 2,34%, utilizando únicamente 1,10M de parámetros entrenables (el 0,22% de la red principal de 494M). La calidad mejora con respecto al modelo base debido a la adaptación LoRA, con un delta de perplejidad de -1,29 en llamadas a herramientas y -1,30 en planificación.

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.