PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ziyue Li, Yang Li, Tianyi Zhou

arXiv:2606.06574v2cs.LG

TL;DR

LLMs normally apply the same fixed-depth, fixed-order layer sequence to every input, leaving open whether this is sufficient for inputs with different computational demands. The paper introduces POLAR to predict input-specific programs that skip or repeat pretrained layers, consistently improving accuracy over standard and prior dynamic-depth inference, often with fewer executed layers.

  • Problem

    It is unclear whether applying the same fixed-depth, fixed-order forward pass to diverse inputs is always sufficient or efficient.

  • Method

    POLAR predicts input-specific execution programs over frozen pretrained layers, dynamically skipping or repeating layer segments without parameter updates.

  • Results

    Across models and in-distribution and out-of-distribution benchmarks, POLAR consistently outperforms standard inference and prior dynamic-depth methods; at pass@5, it achieves 51.2% versus Base sampling’s 43.8%, a +7.4 percentage-point gain.

  • Takeaways & Limitations

    Fixed-depth execution captures only a narrow subset of an LLM’s latent reasoning capacity.

  • Takeaways & Limitations

    Sequential search over the exponentially large execution-program space is impractical for inference, motivating POLAR’s direct program prediction.

Abstract

from arXiv · show

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.

1. Introduction

The paper challenges fixed-depth, fixed-order inference by treating pretrained LLM layers as reusable functions that can form input-specific programs with skipped or repeated layers. It shows that such programs often improve accuracy and reduce execution depth, and introduces a predictor that generates them efficiently and generalizes across benchmarks.

  • Motivation: Generalist foundation models apply the same static architecture to inputs with widely varying complexity and difficulty.The introduction contrasts this uniform deployment with the flexibility of problem-specific programs.
  • Program-of-layers formulation: POLAR represents pretrained LLM layers as atomic functions callable in arbitrary orders and repetitions, defining an input-specific dynamic architecture.Its execution programs can skip layers and recurrently apply them, rather than following the standard fixed sequence.
  • Empirical investigation: MCTS finds that better programs—more accurate, shorter, or both—almost always exist across evaluated input tasks.The study examines programs beyond the standard forward pass and combines layer skipping with recurrence.
  • Efficient inference: POLAR predicts execution programs over frozen pretrained layers, making dynamic inference feasible without expensive per-input search.The predicted program selectively skips or recurrently applies layers and executes once to produce the final output.
  • Results: POLAR consistently improves accuracy over standard inference and prior dynamic-depth methods, often while executing fewer layers, with gains scaling to more candidate programs and generalizing out of distribution.Evaluations use multiple pretrained LLMs on mathematical reasoning benchmarks, including out-of-distribution benchmarks across diverse domains.

2. Dynamic Inference as a Program-of-Layers (POLAR) in Large Language Models

POLAR treats LLM inference as an input-dependent program over pretrained layers, allowing layers to be skipped or repeated instead of following one fixed forward pass. Search results show that multiple valid programs exist, often with shorter executions, while harder inputs benefit from greater latent execution complexity.

  • Program-of-Layers: POLAR models inference as variable-length, input-dependent sequences of fixed pretrained layer functions that can skip or repeat layers.A program composes selected layer functions, and it is valid when it produces a correct prediction for the input.
  • Program-of-Layers: MCTS searches the large, discrete program space to verify valid executions, but its sequential search is impractical for inference.This motivates directly predicting programs with a lightweight network rather than searching at test time.
  • Program Operators: Skip&Loop achieves the highest accuracy in every reported setting, outperforming Loop and Skip individually across models and difficulty levels.Layer recurrence is stronger than skipping alone, while combining both operations provides complementary benefits.
  • Valid Programs: 71.9% of C→C inputs admit shorter valid programs, while 34.0% of W→C inputs admit shorter programs that correct the model.These findings indicate that standard inference often over-computes and that alternative latent executions can repair incorrect predictions.
  • Latent Execution Complexity: Increasing recurrence, execution depth, and structural flexibility expands the valid-program space and improves accuracy, especially for harder inputs.The fraction of inputs requiring recurrence or skipping generally rises with dataset difficulty, and valid-program accuracy increases with total execution depth.

3. Learning Program-of-Layers (POLAR) in Large Language Models

