Skip a Layer or Loop It? Learning Program-of-Layers in LLMs — part 3/4

Cleaned-text code points [31361, 46994). Bibliography omitted here; source.md includes it.

Transformers process inputs through sequential layer stacks, making layer-level computation reduction a critical research direction. Early-exit and layer skipping methods ( Liu et al., 2020 ; Xin et al., 2020 ; Zhou et al., 2020 ; Liu et al., 2021a ) dynamically terminate computation at intermediate layers using auxiliary classifiers and confidence metrics, allowing easy inputs to exit early. LayerSkip ( Elhoushi et al., 2024 ) shares classifiers across layers to reduce overhead. LayerDrop ( Fan et al., 2019 ) trains models so arbitrary layer subsets can be skipped during inference. ShortGPT ( Men et al., 2025 ) assesses layer importance based on input-output similarity and drops low-importance layers. LaCo ( Yang et al., 2024 ) merges layers using weight arithmetic. Recent work introduces learned routing for adaptive skipping: FlexiDepth ( Luo et al., 2025 ) and MindSkip ( He et al., 2024 ) attach lightweight routers to pretrained models for input-adaptive layer skipping.

In addition to skipping layers, another line of research explores layer reuse and recurrence. Universal Transformers ( Dehghani et al., 2018 ) apply self-attention blocks recurrently with halting mechanisms to adapt depth per token. Recent looped transformers ( Fan et al., 2024 ; Yang et al., 2023 ) repeatedly apply single blocks to achieve better length generalization on algorithmic tasks by adjusting loop counts during inference. The Inner Thinking Transformer ( Chen et al., 2025 ) interleaves adaptive loops with residual ”thinking” connections and per-token routing, devoting extra computation to difficult tokens. While these approaches demonstrate the value of recurrence, they require architectural redesign and training from scratch.

Li et al. (2025) studies test-time depth adaptation by using search to dynamically skip or repeat pretrained transformer layers without finetuning. Their work demonstrates that alternative execution paths can improve inference, but the method remains search-based and requires expensive per-input program discovery. In contrast, we use MCTS only as an offline diagnostic tool to characterize the structure of the program space, and then replace search with a learned predictor that generates execution programs in a single shot.

Following this direction, DR.LLM ( Heakl et al., 2025 ) learns routing policies from MCTS-generated supervision and supports both skipping and repeating layers. However, DR.LLM performs sequential layer-wise routing, where each decision is made locally during the forward pass and depends on intermediate hidden states. In contrast, PoLar predicts the entire execution program upfront, before executing the frozen LLM. This avoids interleaving routing with layer execution and enables more efficient inference. Moreover, DR.LLM is limited to single-layer recurrence, whereas PoLar operates on contiguous layer segments; for example, PoLar can represent multi-layer recurrent modules such as 4 → 5 → 4 → 5 4{\rightarrow}5{\rightarrow}4{\rightarrow}5 , which are outside the single-layer routing space. Thus, PoLar provides a more coordinated and expressive program space while preserving fully frozen base model parameters.

## 6 Conclusion

We show that inference in LLMs need not be limited to a fixed-depth forward pass. By viewing pretrained transformer layers as reusable functions, we uncover multiple valid execution programs for a single input, many of which are shorter than standard execution and can correct model errors. Motivated by this insight, we introduce PoLar , a lightweight framework that predicts input-dependent execution programs by selectively skipping or repeating contiguous layer segments at inference time, without modifying model parameters. Across models and both in-distribution and out-of-distribution benchmarks, PoLar consistently outperforms standard inference and prior dynamic-depth methods. These findings suggest that fixed-depth execution captures only a narrow subset of an LLM’s latent reasoning capacity. Enabling flexible, programmatic execution over pretrained layers reallocates computation at inference time, offering a simple and effective route to more expressive and efficient inference in foundation models.

## Impact Statement

This work adapts the internal computation of pretrained LLMs at inference time by dynamically skipping or repeating layer segments. Its main potential benefit is improved efficiency: PoLar can reduce unnecessary computation on easier inputs while allocating more latent computation to harder ones, lowering inference cost, latency, and energy use without retraining the base model. This may make capable LLMs more accessible to researchers and practitioners with limited compute, and may support more sustainable deployment of foundation models. As with other methods that improve LLM capability or efficiency, broader deployment may amplify both beneficial and harmful uses. Potential benefits include education, scientific reasoning, and software assistance, while potential misuse includes scalable generation of misleading or harmful content. These risks largely arise from the underlying pretrained models and their applications rather than from dynamic execution itself. Since PoLar changes the execution path in an input-dependent manner, future work may further study how such paths can be audited or interpreted. Our experiments focus on mathematical reasoning and related benchmarks. Before applying dynamic execution in high-stakes domains, future work should evaluate robustness, calibration, interpretability, and safety alongside accuracy and efficiency. We hope this work encourages more responsible test-time computation methods that improve model performance while making compute allocation more transparent and efficient.

