Pular uma camada ou repeti-la? Adaptação de profundidade em tempo de teste de LLMs pré-treinados

Skip a Layer or Loop it? Test-Time Depth Adaptation of Pretrained LLMs

July 10, 2025

Autores: Ziyue Li, Yang Li, Tianyi Zhou

cs.AI

Resumo

Uma rede neural pré-treinada pode adaptar sua arquitetura para diferentes entradas sem qualquer ajuste fino? Precisamos de todas as camadas para tarefas simples, e elas são adequadas para tarefas desafiadoras? Descobrimos que as camadas de um modelo de linguagem grande (LLM) pré-treinado podem ser manipuladas como módulos separados para construir um modelo melhor e até mais raso, personalizado para cada amostra de teste. Em particular, cada camada do modelo pré-treinado pode ser ignorada/podada ou repetida várias vezes como redes neurais recorrentes (RNN), e empilhada com outras em ordens arbitrárias, resultando em uma cadeia de camadas (CoLa) por amostra. Este espaço composicional expande consideravelmente o escopo dos trabalhos existentes sobre módulos pré-treinados em loop/recorrentes, poda de camadas ou redes de saída antecipada. Desenvolvemos um protocolo de Monte Carlo Tree Search (MCTS) para explorar e identificar a CoLa ótima para cada amostra de benchmarks de raciocínio matemático e de senso comum. Comparado a um modelo estático de profundidade fixa, a CoLa permite caminhos curtos (pensamento rápido), recorrência da(s) mesma(s) camada(s) (pensamento lento) e a combinação de ambos, oferecendo arquiteturas mais flexíveis e dinâmicas para diferentes entradas. Realizamos uma análise extensa da CoLa otimizada por MCTS, que leva a duas descobertas principais: (1) Para >75% das amostras com previsões corretas pelo LLM original, podemos encontrar CoLa mais curtas, sugerindo um grande espaço para melhorar a eficiência de inferência; (2) Para >60% das amostras com previsões originalmente incorretas, podemos identificar CoLa que alcançam previsões corretas, sugerindo um grande espaço de melhoria de desempenho. Nossos resultados destacam as deficiências de usar uma arquitetura fixa de LLMs pré-treinados para inferência em diferentes amostras e abrem caminho para desbloquear o poder de generalização da adaptação de profundidade no momento do teste.

English

Can a pretrained neural network adapt its architecture to different inputs without any finetuning? Do we need all layers for simple tasks, and are they adequate for challenging tasks? We found that the layers of a pretrained large language model (LLM) can be manipulated as separate modules to build a better and even shallower model customized for each test sample. In particular, each layer from the pretrained model can be skipped/pruned or repeated multiple times as recurrent neural networks (RNN), and stacked with others in arbitrary orders, yielding a chain-of-layers (CoLa) per sample. This compositional space greatly expands the scope of existing works on looped/recurrent pretrained modules, layer pruning, or early-exit networks. We develop a Monte Carlo Tree Search (MCTS) protocol to explore and identify the optimal CoLa for each sample from math and commonsense reasoning benchmarks. Compared to a static model of a fixed depth, CoLa allows shortcut paths (fast thinking), recurrence of the same layer(s) (slow thinking), and combining both, offering more flexible, dynamic architectures for different inputs. We conduct an extensive analysis of the MCTS-optimized CoLa, which leads to two key findings: (1) For >75% of samples with correct predictions by the original LLM, we can find shorter CoLa, suggesting a large space for improving inference efficiency; (2) For >60% of samples with originally incorrect predictions, we can identify CoLa achieving correct predictions, suggesting a large space of performance enhancement. Our results highlight the shortcomings of using a fixed architecture of pre-trained LLMs for inference on different samples and pave the way to unlock the generalization power of test-time depth adaptation.