cola-test-time-depth-adaptation-llm

name	cola-test-time-depth-adaptation-llm
title	Skip a Layer or Loop it? Test-Time Depth Adaptation of Pretrained LLMs
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2507.07996
keywords	["LLM Architecture","Depth Adaptation","Test-Time Optimization","Monte Carlo Tree Search","Layer Reordering"]
description	Dynamically adapt LLM depth per input at test time by skipping, repeating, or reordering layers using MCTS search, correcting 60% of initially wrong predictions and processing 75% of correct predictions with shorter architectures without retraining.

CoLa: Chain-of-Layers Architecture Search for Test-Time Adaptation

Standard language models apply all layers sequentially. But not every input needs every layer. Difficult examples benefit from depth, easy examples waste computation. CoLa (Chain-of-Layers) reframes inference as an architecture search problem: for each test input, find the optimal layer sequence by skipping unnecessary layers, repeating useful ones, or reordering them. Monte Carlo Tree Search efficiently explores this space without training, finding custom architectures that improve accuracy or reduce latency.

The method reveals that over 75% of correctly predicted samples could be processed through shorter paths, and over 60% of initially incorrect predictions become correct through architectural reconfiguration. This enables substantial efficiency gains or accuracy improvements depending on your objectives.

Core Concept

The key insight is that transformer layers are relatively independent modules. A sequence of layers is not sacred; the same model can process inputs via different layer orderings. CoLa searches for the optimal "chain of layers" for each input: which layers to include, which to skip, which to repeat. This search happens at test time via MCTS, a planning algorithm that balances exploration (trying new architectures) and exploitation (focusing on promising ones).

For easy inputs, this finds shorter paths. For hard inputs, it identifies beneficial layer reorderings or repetitions. The method imposes no training overhead and works with frozen pretrained models.

Architecture Overview

• Layer Manipulation Space: Skip, repeat, or reorder individual layers or layer blocks
• MCTS Search: Exploration-exploitation via Upper Confidence Bound over architecture space
• State Representation: Current layer index, already-processed layers, input progress
• Action Space: Skip 1-4 layers, repeat 1-4 layers, continue with next layer
• Reward Signal: Model accuracy on input (did you predict correctly?)
• Simulation Budget: 200 MCTS simulations per input, ~5x inference time typical
• Compatible Layers: Works with LLaMA-3, OLMoE, and other transformer families

Implementation

Step 1: Define Layer Manipulation Operations

Create the action space: all valid ways to modify layer sequences:

import torch
import torch.nn as nn
from typing import List, , 
 dataclasses  dataclass
 enum  Enum

 ():
    
    CONTINUE =       
    SKIP =               
    REPEAT =           
    REORDER =         


 :
    
    current_layer_idx:       
    used_layers: []      
    hidden_state: torch.Tensor  
    num_steps:              

 :
     ():
        .model = model
        .num_layers = num_layers
        .max_skip = max_skip
        .max_repeat = max_skip

     () -> [[LayerAction, ]]:
        
        actions = []

        
         state.current_layer_idx < .num_layers - :
            actions.append((LayerAction.CONTINUE, ))

        
         k  (, .max_skip + ):
             state.current_layer_idx + k < .num_layers:
                actions.append((LayerAction.SKIP, k))

        
         state.current_layer_idx < .num_layers:
             k  (, .max_repeat + ):
                actions.append((LayerAction.REPEAT, k))

        
         k  (, .max_skip + ):
             state.current_layer_idx + k < .num_layers:
                actions.append((LayerAction.REORDER, k))

         actions

     () -> LayerState:
        
        new_used_layers = state.used_layers.copy()
        new_idx = state.current_layer_idx

         action == LayerAction.CONTINUE:
            new_used_layers.append(new_idx)
            new_idx += 

         action == LayerAction.SKIP:
            new_idx += param

         action == LayerAction.REPEAT:
            new_used_layers.extend([new_idx] * param)

         action == LayerAction.REORDER:
            
             i  (param):
                new_used_layers.append(new_idx + param -  - i)
            new_idx += param

        
        hidden = state.hidden_state
         layer_idx  new_used_layers[(state.used_layers):]:
            layer = .model.model.layers[layer_idx]
             torch.no_grad():
                hidden = layer(hidden)[]

         LayerState(
            current_layer_idx=new_idx,
            used_layers=new_used_layers,
            hidden_state=hidden,
            num_steps=state.num_steps + 
        )