Sauter une couche ou la boucler ? Adaptation en profondeur à l'exécution des LLM préentraînés

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

July 10, 2025

papers.authors: Ziyue Li, Yang Li, Tianyi Zhou

cs.AI

papers.abstract

Un réseau de neurones pré-entraîné peut-il adapter son architecture à différents entrées sans aucun ajustement fin ? Avons-nous besoin de toutes les couches pour des tâches simples, et sont-elles adéquates pour des tâches plus complexes ? Nous avons découvert que les couches d'un grand modèle de langage (LLM) pré-entraîné peuvent être manipulées comme des modules distincts pour construire un modèle meilleur et même plus superficiel, personnalisé pour chaque échantillon de test. En particulier, chaque couche du modèle pré-entraîné peut être ignorée/élaguée ou répétée plusieurs fois comme dans les réseaux de neurones récurrents (RNN), et empilée avec d'autres dans des ordres arbitraires, produisant une chaîne de couches (CoLa) par échantillon. Cet espace compositionnel élargit considérablement la portée des travaux existants sur les modules pré-entraînés en boucle/récurrents, l'élagage de couches ou les réseaux à sortie précoce. Nous développons un protocole de recherche arborescente Monte Carlo (MCTS) pour explorer et identifier la CoLa optimale pour chaque échantillon à partir de benchmarks de raisonnement mathématique et de bon sens. Par rapport à un modèle statique de profondeur fixe, CoLa permet des chemins raccourcis (pensée rapide), la récurrence de la même couche ou des mêmes couches (pensée lente), et la combinaison des deux, offrant des architectures plus flexibles et dynamiques pour différentes entrées. Nous menons une analyse approfondie de la CoLa optimisée par MCTS, qui conduit à deux conclusions clés : (1) Pour >75 % des échantillons avec des prédictions correctes par le LLM original, nous pouvons trouver des CoLa plus courtes, suggérant un large espace pour améliorer l'efficacité de l'inférence ; (2) Pour >60 % des échantillons avec des prédictions initialement incorrectes, nous pouvons identifier des CoLa atteignant des prédictions correctes, suggérant un large espace d'amélioration des performances. Nos résultats mettent en évidence les limites de l'utilisation d'une architecture fixe de LLM pré-entraînés pour l'inférence sur différents échantillons et ouvrent la voie à l'exploitation du pouvoir de généralisation de l'adaptation de la profondeur au moment du test.

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.