레이어를 건너뛸 것인가, 루프를 돌 것인가? 사전 학습된 대형 언어 모델의 테스트 시점 깊이 적응

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)의 각 레이어를 별도의 모듈로 조작하여 각 테스트 샘플에 맞춤화된 더 나은, 심지어 더 얕은 모델을 구축할 수 있음을 발견했다. 특히, 미리 학습된 모델의 각 레이어는 건너뛰거나(pruned) 반복 신경망(RNN)처럼 여러 번 반복될 수 있으며, 임의의 순서로 다른 레이어와 결합되어 샘플별로 체인 오브 레이어(CoLa)를 형성할 수 있다. 이 조합 공간은 기존의 루프/반복 미리 학습 모듈, 레이어 프루닝, 또는 조기 종료 네트워크 연구의 범위를 크게 확장한다. 우리는 수학 및 상식 추론 벤치마크에서 각 샘플에 대한 최적의 CoLa를 탐색하고 식별하기 위해 몬테카를로 트리 탐색(MCTS) 프로토콜을 개발했다. 고정 깊이의 정적 모델과 비교했을 때, CoLa는 단축 경로(빠른 사고), 동일한 레이어의 반복(느린 사고), 그리고 둘을 결합하는 것을 허용하여 다양한 입력에 대해 더 유연하고 동적인 아키텍처를 제공한다. 우리는 MCTS로 최적화된 CoLa에 대한 광범위한 분석을 수행했으며, 두 가지 주요 발견을 얻었다: (1) 원래 LLM이 정확히 예측한 샘플의 75% 이상에서 더 짧은 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.