Пропустить слой или зациклить? Адаптация глубины предобученных больших языковых моделей во время тестирования

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

July 10, 2025

Авторы: Ziyue Li, Yang Li, Tianyi Zhou

cs.AI

Аннотация

Может ли предобученная нейронная сеть адаптировать свою архитектуру к различным входным данным без дополнительного дообучения? Нужны ли все слои для простых задач, и достаточно ли они эффективны для сложных? Мы обнаружили, что слои предобученной крупной языковой модели (LLM) можно манипулировать как отдельные модули для создания более эффективной и даже более компактной модели, адаптированной для каждого тестового образца. В частности, каждый слой предобученной модели можно пропустить/удалить или повторить несколько раз, как в рекуррентных нейронных сетях (RNN), и комбинировать с другими в произвольном порядке, формируя цепочку слоев (CoLa) для каждого образца. Это композиционное пространство значительно расширяет возможности существующих подходов, таких как зацикленные/рекуррентные предобученные модули, удаление слоев или сети с ранним выходом. Мы разработали протокол поиска по дереву Монте-Карло (MCTS) для исследования и определения оптимальной CoLa для каждого образца из тестов на математические и логические рассуждения. По сравнению со статической моделью фиксированной глубины, CoLa позволяет использовать короткие пути (быстрое мышление), повторение одних и тех же слоев (медленное мышление) или их комбинацию, предлагая более гибкие и динамичные архитектуры для различных входных данных. Мы провели детальный анализ оптимизированных с помощью MCTS CoLa, что привело к двум ключевым выводам: (1) Для более чем 75% образцов, правильно классифицированных исходной LLM, можно найти более короткие CoLa, что указывает на значительный потенциал повышения эффективности вывода; (2) Для более чем 60% образцов, изначально классифицированных неправильно, можно найти CoLa, обеспечивающие правильные предсказания, что указывает на значительный потенциал повышения производительности. Наши результаты подчеркивают недостатки использования фиксированной архитектуры предобученных LLM для вывода на различных образцах и прокладывают путь к раскрытию обобщающей способности адаптации глубины на этапе тестирования.

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.