¿Saltar una capa o repetirla? Adaptación de profundidad en tiempo de prueba de LLMs preentrenados

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

Resumen

¿Puede una red neuronal preentrenada adaptar su arquitectura a diferentes entradas sin necesidad de ajuste fino? ¿Necesitamos todas las capas para tareas simples, y son adecuadas para tareas desafiantes? Descubrimos que las capas de un modelo de lenguaje grande (LLM) preentrenado pueden manipularse como módulos separados para construir un modelo mejor e incluso más superficial personalizado para cada muestra de prueba. En particular, cada capa del modelo preentrenado puede omitirse/podarse o repetirse múltiples veces como redes neuronales recurrentes (RNN), y apilarse con otras en órdenes arbitrarios, generando una cadena-de-capas (CoLa) por muestra. Este espacio composicional amplía enormemente el alcance de trabajos existentes sobre módulos preentrenados en bucle/recurrentes, poda de capas o redes de salida temprana. Desarrollamos un protocolo de Búsqueda de Árbol de Monte Carlo (MCTS) para explorar e identificar la CoLa óptima para cada muestra en benchmarks de razonamiento matemático y de sentido común. En comparación con un modelo estático de profundidad fija, CoLa permite rutas cortas (pensamiento rápido), recurrencia de la(s) misma(s) capa(s) (pensamiento lento) y la combinación de ambos, ofreciendo arquitecturas más flexibles y dinámicas para diferentes entradas. Realizamos un análisis extenso de la CoLa optimizada por MCTS, lo que lleva a dos hallazgos clave: (1) Para >75% de las muestras con predicciones correctas por el LLM original, podemos encontrar CoLa más cortas, sugiriendo un gran espacio para mejorar la eficiencia de inferencia; (2) Para >60% de las muestras con predicciones originalmente incorrectas, podemos identificar CoLa que logran predicciones correctas, sugiriendo un gran espacio de mejora de rendimiento. Nuestros resultados destacan las limitaciones de usar una arquitectura fija de LLMs preentrenados para inferencia en diferentes muestras y allanan el camino para desbloquear el poder de generalización de la adaptación de profundidad en tiempo de prueba.

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.