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

Abstract: 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.

• The paper demonstrates that dynamic execution programs, using skip and repeat operations, significantly improve accuracy and efficiency in LLMs.
• It employs Monte Carlo Tree Search and a PoLar prediction network to predict optimal layer execution paths without modifying base model parameters.
• Results reveal up to +20.8 pass@5 accuracy improvement and a 17% runtime reduction, highlighting beneficial computation-accuracy tradeoffs.

Program-of-Layers: Dynamic Execution Programs for LLMs

Motivation and Problem Formulation

Inference in LLMs is universally implemented via a fixed-depth, strictly ordered feed-forward pass through all transformer layers, regardless of the input's complexity or difficulty. This paper investigates the hypothesis that this paradigm significantly underutilizes the latent reasoning capacity of pretrained LLMs. Specifically, transformer layers are conceptualized as a library of atomic functions, allowing for the composition of dynamic, input-specific programs-of-layers (PoLar), which skip or repeat layers to form latent execution programs for each input.

Figure 1 visualizes this idea:

Figure 1: Program-of-layers for two distinct inputs, each defining its own execution path over the model's layers by selectively skipping or repeating them.

Empirical investigation shows that, for most tasks, valid latent computations exist beyond the standard forward pass—distinct layer execution programs that are shorter, correct erroneous outputs, or yield higher accuracy. This expands the conventional search space of LLM inference beyond static per-layer application, motivating both a search-based and a learned approach for discovering such alternate programs.

Monte Carlo Search and Emergent Execution Properties

The authors first explore the PoLar hypothesis with Monte Carlo Tree Search (MCTS), allowing discrete choices of skip and repeat operations over variable-length, contiguous segments. The MCTS search empirically substantiates four key findings:

1. Layer recurrence is more effective than skipping, but their combination yields the highest accuracy. This joint action space (skip + recurrence) consistently surpasses prior methods constrained to a single type of dynamic control.
2. Occam’s razor effect: Valid execution programs are frequently shorter than the standard forward pass, leading to computational savings without sacrificing accuracy. Most correctly answered examples admit even shorter solution programs.
3. Computation–accuracy tradeoff: Increasing the allowance for recursion and latent execution depth expands the set of valid programs and systematically improves the accuracy, especially on harder tasks (Figure 2, Figure 3, Figure 4).

   Figure 2: MCTS shows that accuracy can be maintained or improved with programs shorter than the original forward pass across increasing difficulty levels.

   

   Figure 3: (a) Allowing more recurrent execution increases valid program probability; (b) complex routing (recurrence/skipping) is required for harder tasks.

   

   Figure 4: Accuracy consistently increases with total layer executions, emphasizing the computation–accuracy tradeoff.

4. Structural simplicity: The valid execution programs most often consist of contiguous layer segments, rarely requiring more than one repeat per segment; arbitrary or deep recurrences are uncommon (Figure 5).

   Figure 5: Most valid programs are built from contiguous segments and employ at most a single recurrence per module.

This diagnostic phase demonstrates the prevalence of alternate, often more efficient and accurate, program-of-layers executions. However, MCTS is impractical for test-time inference due to its combinatorial complexity (Figure 6).

Figure 6: MCTS (left) imposes high search costs; PoLar (right) directly predicts execution programs in a single shot.

The PoLar Network: Learning to Predict Execution Programs

To make dynamic execution feasible in practice, the paper introduces the PoLar prediction network. It reframes execution path selection as a structured prediction problem: for each input, predict a segmentation of the pretrained layers into modules, and then assign to each segment an operation—skip, keep, or repeat. This representation captures the empirical regularities: most segments are short and local; only a single recurrence is allowed.

The PoLar network employs a lightweight cross-attention-based encoder over the input, producing per-layer scores for segment boundaries and operation selection. The training signal is provided by MCTS-discovered valid programs, parsed into segmentation and operation labels. Inference combines deterministic segmentation with a beam search over per-segment operations, yielding a ranked set of candidate execution programs.

This design generalizes all prior dynamic inference methods: it unifies and extends skipping and recurrence in a single framework, enables flexible compute scaling without modifying base model parameters (zero-update/frozen setting), and shifts from expensive online search to amortized, test-time prediction.

Experimental Results

Benchmarks and Setup: Experiments are conducted on DART-Math (five difficulty levels), arithmetic word problem datasets (ASDiv, MAWPS), and MMLU-Pro covering a wide range of domains. Models span multiple LLM families (LLaMA, Qwen, etc.), all evaluated in a fully frozen setting.

In-distribution results: PoLar consistently and substantially outperforms both the static forward pass and prior dynamic-depth methods. For LLaMA-3.2-3B-Instruct, PoLar yields up to +20.8 pass@5 accuracy improvement over the strongest stochastic decoding baseline and demonstrates monotonic accuracy improvements as more candidate programs are considered, showing effective test-time scaling (Figure 7).

Figure 7: (a) Pass@k accuracy for PoLar dominates base model on LLaMA-3.2-3B-Instruct; (b) most PoLar solutions require fewer unique layers than the full model.

Efficiency: PoLar often solves the same inputs with fewer unique executed layers than the original model path, achieving up to a 17% reduction in runtime, and the overhead of the prediction network is negligible (<1% of the model’s compute cost).

Out-of-distribution generalization: When trained only on DART-Math, PoLar demonstrably generalizes its learned execution strategies to unseen domains (arithmetic and general knowledge benchmarks), yielding robust, transferable gains.

Contrast to prior methods: Simplistic layer-skipping (ShortGPT, MindSkip, FlexiDepth) and layer-wise routers (DR.LLM) are consistently outperformed, especially as the program search space and program complexity increase. PoLar’s segment-level control and program-level view enables both improved accuracy and greater parameter efficiency.

Theoretical and Practical Implications

The results unambiguously show that fixed-depth inference captures only a restricted subset of an LLM’s latent reasoning capacity. The existence of multiple valid—often shorter, sometimes corrective—execution programs implies that a pretrained LLM encodes an under-exploited computational basis. By enabling program-level dynamic routing over the frozen layers, PoLar unleashes this latent modularity with minimal additional compute or overhead.

Practically, PoLar supports efficient, input-adaptive inference, translating to lower compute, reduced latency, and potentially improved interpretability of model internals. The flexibility to reallocate computation at inference time, without any model finetuning, also improves generalization to out-of-distribution domains.

Theoretically, this work directly challenges the view of transformer networks as strictly monolithic and opens investigation into modular, reusable, and adaptive architectures—even in the absence of explicit modular training objectives. It also raises open questions regarding the optimality of learned versus searched execution programs, tradeoffs in program space regularization, and the interplay with emerging modular and expert-based transformer families.

Future Directions

Among several promising avenues are: (1) extending PoLar to non-contiguous modular compositions for richer control flow; (2) integrating semantic or task descriptors into program generation for explicit task adaptation; (3) interpreting learned execution programs for grounding mechanistic understandings of LLM inference; (4) combining program-of-layers with MoE or architecture search in a unified framework; and (5) exploring dynamic execution in high-risk, safety-critical domains with emphasis on transparency and calibration.

Conclusion

This paper demonstrates that standard fixed-depth transformer inference is an arbitrary point in a vast, rich space of valid latent computation programs. By empirically uncovering and learning to predict dynamic program-of-layers—a selective, per-input, skip/repeat execution program—PoLar unlocks significant gains in accuracy, efficiency, and transfer, all without modifying the underlying parameters of pretrained LLMs. This points to a broader principle: flexible, programmatic execution is critical to fully leveraging the capacity of foundation models.