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

• URL: http://arxiv.org/abs/2507.07996v1
• Date: Thu, 10 Jul 2025 17:59:53 GMT
• Title: Skip a Layer or Loop it? Test-Time Depth Adaptation of Pretrained LLMs
• Authors: Ziyue Li, Yang Li, Tianyi Zhou,
• Abstract summary: We find that 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.<n>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.
• Score: 21.541258368039955
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Learning to Reason Across Parallel Samples for LLM Reasoning [45.60752271688715]
  Scaling test-time compute brings substantial performance gains for large language models.<n>We propose a new way to leverage such multiple sample set.<n>We train a compact LLM, that takes a sequence of multiple samples and output the final answer.
  arXiv  Detail & Related papers  (2025-06-10T17:42:35Z)
• C3PO: Critical-Layer, Core-Expert, Collaborative Pathway Optimization for Test-Time Expert Re-Mixing [21.119495676190127]
  Mixture-of-Experts (MoE) Large Language Models (LLMs) suffer from severely sub-optimal expert pathways.<n> naive expert selection learned from pretraining leaves a surprising 10-20% accuracy gap for improvement.<n>We develop a novel class of test-time optimization methods to re-weight or "re-mixing" the experts in different layers jointly for each test sample.
  arXiv  Detail & Related papers  (2025-04-10T17:59:56Z)
• LESA: Learnable LLM Layer Scaling-Up [57.0510934286449]
  Training Large Language Models (LLMs) from scratch requires immense computational resources, making it prohibitively expensive.<n>Model scaling-up offers a promising solution by leveraging the parameters of smaller models to create larger ones.<n>We propose textbfLESA, a novel learnable method for depth scaling-up.
  arXiv  Detail & Related papers  (2025-02-19T14:58:48Z)
• Scaling LLM Inference with Optimized Sample Compute Allocation [56.524278187351925]
  We propose OSCA, an algorithm to find an optimal mix of different inference configurations. Our experiments show that with our learned mixed allocation, we can achieve accuracy better than the best single configuration. OSCA is also shown to be effective in agentic beyond single-turn tasks, achieving a better accuracy on SWE-Bench with 3x less compute than the default configuration.
  arXiv  Detail & Related papers  (2024-10-29T19:17:55Z)
• Initialization Matters: On the Benign Overfitting of Two-Layer ReLU CNN with Fully Trainable Layers [20.25049261035324]
  We extend the analysis to two-layer ReLU convolutional neural networks (CNNs) with fully trainable layers. Our results show that the scaling of the output layer is crucial to the training dynamics. In both settings, we provide nearly matching upper and lower bounds on the test errors.
  arXiv  Detail & Related papers  (2024-10-24T20:15:45Z)
• Large Language Models aren't all that you need [0.0]
  This paper describes the architecture and systems built towards solving the SemEval 2023 Task 2: MultiCoNER II. We evaluate two approaches (a) a traditional Random Fields model and (b) a Large Language Model (LLM) fine-tuned with a customized head and compare the two approaches.
  arXiv  Detail & Related papers  (2024-01-01T08:32:50Z)
• RGM: A Robust Generalizable Matching Model [49.60975442871967]
  We propose a deep model for sparse and dense matching, termed RGM (Robust Generalist Matching) To narrow the gap between synthetic training samples and real-world scenarios, we build a new, large-scale dataset with sparse correspondence ground truth. We are able to mix up various dense and sparse matching datasets, significantly improving the training diversity.
  arXiv  Detail & Related papers  (2023-10-18T07:30:08Z)
• Test-Time Adaptation with CLIP Reward for Zero-Shot Generalization in Vision-Language Models [76.410400238974]
  We propose TTA with feedback to rectify the model output and prevent the model from becoming blindly confident. A CLIP model is adopted as the reward model during TTA and provides feedback for the VLM. The proposed textitreinforcement learning with CLIP feedback(RLCF) framework is highly flexible and universal.
  arXiv  Detail & Related papers  (2023-05-29T11:03:59Z)
• OFA$^2$: A Multi-Objective Perspective for the Once-for-All Neural Architecture Search [79.36688444492405]
  Once-for-All (OFA) is a Neural Architecture Search (NAS) framework designed to address the problem of searching efficient architectures for devices with different resources constraints. We aim to give one step further in the search for efficiency by explicitly conceiving the search stage as a multi-objective optimization problem.
  arXiv  Detail & Related papers  (2023-03-23T21:30:29Z)
• LV-BERT: Exploiting Layer Variety for BERT [85.27287501885807]
  We introduce convolution into the layer type set, which is experimentally found beneficial to pre-trained models. We then adopt an evolutionary algorithm guided by pre-training accuracy to find the optimal architecture. LV-BERT model obtained by our method outperforms BERT and its variants on various downstream tasks.
  arXiv  Detail & Related papers  (2021-06-22T13:20:14Z)
• MTL-NAS: Task-Agnostic Neural Architecture Search towards General-Purpose Multi-Task Learning [71.90902837008278]
  We propose to incorporate neural architecture search (NAS) into general-purpose multi-task learning (GP-MTL) In order to adapt to different task combinations, we disentangle the GP-MTL networks into single-task backbones. We also propose a novel single-shot gradient-based search algorithm that closes the performance gap between the searched architectures.
  arXiv  Detail & Related papers  (2020-03-31T09:49:14Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.