Een Laag Overslaan of Herhalen? Testtijd Diepte-aanpassing van Vooraf Getrainde LLM's

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

July 10, 2025

Auteurs: Ziyue Li, Yang Li, Tianyi Zhou

cs.AI

Samenvatting

Kan een vooraf getraind neuraal netwerk zijn architectuur aanpassen aan verschillende invoerwaarden zonder enige finetuning? Hebben we alle lagen nodig voor eenvoudige taken, en zijn ze voldoende voor uitdagende taken? We ontdekten dat de lagen van een vooraf getraind groot taalmodel (LLM) als afzonderlijke modules kunnen worden gemanipuleerd om een beter en zelfs ondieper model te bouwen dat is aangepast aan elk testsample. In het bijzonder kan elke laag van het vooraf getrainde model worden overgeslagen/gesnoeid of meerdere keren worden herhaald als recurrent neural networks (RNN), en in willekeurige volgorden worden gestapeld met andere lagen, wat resulteert in een keten-van-lagen (CoLa) per sample. Deze compositorische ruimte breidt de reikwijdte van bestaande werken aan geloopte/recurrente vooraf getrainde modules, laagsnoei of early-exit netwerken aanzienlijk uit. We ontwikkelen een Monte Carlo Tree Search (MCTS)-protocol om de optimale CoLa voor elk sample te verkennen en te identificeren op basis van benchmarks voor wiskundig en gezond verstand redeneren. Vergeleken met een statisch model van een vaste diepte, biedt CoLa snelkoppelingen (snel denken), herhaling van dezelfde laag(en) (langzaam denken), en een combinatie van beide, wat meer flexibele, dynamische architecturen biedt voor verschillende invoerwaarden. We voeren een uitgebreide analyse uit van de MCTS-geoptimaliseerde CoLa, wat leidt tot twee belangrijke bevindingen: (1) Voor >75% van de samples met correcte voorspellingen door het originele LLM, kunnen we een kortere CoLa vinden, wat wijst op een grote ruimte voor het verbeteren van de inferentie-efficiëntie; (2) Voor >60% van de samples met oorspronkelijk incorrecte voorspellingen, kunnen we CoLa identificeren die correcte voorspellingen bereiken, wat wijst op een grote ruimte voor prestatieverbetering. Onze resultaten benadrukken de tekortkomingen van het gebruik van een vaste architectuur van vooraf getrainde LLM's voor inferentie op verschillende samples en banen de weg om de generalisatiekracht van diepte-aanpassing tijdens testtijd te ontsluiten.

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.