POLAR programs pretrained language models at inference time by dynamically composing contiguous layer modules and assigning each segment a skip, keep, or repeat operation. A lightweight predictor learns input-specific programs from offline valid executions and decodes globally consistent candidates at inference time.

  • Program representation: POLAR dynamically segments pretrained layers into reusable modules and composes them into input-specific execution programs without parameter updates.Each program specifies both a layer segmentation and an operation for each segment.
  • Program representation: Each layer segment receives one of three operations—skip, keep, or repeat—respectively omitting, preserving, or adding one execution pass.The repeat operator permits at most one additional execution within a segment.
  • Program prediction: The predictor encodes inputs with a frozen embedding model, uses layer-specific queries and cross-attention, and applies cross-layer self-attention for depth-aware decisions.The cross-layer encoder produces input-conditioned representations for each pretrained layer index.
  • Program prediction: Two linear heads predict segmentation boundaries and segment operations, supervised by offline valid programs collected through MCTS.Segmentation uses boundary labels, while operation prediction is supervised only at segment start positions; the final loss is L = Lseg + Lop.
  • Program decoding: At inference, POLAR thresholds boundary logits, enforces the maximum segment length, and uses beam search to produce globally consistent ranked execution programs.Beam search accounts for non-local interactions between segment-level operation choices before mapping candidates to concrete execution paths.

4. Experiments

Experiments evaluate POLAR on frozen LLMs across in-distribution and substantially different out-of-distribution mathematical benchmarks. POLAR improves accuracy, supports test-time scaling through program exploration, transfers across domains, and adds little inference overhead.

  • Experimental setup: POLAR is evaluated on four fully frozen instruction-tuned LLMs spanning architectures and scales, including LLaMA-3.2-3B-Instruct and Qwen3-8B.The models are LLaMA-3.2-3B-Instruct, Qwen1.5-MoE-A2.7B-Chat, Qwen2.5-3B-Instruct, and Qwen3-8B, with no parameter updates.
  • Experimental setup: The evaluation uses DART-Math for in-distribution testing and ASDiv, MAWPS, and MMLU-Pro subject subsets for out-of-distribution testing across differing formats and domains.OOD evaluation trains on the union of DART-Math training data and evaluates zero-shot, testing transfer beyond dataset- or difficulty-specific heuristics.
  • In-distribution results: At pass@1 on DART-Math DM-1, POLAR improves accuracy from 48.9% to 54.6% over Base (sampling), a +5.7 percentage-point gain.Because pass@1 evaluates one decoded output, the gain reflects latent execution selection rather than output-space diversity.
  • In-distribution results: At pass@5, Base (sampling) reaches 43.8% while POLAR reaches 51.2%, after improving from 32.8% and 35.1% at pass@1, respectively.Increasing the number of candidate execution programs monotonically improves POLAR, demonstrating test-time scaling through execution-program exploration.
  • Efficiency: The encoder, predictor head, and beam search add 3.05 ms, or 0.8% of a standard forward pass and approximately 0.23 LLM layers.This overhead is described as negligible relative to the latency reduction from executing fewer layers.

5. Related Works

Prior work reduces transformer computation through early exit, layer skipping, recurrence, or search-based depth adaptation. POLAR instead predicts a complete execution program before running the frozen LLM, avoiding sequential routing during layer execution.

  • Layer skipping: Early-exit and layer-skipping methods dynamically terminate computation at intermediate layers using auxiliary classifiers or confidence metrics.LayerSkip shares classifiers across layers to reduce overhead.
  • Layer reuse and recurrence: Universal Transformers and looped transformers reuse layers recurrently, adjusting depth or loop counts for adaptive computation and better algorithmic length generalization.Universal Transformers use halting mechanisms to adapt depth per token.
  • Layer reuse and recurrence: Adaptive-loop architectures such as the Inner Thinking Transformer route extra computation to difficult tokens but require architectural redesign and training from scratch.The method combines residual thinking connections with per-token routing.
  • Test-time depth adaptation: Li et al. use search to skip or repeat pretrained layers without finetuning, but require expensive per-input program discovery.POLAR uses MCTS offline to characterize the program space before replacing search with a learned predictor.
  • Learned routing: DR.LLM learns skipping and repeating policies from MCTS supervision, whereas POLAR predicts the entire execution program upfront before executing the frozen LLM.DR.LLM makes sequential layer-wise decisions using intermediate hidden states; POLAR avoids interleaving routing with layer execution.

6. Conclusion

The conclusion argues that LLM inference can use multiple valid, input-specific execution programs rather than a fixed-depth forward pass. It introduces POLAR to predict these programs by skipping or repeating contiguous layer segments at inference time.

  • Pretrained transformer layers can be viewed as reusable functions that form multiple valid execution programs for a single input.
  • Many alternative programs are shorter than standard execution and can correct model errors.
  • POLAR predicts input-dependent programs by selectively skipping or repeating contiguous layer segments at inference time.

