Eine Schicht überspringen oder sie wiederholen? Testzeitliche Tiefenanpassung vortrainierter großer Sprachmodelle

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

Kann ein vortrainiertes neuronales Netzwerk seine Architektur an verschiedene Eingaben anpassen, ohne Feinabstimmung? Benötigen wir alle Schichten für einfache Aufgaben, und sind sie für anspruchsvolle Aufgaben ausreichend? Wir fanden heraus, dass die Schichten eines vortrainierten großen Sprachmodells (LLM) als separate Module manipuliert werden können, um ein besseres und sogar flacheres Modell zu erstellen, das für jede Testprobe maßgeschneidert ist. Insbesondere kann jede Schicht des vortrainierten Modells übersprungen/beschnitten oder mehrfach als rekurrentes neuronales Netzwerk (RNN) wiederholt und in beliebiger Reihenfolge mit anderen gestapelt werden, wodurch eine Kette von Schichten (CoLa) pro Probe entsteht. Dieser kompositionelle Raum erweitert den Umfang bestehender Arbeiten zu geloopten/rekurrenten vortrainierten Modulen, Schichtbeschnitt oder Early-Exit-Netzwerken erheblich. Wir entwickeln ein Monte-Carlo-Baumsuchverfahren (MCTS), um den optimalen CoLa für jede Probe aus mathematischen und Common-Sense-Reasoning-Benchmarks zu erkunden und zu identifizieren. Im Vergleich zu einem statischen Modell mit fester Tiefe ermöglicht CoLa Abkürzungspfade (schnelles Denken), Wiederholung derselben Schicht(en) (langsames Denken) und die Kombination beider, wodurch flexiblere, dynamischere Architekturen für verschiedene Eingaben angeboten werden. Wir führen eine umfangreiche Analyse des MCTS-optimierten CoLa durch, die zu zwei zentralen Erkenntnissen führt: (1) Für >75 % der Proben mit korrekten Vorhersagen durch das ursprüngliche LLM können wir kürzere CoLa finden, was auf einen großen Raum zur Verbesserung der Inferenzeffizienz hinweist; (2) Für >60 % der Proben mit ursprünglich falschen Vorhersagen können wir CoLa identifizieren, die korrekte Vorhersagen erzielen, was auf einen großen Raum zur Leistungssteigerung hindeutet. Unsere Ergebnisse verdeutlichen die Nachteile der Verwendung einer festen Architektur vortrainierter LLMs für die Inferenz bei verschiedenen Proben und ebnen den Weg, um das Generalisierungspotenzial der Tiefenanpassung zur Testzeit freizusetzen.

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.