## Appendix A Related Work

#### Layer Pruning and Early-Exit Neural Networks

Many works aim to accelerate large Transformers by statically pruning weights or dynamically halting computation. Static pruning typically removes redundant neurons, heads, or layers after training. For example, Liu et al. (2021b) demonstrate that a significant fraction of BERT’s attention heads can be dropped with negligible performance loss, and Gordon et al. (2020) investigate fine-grained weight pruning in BERT. ( Fan et al., 2019 ) introduce LayerDrop, a structured dropout technique that effectively trains models so arbitrary subsets of layers can be skipped during inference without requiring fine-tuning. These methods produce smaller models that trade computation for a small accuracy loss.

By contrast, early-exit or input-adaptive methods add auxiliary classifiers at intermediate layers so that ”easy” inputs exit early. Notable examples include FastBERT ( Liu et al., 2020 ) and DeeBERT ( Xin et al., 2020 ) , which insert classifiers after each block and use confidence or entropy metrics to decide when to stop. PABEE ( Zhou et al., 2020 ) employs a patience criterion to halt when predictions stabilize. DACT-BERT ( Eyzaguirre et al., 2022 ) adopts a differentiable Adaptive Computation Time mechanism to learn how many Transformer layers to run for each example. Liu et al. (2021a) estimate input ”hardness” via mutual information or reconstruction error to pre-determine the number of Transformer layers to use.

These early-exit networks achieve significant speedups on NLP tasks by adaptively reducing depth per input. More recently, early-exit ideas have been extended to vision and multimodal Transformers. Xu et al. (2023) propose LGViT, which adds heterogeneous exit heads (local and global) to ViT so that vision transformers can terminate early with minimal feature loss. Tang et al. (2023) introduce MuE (“Multiple Exiting”), a strategy for unified vision-language models that dynamically skips layers in both encoder and decoder based on input similarity. These works demonstrate that later layers can be skipped to allow image and vision-language models to adapt computation per sample with minimal accuracy drop. Our work generalizes this approach by allowing skipping of arbitrary layers and enabling reuse of certain layers.

#### Looped Transformer and Recurrent Depth

Another line of research makes Transformer depth adaptive by looping or repeating layers. The Universal Transformer ( Dehghani et al., 2018 ) was an early example: it applies the same self-attention block recurrently and uses a halting mechanism to determine when each position is “done” (adapting depth per token). Building on these ideas, recent work explicitly introduces loops in model architectures. Fan et al. (2024) demonstrate that a Looped Transformer – a single Transformer block applied repeatedly – can achieve much better length generalization on algorithmic tasks by adjusting the number of loops during inference. Similarly, Yang et al. (2023) note that looped architectures excel at learning algorithms by explicitly incorporating iterative characteristics into the transformer architecture. More sophisticated variants like the Inner Thinking Transformer Chen et al. (2025) interleave adaptive loops with residual “thinking” connections and per-token routing, enabling the model to devote extra computation only to particularly difficult tokens. In summary, these approaches explore recurrent or elastic depth via explicit loops to tailor the number of applied layers to each input’s complexity. Unlike our approach, they require special architecture design and training from scratch, whereas our work focuses on pure test-time adaptation.

#### Dynamic Routing and Modular Inference

A third theme treats networks as collections of modules or experts with dynamically chosen pathways per sample. Mixture-of-Experts (MoE) Transformer layers are a well-known example: they maintain multiple sub-networks (“experts”) and route each token to a subset. Wu et al. (2024) introduce Routing Experts (RoE) for multimodal LLMs, retrofitting trained models into a mixture-of-experts style by learning input-dependent shortcut routes through layers, guided by sparsity regularizers. Jain et al. (2024) present Mixture of Nested Experts (MoNE): experts organized in a hierarchy of increasing capacity, where tokens are sent to smaller experts when sufficient. MoNE learns to prioritize easy tokens through low-cost experts and reserve full models for hard cases, halving inference compute on ImageNet/Video tasks.

These methods exemplify sample-wise routing: at inference time, the model conditionally activates different sub-modules or experts for each input. Similarly, neural module networks ( Andreas et al., 2016 ) assemble task-specific computation graphs from a library of modules. In modern LLMs/VLMs, these routing approaches – whether through gating experts, skipping layers, or assembling modules – form a spectrum of modular inference techniques that adapt the computation graph on a per-sample basis to balance cost and accuracy. Interestingly, our work suggests that transformer layers can function effectively as modules even without being specifically trained for that purpose.

