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

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

# Skip a Layer or Loop It? Learning Program-of-Layers in LLMs

###### Abstract

Large language models (LLMs) perform inference by following a fixed depth and order, non-recurrent execution of all layers. We reveal the wide existence of training-free, flexible, dynamic “program-of-layers ( PoLar )”, where pretrained layers can be packed as modules and then skipped or looped to form a customized program for each input. For most inputs, substantially shorter program executions can achieve the same or better accuracy, while incorrect predictions of the original LLM can be corrected by alternative programs with fewer layers. These observations indicate that inference admits multiple valid latent computations beyond the standard forward pass. To efficiently achieve PoLar in practice, we propose a lightweight PoLar prediction network, which learns to generate execution programs that dynamically skip or repeat pretrained layers for each input. Experiments on mathematical reasoning benchmarks demonstrate that PoLar consistently improves accuracy over standard inference and prior dynamic-depth methods, often while executing fewer layers, and that these gains persist under out-of-distribution evaluation. Our results suggest that fixed-depth execution captures only a narrow subset of an LLM’s latent reasoning capacity.

###### Keywords:

Project: https://github.com/tianyi-lab/PoLar

## 1 Introduction

Generalist foundation models, e.g., LLMs and VLMs, uniformly deploy a static, pre-defined architecture to all inputs, despite their diversity and high variance in complexity and difficulty ( Liu et al., 2020 ; Xin et al., 2020 ; Zhou et al., 2020 ; Liu et al., 2021a ) . In contrast, conventional problem solving by programs can be more flexible and adaptive in algorithmic structures and complexity. For example, an experienced programmer can save more steps and compute on easier tasks, and meanwhile knows how to scale up the space/time complexity to address more challenging problems. However, these programs are specifically designed and optimized for every problem class, so they are not as general as LLMs. This raises the questions: Is it always optimal and efficient to apply the same architecture or “program”, i.e., forward pass through all the layers in a fixed order, to different tasks? Can a generalist model further optimize its “program” applied to each input?

In this paper, we formulate layers in a pretrained LLM as a library of atomic functions that a program can call in arbitrary order for arbitrary times. This formulation allows us to represent a dynamic model architecture during inference as a program-of-layers ( PoLar ) for each input, as illustrated in Figure 1 . As the first empirical study of its kind, we investigate PoLar beyond the standard forward pass by Monte-Carlo Tree Search (MCTS) and find that better (more accurate and/or shorter) programs almost always exist for every input task evaluated. Unlike previous works on layer-skipping/recurrence, early exit, and looped transformer ( Liu et al., 2020 ; Xin et al., 2020 ; Zhou et al., 2020 ; Fan et al., 2019 ; Fan et al., 2024 ; Yang et al., 2023 ) , which only adopt one operation (either skip or repeat) to produce architectures of dynamic depths, our empirical study on the MCTS-searched programs reveals that searching in a joint space of layer-skip/repeat often discovers much better programs than those found in separate spaces. While most effective programs can be shorter than the default, increasing program complexity via skip/repeat operations can substantially improve the output quality, especially on more difficult tasks. In addition, most successful programs are predominantly composed of contiguous layer segments. These observations not only verify the broad existence of better PoLar without requiring any training, but also motivate a practical PoLar prediction method that avoids the expensive cost of MCTS in PoLar ’s large search space. In particular, we aim to replace search-based program discovery with a direct, inference-time mechanism for generating execution programs. Instead of enumerating or exploring execution paths for each input ( Li et al., 2025 ) , our goal is to predict an input-specific program-of-layers that determines how pretrained layers are executed during inference. This shifts program selection from an online search problem to a single-shot prediction problem, enabling practical deployment of program-of-layers inference in LLMs.

To this end, we propose a PoLar algorithm that predicts execution programs over frozen pretrained layers at inference time. The predicted execution program specifies how pretrained layers are selectively skipped or recurrently applied and is executed once to produce the final output. This design offers several advantages. First, it makes program-of-layers inference computationally feasible by eliminating the need for expensive per-input search. Second, by jointly supporting layer skipping and recurrence within a unified execution framework, PoLar strictly generalizes prior dynamic-depth methods that are limited to a single form of execution control. Third, it enables flexible test-time computation scaling in fully frozen models, allowing inference to adapt to input difficulty while preserving model generality.

