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

Resumo Técnico: LayerRoute – Pulo de Camada Adaptativo Condicionado ao Input via Ajuste Fino de LoRA

Problema

Sistemas de IA agêntica modernos orquestram Grandes Modelos de Linguagem (LLMs) através de fluxos de trabalho heterogêneos que compreendem dois tipos distintos de etapas:

1. Chamadas de ferramentas estruturadas: Sequências curtas e determinísticas com baixa perplexidade (ex: consultas a bancos de dados).
2. Planejamento/Raciocínio aberto: Sequências longas e complexas com alta perplexidade que exigem raciocínio de múltiplos passos.

Sistemas de inferência atuais aplicam recursos computacionais idênticos a ambos os tipos de etapa, roteando cada entrada por todas as camadas do transformer. Isso é ineficiente porque a alta confiança nos estados iniciais do fluxo residual para chamadas de ferramentas sugere que camadas profundas contribuem marginalmente para a predição final. O artigo aborda a necessidade de um sistema de inferência que possa alocar computação dinamicamente com base na complexidade da entrada sem sacrificar a qualidade do modelo.

Metodologia

Os autores propõem o LayerRoute, um adaptador leve projetado para o modelo Qwen2.5-0.5B-Instruct (24 camadas, tamanho oculto 896). A arquitetura aumenta o backbone congelado com dois componentes primários por bloco de transformer:

1. Roteadores por Camada: Cada bloco $i$ contém uma camada linear leve ($w_i^T \bar{h}_i + b_i$) que produz um valor sigmoide suave $\sigma_i$. Este é convertido em um portão binário rígido $g_i \in \{0, 1\}$ via um limiar de 0,5.

   • Passagem de Forward: Se $g_i=1$, o bloco é executado normalmente (com adaptadores LoRA ativos). Se $g_i=0$, o estado oculto passa inalterado, incorrendo em zero computação para aquele bloco.
   • Treinamento: O portão rígido não diferenciável é tratado via Estimador de Passagem Direta (Straight-Through Estimator - STE), permitindo que os gradientes fluam através da sigmoide durante a retropropagação enquanto mantém decisões rígidas durante a inferência.

2. Adaptadores LoRA: Adaptadores de baixo rank (Rank 8) são aplicados às projeções de Query, Key, Value e Output ($W_Q, W_K, W_V, W_O$) de todos os 24 blocos. Os pesos do backbone permanecem congelados.

Estratégia de Treinamento e Inovações:

• Otimização Conjunta: Os pesos do roteador e as matrizes LoRA são treinados conjuntamente usando uma função de perda combinada: $L = L_{LM} + \lambda \cdot \frac{1}{L}\sum \sigma(s_i)$.
• Regularização de Portão: Um termo de regularização ($\lambda=1.0$) penaliza valores altos de portão suave, evitando que o roteador colapse para uma solução trivial onde todos os portões permanecem abertos.
• Inicialização Enviesada: Para quebrar a simetria e evitar o problema do "ovo e da galinha", onde os portões não conseguem diferenciar até que pulem, os autores inicializam os roteadores das camadas médias (camadas 8–16) com um viés de $-1,0$ (começando fechados/pulados) e as camadas externas com $+1,0$ (começando abertos). Isso força a diferenciação imediata baseada na perda de modelagem de linguagem.

Principais Contribuições

O artigo descreve cinco contribuições específicas:

1. Conexões de salto com portão rígido por camada usando STE para eliminar a discrepância entre treino e inferência.
2. Ajuste fino de LoRA conjunto das projeções de atenção junto com o aprendizado do roteador, permitindo que a política de roteamento evolua em conjunto com os pesos adaptados.
3. Regularização de portão para forçar a descoberta de blocos passíveis de salto em vez de manter a computação total.
4. Inicialização enviesada dos roteadores de camadas médias para permitir a diferenciação baseada em gradiente desde o primeiro passo de treinamento.
5. Validação empírica no Qwen2.5-0.5B-Instruct demonstrando um diferencial de salto significativo entre chamadas de ferramentas e etapas de planejamento, treinado em menos de 7 minutos.

Resultados Experimentais

O modelo foi treinado em um dataset agêntico misto (Hermes, Glaive, GSM8K, Turing) por 3.000 passos (aprox. 6,4 minutos em uma A100 40GB) usando apenas 1,10M de parâmetros treináveis (0,22% do backbone de 494M).

• Diferencial de Salto: O LayerRoute alcançou um diferencial de salto de 12,91%. Chamadas de ferramentas pularam 15,25% dos FLOPs, enquanto etapas de planejamento pularam apenas 2,34%.
• Melhoria de Qualidade: Contrariando os típicos trade-offs de eficiência, o modelo com portão mostrou melhoria na perplexidade em relação ao baseline de camada total.
  • Chamadas de ferramentas: Delta de perplexidade de −1,29.
  • Planejamento: Delta de perplexidade de −1,30.
  • Essa melhoria é atribuída à adaptação LoRA compensando o custo de qualidade do salto.
• Estudos de Ablação:
  • LayerRoute-NoReg: Colapsou para um salto próximo de zero (0,1%), confirmando que a regularização é necessária.
  • LayerRoute-UniformInit: Inverteu o diferencial de salto (planejamento pulou mais do que chamadas de ferramentas), provando o papel crítico da inicialização enviesada.
  • LayerRoute-BCE: Treinado com rótulos explícitos em vez de otimização conjunta, resultando em diferenciais de salto significativamente menores (0,3%).

Estrutura do Portão: Os portões aprendidos convergiram para uma estrutura estável de dois clusters: as camadas 0–7 e 17–23 permaneceram confiavelmente abertas, enquanto as camadas 8–16 (o cluster central) foram confiavelmente puladas para chamadas de ferramentas, mas retidas para planejamento.

Significado e Alegações

O artigo afirma que o tipo de etapa agêntica pode ser aprendido implicitamente a partir do objetivo de modelagem de linguagem sem a necessidade de rótulos explícitos de tipo de tarefa ou um cabeçalho de classificação separado.

• Eficiência: Alcança uma redução de 15,2% nos FLOPs para chamadas de ferramentas enquanto preserva (e até melhora ligeiramente) a qualidade.
• Adaptabilidade: O sistema distingue com sucesso sequências curtas e estruturadas de sequências longas e densas em raciocínio, aplicando computação apenas onde necessário.
• Eficiência de Parâmetros: Todo o processo de adaptação requer menos de 0,22% dos parâmetros do backbone e pode ser concluído em minutos.

Limitações Reconhecidas:
Os autores observam que o cluster de camada média (8–16) forma um bloco de salto monolítico sem diferenciação interna, limitando o potencial de economia de FLOPs. Além disso, a taxa de salto para etapas de planejamento permanece baixa (2,34%), oferecendo benefício mínimo para cargas de trabalho intensivas em planejamento. A avaliação é atualmente limitada ao tamanho do modelo 0.5B.