LayerRoute : Saut de couches adaptatif conditionné par l'entrée via le fine-tuning LoRA pour les modèles de langage agentiques

Résumé

Les systèmes de modèles de langage agentiques alternent entre deux types d'étapes structurellement distincts : les appels d'outils structurés (courts, déterministes, à faible perplexité) et les étapes de planification/raisonnement ouvertes (longues, complexes, à forte perplexité). Malgré cette hétérogénéité, les systèmes d'inférence actuels appliquent un calcul identique à chaque étape. Nous introduisons LayerRoute, un adaptateur léger qui apprend à ignorer sélectivement des blocs de transformeur par entrée. LayerRoute enrichit chacun des 24 blocs de transformeur de Qwen2.5-0.5B-Instruct avec : (1) un routeur par couche (~897 paramètres, Linear(896,1)) qui produit une porte binaire dure via l'estimateur à passage direct, et (2) des adaptateurs LoRA (rang 8, ~1,08M paramètres) sur les projections d'attention Q/K/V/O. Les poids du backbone restent gelés. Un unique passage d'entraînement de bout en bout sur des données agentiques (Hermes, Glaive, GSM8K, Turing) avec un terme de régularisation de porte force le système à découvrir quels blocs peuvent être ignorés par type d'entrée. Après 3 000 pas (6,4 minutes sur un A100 40 Go), LayerRoute atteint un différentiel de saut de 12,91 % : les appels d'outils sautent 15,25 % des FLOPs tandis que les étapes de planification n'en sautent que 2,34 %, en utilisant seulement 1,10 M de paramètres entraînables (0,22 % des 494 M du backbone). La qualité s'améliore par rapport au modèle de base grâce à l'adaptation LoRA, avec un delta de perplexité de -1,29 pour les appels d'outils et de -1,30 pour la planification.

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.