We evaluate PoLar on a range of mathematical reasoning benchmarks using multiple pretrained LLMs. Our results show that PoLar consistently improves accuracy over standard inference and prior dynamic-depth methods, often while executing fewer layers on average. Moreover, increasing the number of candidate execution programs yields strong test-time computation scaling, and execution programs learned on in-distribution data generalize effectively to out-of-distribution benchmarks across diverse domains.

## 2 Dynamic Inference as a Program-of-Layers ( PoLar ) in Large Language Models

Inference in pretrained LLMs is implemented as a fixed-depth, fixed-order forward pass: every input is processed by executing the same sequence of transformer layers. Yet inputs to LLMs vary dramatically in difficulty. Some are answered correctly with minimal reasoning, while others require complex, multi-step computation. This discrepancy raises a basic question: Is the standard forward pass sufficient for correct inference across diverse inputs?

One possibility is that this fixed computation is indeed sufficient for all cases. Another is that correct prediction requires input-dependent variation in computation. In this work, we investigate the latter possibility. Such variation can occur either in token space, through longer and more explicit chains of thought, or within the model’s hidden states, a form of computation we refer to as latent reasoning .

Inference as the execution of a program . In this view, inference is a step-by-step procedure that selects and composes pretrained modules. The execution may vary across inputs in both length and order, while each module remains a fixed, pretrained function.

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 } . \textstyle f_{i}:\mathbb{R}^{T\times d}\rightarrow\mathbb{R}^{T\times d},\quad i\in\{0,\ldots,D-1\}. A program is defined as a finite sequence of layer indices π = ( i 1 , i 2 , … , i K ) , i k ∈ { 0 , … , D − 1 } , \textstyle\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 . \textstyle F_{\pi}=f_{i_{K}}\circ\cdots\circ f_{i_{1}}. Executing a program applies this composition to the input and produces a prediction. A program is considered valid if it yields a correct prediction for a given input.

Searching for valid execution programs. We explore the space of execution programs using MCTS. Execution programs are variable-length sequences over pretrained transformer layers, allowing both skipping and repetition. This space is large, discrete, and highly non-convex, making exhaustive search infeasible. MCTS provides a principled way to prioritize promising partial programs, enabling us to verify the existence of valid programs and analyze their structural properties. We use MCTS strictly as a diagnostic tool rather than a practical inference-time method; implementation details are given in Appendix B . All experiments are conducted on DART-Math ( Tong et al., 2024 ) , a structured mathematical reasoning benchmark with five difficulty levels (DM-1 to DM-5). We evaluate four pretrained transformer models: LLaMA-3.2-3B-Instruct , Qwen1.5-MoE-A2.7B-Chat , Qwen2.5-3B-Instruct , and Qwen3-8B .

We summarize our empirical findings below.

As shown in Table 1 , execution programs that allow layer recurrence ( Loop ) consistently outperform those that allow only layer skipping ( Skip ) across all evaluated models and difficulty levels. Moreover, combining skipping with recurrence ( Skip&Loop ) yields substantially larger gains than either operation alone, achieving the highest accuracy in every setting reported in Table 1 . These results indicate that recurrence provides a stronger mechanism for improving inference than skipping alone, while the two operations play complementary roles when combined.

As shown in Figure 3 , which reports the best accuracy obtained by MCTS under explicit budgets on total layer executions, many inputs remain solvable even when the overall computation is constrained to be significantly shorter than the standard forward pass. Consistent with this trend, Figure 4 shows that across models we frequently discover valid execution programs that require fewer layer applications than standard inference. In particular, among inputs already solved correctly by standard inference (C → \rightarrow C), 71.9% admit shorter valid programs. Even for inputs initially solved incorrectly (W → \rightarrow C), 34.0% admit shorter programs that correct the model’s prediction. These results indicate that standard inference often over-computes, and that correct inference can frequently be achieved with substantially fewer latent computation steps.

While simple inputs often admit short execution programs (Finding 2), harder inputs demand greater test-time execution complexity. Across models and datasets, increasing execution depth and structural flexibility systematically expands valid program space and improves inference accuracy.

