Saltare un livello o ripeterlo? Adattamento della profondità in fase di test per LLM pre-addestrati

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

July 10, 2025

Autori: Ziyue Li, Yang Li, Tianyi Zhou

cs.AI

Abstract

Una rete neurale preaddestrata può adattare la propria architettura a diversi input senza alcun fine-tuning? Abbiamo bisogno di tutti i livelli per compiti semplici, e sono adeguati per compiti complessi? Abbiamo scoperto che i livelli di un modello linguistico preaddestrato di grandi dimensioni (LLM) possono essere manipolati come moduli separati per costruire un modello migliore e persino più superficiale, personalizzato per ogni campione di test. In particolare, ogni livello del modello preaddestrato può essere saltato/prunato o ripetuto più volte come nelle reti neurali ricorrenti (RNN), e combinato con altri in ordini arbitrari, producendo una catena di livelli (CoLa) per ogni campione. Questo spazio compositivo amplia notevolmente l'ambito dei lavori esistenti sui moduli preaddestrati a ciclo/ricorrenti, sulla potatura dei livelli o sulle reti con uscita anticipata. Abbiamo sviluppato un protocollo di Monte Carlo Tree Search (MCTS) per esplorare e identificare la CoLa ottimale per ogni campione proveniente da benchmark di ragionamento matematico e di senso comune. Rispetto a un modello statico di profondità fissa, CoLa consente percorsi abbreviati (pensiero veloce), la ripetizione dello stesso livello/i (pensiero lento) e la combinazione di entrambi, offrendo architetture più flessibili e dinamiche per diversi input. Abbiamo condotto un'analisi approfondita della CoLa ottimizzata con MCTS, che ha portato a due risultati chiave: (1) Per oltre il 75% dei campioni con previsioni corrette da parte dell'LLM originale, possiamo trovare CoLa più brevi, suggerendo un ampio spazio per migliorare l'efficienza dell'inferenza; (2) Per oltre il 60% dei campioni con previsioni originariamente errate, possiamo identificare CoLa che raggiungono previsioni corrette, suggerendo un ampio spazio di miglioramento delle prestazioni. I nostri risultati evidenziano le carenze nell'utilizzo di un'architettura fissa di LLM preaddestrati per l'inferenza su diversi campioni e aprono la strada per sbloccare il potere di generalizzazione dell'adattamento della profondità al momento del test.

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.