## Appendix B Searching the Execution Program Space

This appendix provides full details of the execution-program search procedure used to test the conjecture in Section 2 . The search is used purely as a diagnostic tool to study the existence and structure of valid execution programs, rather than as a practical inference-time method.

### B.1 Execution Program Space

We follow the formalization in the main text and represent inference as the execution of a variable-length program that composes pretrained transformer layers. Consider a pretrained LLM with D D transformer layers, where each layer defines a fixed computation function f i : ℝ T × d → ℝ T × d , i ∈ { 0 , … , D − 1 } . f_{i}:\mathbb{R}^{T\times d}\rightarrow\mathbb{R}^{T\times d},\quad i\in\{0,\ldots,D-1\}. An execution program is defined as a finite sequence of layer indices π = ( i 1 , i 2 , … , i K ) , i k ∈ { 0 , … , D − 1 } , \pi=(i_{1},i_{2},\ldots,i_{K}),\quad i_{k}\in\{0,\ldots,D-1\}, which induces the composed computation F π = f i K ∘ ⋯ ∘ f i 1 . F_{\pi}=f_{i_{K}}\circ\cdots\circ f_{i_{1}}. Executing a program applies this composition to the input representation and produces a prediction. A program is considered valid for a given input if it yields a correct prediction.

Programs may be shorter than the standard forward pass through layer skipping, or longer through layer repetition. Increasing program length corresponds to increasing the number of latent reasoning steps.

### B.2 Search Space Constraints

The unconstrained space of programs grows exponentially with program length. To make search tractable while preserving expressiveness, we restrict the action space to structured operations on contiguous subsequences of layer indices. Specifically, we allow two classes of actions: • Skip : remove a contiguous block of k k indices from the program;

• Repeat : duplicate a contiguous block of k k indices for r r repetitions.

In all experiments, block size k k and repetition count r r are bounded by small constants ( k , r ≤ 4 k,r\leq 4 ). These constraints significantly reduce the branching factor while retaining the ability to realize layer skipping, recurrence, and emergent reordering patterns.

### B.3 Monte Carlo Tree Search Formulation

We formulate program discovery as a sequential decision process and employ Monte Carlo Tree Search (MCTS) to explore the constrained program space.

#### State and Actions.

Each MCTS node corresponds to a partial or complete execution program π \pi . Actions modify the current program by applying a valid skip or repeat operation, yielding a new program.

#### Reward.

For a completed program π \pi and input x x with ground-truth answer y y , we define a binary reward r ( π , x ) = 𝟏 { F π ( x ) = y } , r(\pi,x)=\mathbf{1}\{F_{\pi}(x)=y\}, where F π ​ ( x ) F_{\pi}(x) denotes executing the composed computation induced by π \pi .

#### Tree Policy.

Tree traversal is guided by a UCB-style objective that balances exploitation, exploration, and program length regularization: UCB ⁡ ( π ) = R ⁡ ( π ) v ⁡ ( π ) + c ​ ln ⁡ V v ⁡ ( π ) − λ ​ | π | D , \mathrm{UCB}(\pi)=\frac{R(\pi)}{v(\pi)}+c\sqrt{\frac{\ln V}{v(\pi)}}-\lambda\frac{|\pi|}{D}, where R ⁡ ( π ) R(\pi) is the cumulative reward, v ⁡ ( π ) v(\pi) is the visit count, V V is the total number of simulations, and λ \lambda penalizes long programs to encourage efficiency.

### B.4 Search Algorithm

Algorithm 1 summarizes the MCTS procedure. We initialize the root node with the standard forward execution π 0 = ( 0 , 1 , … , D − 1 ) \pi_{0}=(0,1,\ldots,D-1) and iteratively perform selection, expansion, simulation, and backpropagation. After a fixed number of simulations, we collect all explored programs with nonzero visit counts and analyze their validity and structural properties.

## Appendix C Experimental Results

This appendix provides additional experimental results that complement the main paper. We report full quantitative comparisons for both in-distribution and out-of-distribution evaluations across multiple pretrained LLMs. Unless otherwise stated, all models are evaluated in a fully frozen setting, and PoLar only predicts execution programs at inference time.

### C.1 In-Distribution Performance

We first report detailed in-distribution results on DART-Math, a structured mathematical reasoning benchmark with five difficulty levels (DM-1 to DM-5). Tables 5 , 6 , and 7 present pass@k accuracy under different inference strategies for Qwen1.5-MoE-A2.7B-Chat, Qwen2.5-3B-Instruct, and Qwen3-8B, respectively.