(a) Test-time scaling expands the space of valid execution programs for latent reasoning. As shown in Figure 5 (a), allocating more test-time computation—via recurrence—monotonically increases the existence of valid execution programs across models. This establishes test-time scaling at latent reasoning: greater computation yields a larger feasible program space and higher correctness.

(b) Harder inputs require more complex execution programs. Figure 5 (b) shows that the fraction of solvable inputs that require recurrence and/or skipping generally increases with dataset difficulty for most models. As task difficulty grows, valid execution programs become more constrained and increasingly rely on non-trivial execution structures rather than standard inference. This indicates that higher latent execution complexity is not merely helpful, but often necessary for solving harder inputs. The different trend observed for LLaMA-3.2-3B-Instruct is explained by a mismatch between dataset-defined difficulty and the model’s effective difficulty (Table 1 ).

(c) Inference accuracy improves systematically with execution depth. As shown in Figure 6 , for all models and difficulty levels, the average accuracy of valid execution programs increases with total execution depth, measured relative to the original forward-pass depth. This reveals a consistent computation–accuracy trade-off: while many inputs admit short execution programs (Finding 2), harder inputs benefit from—and often require— deeper or recurrent execution to achieve correct inference.

Latent execution programs reveal a continuum of test-time inference behaviors that standard inference cannot access. Allocating more execution complexity enables harder inputs to be solved and yields higher accuracy.

As shown in Figure 7 , valid execution programs discovered by pretrained models exhibit a strong structural bias toward simplicity. A segment denotes a set of layers executed as a unit and need not be contiguous, while recurrence always corresponds to re-execution of the same segment. We therefore analyze segment structure by measuring the number of consecutive layers within each segment. Most valid programs are dominated by highly local segments and involve at most a single recurrence; long-range jumps and deep iterative reuse are rare. Figure 7 (a) shows that 57.7% of segments consist of a single layer, and over two-thirds contain at most two consecutive layers, whereas segments with predominantly non-consecutive layers account for less than 2.9% of cases. Consistently, Figure 7 (b) indicates that most segments are repeated at most once. Together, these results reveal an inherent limitation of pretrained models as execution-program generators: their training objectives favor short-range, local reuse over rich program composition and complex control flow.

These findings show that standard inference selects only one execution from a vast space of valid latent programs. While MCTS reveals this space, its reliance on sequential search over an exponentially large program space makes it impractical for inference. This motivates a different approach: rather than searching over programs at test time, we ask whether a lightweight model can directly predict execution programs. Figure 2 contrasts the MCTS-based sequential search with our proposed direct program prediction approach. In the remainder of this work, we pursue this learning-based alternative, retaining the benefits of latent program selection uncovered by MCTS while eliminating sequential search.

## 3 Learning Program-of-Layers ( PoLar ) in Large Language Models

Building on our empirical analysis (Section 2 ), we propose PoLar , a method for programming pretrained language models at inference time by predicting input-specific execution programs (Figure 2 ). PoLar dynamically segments and composes pretrained layers into reusable modules, enabling flexible computation without parameter updates.

### 3.1 Program Representation

We instantiate the function library using packed modules , which segment contiguous pretrained transformer layers into reusable computation units. For a pretrained model of depth D D , an execution program specifies (i) a segmentation of layers into modules and (ii) an operation applied to each segment. Each execution program is represented by two discrete structures: a binary boundary mask encoding the segmentation, and an operation label vector specifying the segment-level operations.

Segmentation. We partition the D D layers of a pretrained model into contiguous segments [ 0 = s 1 , s 2 ) , [ s 2 , s 3 ) , … , [ s M , s M + 1 = D ) , \textstyle[0=s_{1},s_{2}),\ [s_{2},s_{3}),\ \ldots,\ [s_{M},s_{M+1}=D), with each segment length bounded by s j + 1 − s j ≤ K max s_{j+1}-s_{j}\leq K_{\max} . Segmentation is represented by a binary boundary mask 𝐳 seg ​ ( x ) ∈ { 0 , 1 } D , \textstyle\mathbf{z}^{\text{seg}}(x)\in\{0,1\}^{D}, where 𝐳 i seg = 1 \mathbf{z}^{\text{seg}}_{i}=1 indicates that layer index i i starts a new segment, and 𝐳 i seg = 0 \mathbf{z}^{\text{seg}}_{i}=0 otherwise.