Impact Statement … B.4. Search Algorithm

The paper frames PoLar as test-time adaptation that dynamically skips or repeats pretrained layers, and details a constrained MCTS procedure for discovering valid execution programs. This approach targets efficient, input-adaptive computation without retraining the base model.

  • Impact Statement: PoLar dynamically skips or repeats pretrained layer segments, reducing unnecessary computation on easy inputs while allocating more computation to hard ones without retraining.The potential benefits include lower inference cost, latency, and energy use, with improved accessibility for compute-limited researchers and practitioners.
  • A. Related Work: Prior work accelerates Transformers through static pruning, early exits, recurrent depth, or modular routing, whereas PoLar focuses on pure test-time adaptation.Early-exit methods add intermediate classifiers, while recurrent-depth approaches and routing systems generally require specialized architecture design or training.
  • B. Searching the Execution Program Space: The search appendix treats execution-program discovery as a diagnostic tool for studying valid program existence and structure, not as a practical inference-time method.It formalizes inference as composing pretrained Transformer layers into variable-length programs and uses correctness on each input to define validity.
  • B.1. Execution Program Space: An execution program is a finite sequence of layer indices whose composition produces a prediction, with shorter programs skipping layers and longer programs repeating them.Increasing program length is interpreted as increasing the number of latent reasoning steps.
  • B.2. Search Space Constraints: The constrained search permits skipping or repeating contiguous index blocks, with block size and repetition count bounded by k, r ≤4.These limits reduce branching while retaining layer skipping, recurrence, and emergent reordering patterns.
  • B.3. Monte Carlo Tree Search Formulation: Program discovery is formulated as a sequential decision process in which MCTS nodes represent partial or complete programs and actions apply valid skip or repeat operations.Completed programs receive binary rewards based on whether execution produces the ground-truth answer.
  • B.3. Monte Carlo Tree Search Formulation: The MCTS tree policy balances exploitation, exploration, and program-length regularization, using λ to penalize long programs and encourage efficiency.The objective depends on cumulative reward, visit counts, total simulations, and the length penalty.
  • B.4. Search Algorithm: MCTS initializes the standard forward program, repeatedly performs selection, expansion, simulation, and backpropagation, then collects explored programs with nonzero visit counts.The collected programs are analyzed for validity and structural properties after a fixed number of simulations.

C. Experimental Results

The appendix reports additional quantitative results for PoLar across in-distribution and out-of-distribution evaluations on multiple pretrained LLMs. Unless stated otherwise, models remain fully frozen and PoLar predicts execution programs only at inference time.

  • The appendix provides additional experimental results complementing the main paper.
  • Quantitative comparisons cover both in-distribution and out-of-distribution evaluations across multiple pretrained LLMs.
  • Unless otherwise stated, all models are fully frozen, with PoLar predicting execution programs only at inference time.

C.1. In-Distribution Performance

On DART-Math, POLAR delivers strong in-distribution performance across models and difficulty levels, generally outperforming standard inference and prior dynamic-depth baselines. Its joint support for layer skipping and recurrence enables computation scaling through execution-program exploration while avoiding the degradation often seen with skipping-only methods.

  • In-Distribution Performance: POLAR generally outperforms standard inference and prior dynamic-depth baselines across models and DART-Math difficulty levels.The evaluation covers Qwen1.5-MoE-A2.7B-Chat, Qwen2.5-3B-Instruct, and Qwen3-8B across DM-1 to DM-5.
  • In-Distribution Performance: Increasing p@k produces monotonic accuracy improvements for POLAR, demonstrating test-time computation scaling through execution-program exploration.At p@5, POLAR also achieves substantial absolute gains over Base (sampling).
  • In-Distribution Performance: Skipping-only methods often degrade in accuracy on harder levels, whereas POLAR combines layer skipping and recurrence to retain or improve accuracy.POLAR explores diverse execution programs rather than relying solely on layer skipping.
  • In-Distribution Performance: POLAR achieves stronger performance than DR.LLM in most settings, particularly at larger p@k.DR.LLM performs layer-level routing, while POLAR uses structured execution programs.

C.2. Out-of-Distribution Generalization · D. Empirical Details

