LayerRoute: Адаптивный пропуск слоев, обусловленный входными данными, посредством тонкой настройки LoRA для агентных языковых моделей

Аннотация

Агентные системы языковых моделей чередуют два структурно различных типа шагов: структурированные вызовы инструментов (короткие, детерминированные, с низкой перплексией) и открытые шаги планирования/рассуждения (длинные, сложные, с высокой перплексией). Несмотря на такую гетерогенность, современные системы инференса применяют одинаковые вычислительные затраты к каждому шагу. Мы представляем LayerRoute — легковесный адаптер, который обучается выборочно пропускать блоки трансформера для каждого входного образца. LayerRoute дополняет каждый из 24 блоков трансформера в Qwen2.5-0.5B-Instruct: (1) маршрутизатором на уровне слоя (~897 параметров, Linear(896,1)), который выдает жесткий бинарный шлюз через прямой оценщик, и (2) LoRA-адаптерами (ранг 8, ~1,08M параметров) на проекциях внимания Q/K/V/O. Веса основной модели остаются замороженными. Один сквозной проход обучения на агентных данных (Hermes, Glaive, GSM8K, Turing) с членом регуляризации шлюзов заставляет систему обнаруживать, какие блоки являются пропускаемыми для каждого типа входных данных. После 3000 шагов (6,4 минуты на A100 40GB) LayerRoute достигает дифференциала пропусков в 12,91%: вызовы инструментов пропускают 15,25% FLOPs, тогда как шаги планирования — лишь 2,34%, используя всего 1,10M обучаемых параметров (0,22% от 494M параметров основной модели). Качество улучшается по сравнению с базовой моделью за счет LoRA-адаптации, с дельтой перплексии -1,29 для вызовов инструментов и -1,30 для планирования.

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.