Skip a Layer or Loop it? Test-Time Depth Adaptation of Pretrained LLMs
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.
68709572c8391850d60978b7
Ziyue Li
68709572c8391850d60978b8
Yang Li
68709572c8391850d60978b9
Tianyi Zhou
claimed_verified
2025-07-11T07:59:58.984Z
647f5af5b0e96764589f3b2a
https://cdn-avatars.huggingface.co/v1/production/uploads/noauth/VJ4cDyjp5M3V5WmI5gPIU.jpeg
Tianyi Zhou
user
zhoutianyi
Skip a Layer or Loop it? Test-Time Depth Adaptation of Pretrained LLMs
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.
A method using chain-of-layers (CoLa) and Monte Carlo Tree Search (MCTS) optimizes the architecture of a pretrained large language model for individual samples, improving inference efficiency and performance.
pretrained large language model chain-of-layers CoLa Monte Carlo Tree Search MCTS looped/recurrent pretrained modules layer pruning early-exit networks shortcut paths recurrence test-time depth adaptation