POLAR generalizes beyond its mathematical training distribution, improving standard inference across arithmetic word-problem benchmarks and most MMLU-Pro subjects. The evaluation covers multiple models and OOD benchmarks, with pass@1 results reported in Tables 8–10.

  • C.2. Out-of-Distribution Generalization: POLAR is evaluated out of distribution on ASDiv, MAWPS, and subject-wise MMLU-Pro subsets using three language models.The evaluated models are LLaMA-3.2-3B-Instruct, Qwen2.5-3B-Instruct, and Qwen3-8B.
  • D. Empirical Details: Tables 8, 9, and 10 report OOD pass@1 accuracy for LLaMA-3.2-3B-Instruct, Qwen2.5-3B-Instruct, and Qwen3-8B, respectively.The tables cover ASDiv, MAWPS, and subject-wise subsets of MMLU-Pro.
  • C.2. Out-of-Distribution Generalization: Across all evaluated models, POLAR consistently improves over standard inference on ASDiv and MAWPS arithmetic word-problem benchmarks.The results indicate transfer from structured mathematical reasoning to natural-language problem settings.
  • C.2. Out-of-Distribution Generalization: On MMLU-Pro, POLAR achieves broad and consistent gains across most subjects spanning mathematics, natural sciences, social sciences, and humanities.MMLU-Pro provides evaluation across diverse domains rather than only mathematical reasoning.
  • C.2. Out-of-Distribution Generalization: The results suggest that POLAR learns general, transferable computation-control strategies rather than dataset-specific heuristics.This interpretation is based on performance across benchmarks differing substantially from DART-Math in format and domain.
  • C.2. Out-of-Distribution Generalization: Despite training on mathematical reasoning data, POLAR generalizes effectively to heterogeneous domains, highlighting the robustness of program-of-layers inference.The reported findings support applicability beyond the original training distribution.
  • D. Empirical Details: Tables 5–7 summarize pass@k accuracy under different inference strategies for Qwen1.5-MoE-A2.7B-Chat, Qwen2.5-3B-Instruct, and Qwen3-8B.These tables provide additional empirical comparisons across inference strategies.

D.1. Dataset Details

The in-distribution evaluation uses deduplicated DART-Math examples across five difficulty levels, with independent difficulty-wise splits and held-out tests matched to each training level.

  • DART-Math comprises five mathematical-reasoning difficulty levels, DM-1 through DM-5, used for all in-distribution experiments.
  • Examples are deduplicated by question text within each difficulty level before independently splitting into training, validation, and test sets.
  • After deduplication, DM-1 to DM-5 contain 565, 1,349, 1,579, 1,537, and 1,577 examples, respectively.Each level uses approximately 62.5% for training, 12.5% for validation, and 25% for testing.
  • The resulting dataset has 4,130 training examples, 826 validation examples, and 1,651 test examples.In-distribution results are reported on held-out test sets corresponding to the same difficulty level used for training.
  • Removing duplicate questions before splitting leaves qualitative conclusions unchanged, while out-of-distribution results remain unaffected.

D.2. Training Configuration · D.3. Predictor Size · D.4. Direct Prompting

The POLAR predictor is trained with validation-selected supervised learning, remains lightweight relative to frozen base LLMs, and is evaluated using direct prompts that isolate latent execution programs from token-level reasoning.

  • D.2. Training Configuration: POLAR prediction networks are trained with supervised learning on the described training splits, with all hyperparameters selected through validation tuning.Optimization uses AdamW, cosine decay with linear warmup, and validation-tuned learning rate, batch size, epochs, and warmup steps.
  • D.2. Training Configuration: The optimization uses AdamW with learning rates from 1e-4 to 3e-3, batch sizes of 32, 128, or 256, and 3 or 10 training epochs.A cosine learning-rate schedule with linear warmup is used, with warmup steps tuned on the validation split.
  • D.3. Predictor Size: Approximately 2.1M parameters make up the learned POLAR predictor across evaluated models.The predictor is trained and used alongside a frozen base LLM.
  • D.3. Predictor Size: The predictor represents only 0.01%–0.06% of the corresponding base model size.This small footprint makes training and inference inexpensive relative to standard LLM fine-tuning or full-model execution.
  • D.3. Predictor Size: Table 11 compares the parameter size of the learned POLAR predictor with each frozen base LLM.The table provides the per-model comparison underlying the reported predictor-size analysis.
  • D.4. Direct Prompting: All experiments use direct prompting without chain-of-thought or intermediate reasoning, instructing models to output only the strictly formatted final answer.The required format is \boxed{ANSWER}.
  • D.4. Direct Prompting: The prompt encloses each math question between problem markers and ends with an answer field, enforcing concise generation and isolating latent execution programs from token-level reasoning strategies.The template includes “### Problem Start {question},” “### Problem End,” and “Answer:”.
Loading 2606.06574v2…