--- title: "Skip a Layer or Loop It? Learning Program-of-Layers in LLMs" authors: - Ziyue Li - Yang Li - Tianyi Zhou affiliations: - University of Maryland, College Park, MD, USA - Mohamed bin Zayed University of Artificial Intelligence (MBZUAI), Abu Dhabi, UAE arxiv_id: "2606.06574" arxiv_version: v2 submitted: 2026-06-04 revised: 2026-08-08 venue: ICML 2026 Oral license: CC BY 4.0 source: https://arxiv.org/abs/2606.06574v2 abstract: | Large language models (LLMs) perform inference by following a fixed depth and order, non-recurrent execution of all layers. The paper identifies training-free, flexible, dynamic programs-of-layers (PoLar), where pretrained layers are treated as modules that can be skipped or looped to form an input-specific execution program. For many inputs, shorter programs achieve the same or better accuracy, and alternative programs can correct predictions made incorrectly by the standard forward pass. The practical method trains a lightweight prediction network from valid execution programs discovered offline by Monte Carlo tree search, then predicts segment boundaries and skip, keep, or repeat operations for a frozen LLM. Experiments on mathematical-reasoning benchmarks report gains over standard inference and prior dynamic-depth methods, including out-of-distribution evaluation. --- # Skip a Layer or Loop It? Learning Program-of-Layers in LLMs - **Authors:** Ziyue Li, Yang Li, Tianyi Zhou - **Affiliations:** University of Maryland, College Park; Mohamed bin Zayed University of Artificial Intelligence (MBZUAI) - **Version:** arXiv:2606.06574v2, revised 2026-08-08 - **Venue:** ICML 2026 Oral - **Project:** https://github.com/tianyi-lab/PoLar # Introduction
Program-of-layers (PoLar) for two different inputs. The D layers in a pretrained LLM define D functions f0, …, fD − 1. Instead of calling them in a static fixed order from f0 to fD − 1, the dynamic inference of PoLar executes an input-specific program π = (i1, …, iK) that calls the functions with layer skipping and recurrence. PoLar enables a training-free architecture of dynamic depth for different inputs, yielding diverse latent computations that cannot be fully covered by existing methods.
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 [@liu2020fastbert; @xin2020deebert; @zhou2020bert; @liu2021faster]. 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 `\ref{fig:search_space}`{=latex}. 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 [@liu2020fastbert; @xin2020deebert; @zhou2020bert; @fan2019reducing; @fan2024looped; @yang2023looped], 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 [@li2025skip], 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. # Dynamic Inference as a Program-of-Layers (PoLar) in Large Language Models {#sec:analysis} 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?**`\looseness-1`{=latex} 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*. ```{=latex} \begin{conjecturebox}[Conjecture: Inference as Program-of-Layers]\textit{ Define the layers in a pretrained LLM as functions, for each input, there can exist multiple distinct executions of programs-of-layers (beyond the standard forward pass) that produce correct predictions. } \end{conjecturebox} ``` **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$ transformer layers, where each layer defines a fixed computation function `\setlength{\abovedisplayskip}{2pt}`{=latex} `\setlength{\belowdisplayskip}{2pt}`{=latex} $$\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 `\setlength{\abovedisplayskip}{2pt}`{=latex} `\setlength{\belowdisplayskip}{2pt}`{=latex} $$\textstyle \pi = (i_1, i_2, \ldots, i_K), \quad i_k \in \{0,\ldots,D-1\},$$ which induces the composed computation `\setlength{\abovedisplayskip}{2pt}`{=latex} `\setlength{\belowdisplayskip}{2pt}`{=latex} $$\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. ```{=latex} \begin{figure*}[t]\begin{center} \includegraphics[width=1.\textwidth]{method_comp.pdf} \end{center} \caption{ \textbf{Sequential MCTS (left) vs. End-to-end \textsc{PoLar} network (right) for prediction of programs.} (a) MCTS in the space of execution programs via sequential iterations of selection, expansion, simulation, and backpropagation. Each node represents a partial or complete execution program, and skip/repeat operations expand the search tree iteratively. This explicit and thorough search is expensive and impractical. (b) Our \textsc{PoLar} trains an end-to-end, lightweight prediction network that directly produces a program representation composed of (i) a binary mask $\mathbf{z}^{seg}(x)$ segmenting layers into modules, and (ii) a vector of operation labels $\mathbf{z}^{op}(x)$ that applies one operation out of \emph{skip}, \emph{keep}, or \emph{repeat} to each module. Our method is scalable in practice and does not require sequential search. } \label{fig:method_comp} \end{figure*} ```
Accuracy of MCTS discovered programs under varying execution-depth budgets across five difficulty levels in DART-Math. We compare the original forward pass (orange) with 90–115% depth-budgeted programs (blue). Shaded regions denote the maximum gain achieved under the highest budget (115%).
**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 `\ref{app:mcts}`{=latex}. All experiments are conducted on **DART-Math** [@tong2024dart], 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. ```{=latex} \begin{findingbox}[title={Finding 1}]\textit{\textbf{Layer recurrence-only performs better than layer skipping-only, but combining the two complementary operators produces the best program-of-layers.}} \end{findingbox} ``` As shown in Table `\ref{tab:overall_accuracy_template_transposed}`{=latex}, 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 `\ref{tab:overall_accuracy_template_transposed}`{=latex}. These results indicate that recurrence provides a stronger mechanism for improving inference than skipping alone, while the two operations play complementary roles when combined.
Latent execution programs often admit shorter valid solutions. We compare the standard forward-pass depth with that of MCTS-discovered valid programs, for initially correct (C→C) and initially incorrect (W→C) inputs. Bars report total execution depth as a fraction of full model depth, with hatched overlays indicating effective depth (the number of unique layers).
(a) Test-time scaling via recurrence over layer segments. Allowing more latent execution steps through segment recurrence leads to a monotonic increase in the probability of discovering valid execution programs across models. (b) Recurrence and skipping are increasingly demanded for harder inputs. The fraction of inputs relying on layer recurrence or skipping to be solved increases with increasing difficulty for most models, except LLaMA-3.2-3B-Instruct, whose deviation is explained by effective difficulty in Table .
Accuracy vs. total layer executions. For each model, we report how the average accuracy of valid execution programs changes with the total number of layer executions (% of base model depth). Across models and difficulty levels, accuracy increases with executed layers, revealing a consistent effect of depth-scaling.
Structural bias of valid execution programs. Valid programs rely primarily on contiguous layer segments as modules (a) and require at most one recurrence of each module (b).
```{=latex} \resizebox{0.43\textwidth}{!}{ \begin{tabular}{lcccc>{\color{blue}}c} \toprule \textbf{Metric} & \textbf{Base} & \textbf{Skip} & \textbf{Loop} & \textbf{Skip\&Loop} & \textbf{Gain} \\ \midrule \rowcolor{gray!15} \textbf{LLaMA-3.2-3B-Instruct} \\ DM-1 & 37.7 & 45.3 & \underline{55.0} & \textbf{85.0} & \textbf{+47.3} \\ DM-2 & 29.1 & 35.9 & \underline{46.6} & \textbf{72.6} & \textbf{+43.5} \\ DM-3 & 23.2 & 29.4 & \underline{38.0} & \textbf{65.6} & \textbf{+42.4} \\ DM-4 & 23.4 & 29.0 & \underline{36.0} & \textbf{58.2} & \textbf{+34.8} \\ DM-5 & 25.8 & 30.2 & \underline{37.5} & \textbf{57.5} & \textbf{+31.7} \\ \midrule \rowcolor{gray!15} \textbf{Qwen1.5-MoE-A2.7B-Chat} \\ DM-1 & 37.2 & 42.7 & \underline{57.3} & \textbf{74.7} & \textbf{+37.5} \\ DM-2 & 27.1 & 33.0 & \underline{44.9} & \textbf{61.5} & \textbf{+34.4} \\ DM-3 & 20.3 & 24.4 & \underline{35.1} & \textbf{52.0} & \textbf{+31.7} \\ DM-4 & 14.6 & 17.2 & \underline{26.0} & \textbf{39.5} & \textbf{+24.9} \\ DM-5 & 10.3 & 12.8 & \underline{19.5} & \textbf{31.3} & \textbf{+21.0} \\ \midrule \rowcolor{gray!15} \textbf{Qwen2.5-3B-Instruct} \\ DM-1 & 26.9 & 47.6 & \underline{60.9} & \textbf{87.8} & \textbf{+60.9} \\ DM-2 & 13.9 & 34.2 & \underline{46.0} & \textbf{76.6} & \textbf{+62.7} \\ DM-3 & 5.3 & 24.6 & \underline{36.9} & \textbf{66.0} & \textbf{+60.7} \\ DM-4 & 2.1 & 16.1 & \underline{24.1} & \textbf{52.8} & \textbf{+50.7} \\ DM-5 & 1.3 & 11.6 & \underline{16.9} & \textbf{41.8} & \textbf{+40.5} \\ \midrule \rowcolor{gray!15} \textbf{Qwen3-8B} \\ DM-1 & 42.3 & 67.1 & \underline{70.3} & \textbf{92.4} & \textbf{+50.1} \\ DM-2 & 29.9 & 52.6 & \underline{58.3} & \textbf{83.5} & \textbf{+53.6} \\ DM-3 & 19.5 & 36.9 & \underline{41.9} & \textbf{68.3} & \textbf{+48.8} \\ DM-4 & 10.9 & 24.9 & \underline{30.2} & \textbf{55.6} & \textbf{+44.7} \\ DM-5 & 8.6 & 18.7 & \underline{22.5} & \textbf{44.2} & \textbf{+35.6} \\ \bottomrule \end{tabular} } ``` ```{=latex} \begin{findingbox}[title={Finding 2}]\textit{\textbf{(Occam's razor) Most valid execution programs are often shorter than the standard forward pass.}} \end{findingbox} ``` As shown in Figure `\ref{fig:length_budget}`{=latex}, 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 `\ref{fig:length}`{=latex} 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. ```{=latex} \begin{findingbox}[title={Finding 3}]\textit{\textbf{Increasing latent execution complexity systematically expands the space of valid programs and improves inference on harder inputs.}} \end{findingbox} ``` 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. `\looseness-1`{=latex} **(a) Test-time scaling expands the space of valid execution programs for latent reasoning.** As shown in Figure `\ref{fig:existence}`{=latex}(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 `\ref{fig:existence}`{=latex}(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 `\ref{tab:overall_accuracy_template_transposed}`{=latex}). **(c) Inference accuracy improves systematically with execution depth.** As shown in Figure `\ref{fig:length_vs_accuracy}`{=latex}, 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. ```{=latex} \begin{findingbox}[title={Finding 4}]\textit{\textbf{Valid execution programs are predominantly composed of contiguous layer segments and typically require at most a single recurrence per segment.}} \end{findingbox} ``` As shown in Figure `\ref{fig:program_structure}`{=latex}, 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 `\ref{fig:program_structure}`{=latex}(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 `\ref{fig:program_structure}`{=latex}(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 `\ref{fig:method_comp}`{=latex} 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.`\looseness-1`{=latex} # Learning Program-of-Layers (PoLar) in Large Language Models {#sec:method} Building on our empirical analysis (Section `\ref{sec:analysis}`{=latex}), we propose PoLar, a method for *programming* pretrained language models at inference time by predicting input-specific execution programs (Figure `\ref{fig:method_comp}`{=latex}). PoLar dynamically segments and composes pretrained layers into reusable modules, enabling flexible computation without parameter updates. ## Program Representation {#sec:prog-repr} 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$, 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$ layers of a pretrained model into contiguous segments `\setlength{\abovedisplayskip}{2pt}`{=latex} `\setlength{\belowdisplayskip}{2pt}`{=latex} $$\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 \le K_{\max}$. Segmentation is represented by a binary boundary mask `\setlength{\abovedisplayskip}{2pt}`{=latex} `\setlength{\belowdisplayskip}{2pt}`{=latex} $$\textstyle \mathbf{z}^{\text{seg}}(x)\in\{0,1\}^D,$$ where $\mathbf{z}^{\text{seg}}_i=1$ indicates that layer index $i$ starts a new segment, and $\mathbf{z}^{\text{seg}}_i=0$ otherwise. We set $K_{\max}=4$ based on empirical evidence. **Finding 4** in Section `\ref{sec:analysis}`{=latex} shows that valid execution programs are dominated by short, contiguous layer segments. Bounding the segment length therefore captures the dominant local execution structures while substantially reducing the complexity of the program space. Although this representation restricts the set of admissible programs, it preserves the most prevalent compositional patterns in practice and enables stable learning with strong empirical performance. **Operations.** For each segment $[s_j,s_{j+1})$, the execution program assigns one of three operations $\{\textsf{skip},\textsf{keep},\textsf{repeat}\}$, which determines how the segment is executed: `\setlength{\abovedisplayskip}{0pt}`{=latex} `\setlength{\belowdisplayskip}{0pt}`{=latex} $$\textstyle \begin{aligned} \textsf{skip} &: \emptyset,\\ \textsf{keep} &: [s_j,\ldots,s_{j+1}-1],\\ \textsf{repeat} &: [s_j,\ldots,s_{j+1}-1,\ s_j,\ldots,s_{j+1}-1]. \end{aligned}$$ The [skip]{.sans-serif} operator omits a segment to reduce computation, while [repeat]{.sans-serif} applies a single additional pass. Operations are represented by a categorical label vector `\setlength{\abovedisplayskip}{2pt}`{=latex} `\setlength{\belowdisplayskip}{2pt}`{=latex} $$\textstyle \mathbf{z}^{\text{op}}(x)\in\{\textsf{skip},\textsf{keep},\textsf{repeat}\}^D,$$ where $\mathbf{z}^{\text{op}}_i$ is defined only when $\mathbf{z}^{\text{seg}}_i=1$ (i.e., at segment start positions); labels at all other positions are ignored. `\looseness-1`{=latex} This operator set is intentionally minimal and empirically grounded. **Finding 4** in Section `\ref{sec:analysis}`{=latex} shows valid execution programs rarely require more than a single re-execution within a segment, and that [skip]{.sans-serif} and [repeat]{.sans-serif} account for the most effective execution patterns, offering strong performance--efficiency trade-offs. The [keep]{.sans-serif} operator preserves the original computation when no modification is needed. Although our implementation allows at most one additional execution through [repeat]{.sans-serif}, the representation is not fundamentally limited to a single recurrence. The operation vocabulary can be extended to $\{\textsf{repeat}\text{-}2,\ldots,\textsf{repeat}\text{-}k\}$ to support multiple recurrences per segment. We use a single-repeat operator because the MCTS traces in Section `\ref{sec:analysis}`{=latex} show that effective programs rarely benefit from deeper repeated execution of the same segment. This choice keeps the prediction space tractable while covering the dominant valid programs. ## Program-of-Layers (PoLar) Prediction Network {#sec:policy} We train a lightweight predictor to output logits for the program representation defined in Section `\ref{sec:prog-repr}`{=latex}. **Architecture.** Given an input $x$, we first encode it using a frozen embedding model (`Qwen3-Embedding-0.6B`), as token-level representations $\textstyle \mathbf{H}=E(x)\in\mathbb{R}^{T\times d_q},$ where $T$ is the token length and $d_q$ is the hidden size of the embedding model. We project token representations to a working dimension $d$: $\textstyle \tilde{\mathbf{H}}=\mathbf{H}\mathbf{W}_h \in \mathbb{R}^{T\times d}.$ *Layer queries.* We associate each pretrained transformer layer index $i\in\{0,\ldots,D-1\}$ with a learnable embedding $\mathbf{e}_i\in\mathbb{R}^{d}$, and stack them as $\mathbf{E}\in\mathbb{R}^{D\times d}$. These embeddings act as layer-specific queries. *Cross-attention.* We apply multi-head cross-attention with layer embeddings as queries and token embeddings as keys/values: $\textstyle \mathbf{X}=\textsc{MHA}(\mathbf{Q},\mathbf{K},\mathbf{V}), \mathbf{Q}=\mathbf{E},\ \mathbf{K}=\tilde{\mathbf{H}},\ \mathbf{V}=\tilde{\mathbf{H}},$ where padding tokens are masked using the input attention mask. The output $\mathbf{X}\in\mathbb{R}^{D\times d}$ provides an input-conditioned representation for each layer index. *Cross-layer encoder.* To model dependencies across model depth, we apply a lightweight transformer encoder over the layer dimension: `\setlength{\abovedisplayskip}{2pt}`{=latex} `\setlength{\belowdisplayskip}{2pt}`{=latex} $$\textstyle \mathbf{X}'=\textsc{Enc}_{\text{layer}}(\mathbf{X})\in\mathbb{R}^{D\times d}.$$ This enables self-attention across layers, allowing decisions at each layer to depend on global depth context. *Prediction heads.* Two linear heads produce logits for segmentation boundaries and operations: `\setlength{\abovedisplayskip}{0pt}`{=latex} `\setlength{\belowdisplayskip}{0pt}`{=latex} $$\boldsymbol{\ell}^{\text{seg}}=\mathbf{X}'\mathbf{W}_{\text{seg}}+\mathbf{b}_{\text{seg}} \in\mathbb{R}^{D}, \boldsymbol{\ell}^{\text{op}}=\mathbf{X}'\mathbf{W}_{\text{op}}+\mathbf{b}_{\text{op}} \in\mathbb{R}^{D\times 3}.$$ **Supervision from Valid Execution Programs.** We supervise training using valid execution programs collected offline via MCTS (Section `\ref{sec:analysis}`{=latex}). Each program is deterministically parsed into program representation, producing ground-truth segmentation and operation labels $\mathbf{z}^{\text{seg}}(x)$ and $\mathbf{z}^{\text{op}}(x)$ in the format defined in Section `\ref{sec:prog-repr}`{=latex}. When multiple valid programs are available for an input and at least one is shorter than the full model depth, we down-weight the loss of the full-depth execution. This choice follows **Finding 2**, which shows that shorter valid programs are preferred while still preserving supervision from the original computation. **Training Objective.** We train the predictor to match the ground-truth execution program, specified by segmentation and operation labels $\big(\mathbf{z}^{\text{seg}*}(x),\mathbf{z}^{\text{op}*}(x)\big)$. Let $p^{\text{seg}}_i=\sigma(\ell^{\text{seg}}_i)$ and $\mathbf{p}^{\text{op}}_i=\textsc{Softmax}(\boldsymbol{\ell}^{\text{op}}_i)$. Segmentation is supervised with binary cross-entropy over boundary indicators: `\setlength{\abovedisplayskip}{2pt}`{=latex} `\setlength{\belowdisplayskip}{2pt}`{=latex} $$\textstyle \mathcal{L}_{\text{seg}} = -\sum_{i=0}^{D-1}\Big[ \mathbf{z}^{\text{seg}*}_i\log p^{\text{seg}}_i +(1-\mathbf{z}^{\text{seg}*}_i)\log(1-p^{\text{seg}}_i) \Big].$$ Operation prediction uses a masked cross-entropy applied only at segment start positions. With mask $m_i=\mathbf{z}^{\text{seg}*}_i$, we compute `\setlength{\abovedisplayskip}{0pt}`{=latex} `\setlength{\belowdisplayskip}{0pt}`{=latex} $$\textstyle \mathcal{L}_{\text{op}} = -\sum_{i=0}^{D-1} m_i\cdot \log \mathbf{p}^{\text{op}}_i\big[\mathbf{z}^{\text{op}*}_i\big].$$ The final objective is $\mathcal{L}=\mathcal{L}_{\text{seg}}+\mathcal{L}_{\text{op}}.$ **Inference-Time Program Decoding.** At inference time, execution programs are decoded in two stages. First, segment boundaries are determined deterministically by thresholding the predicted segmentation logits $\boldsymbol{\ell}^{\text{seg}}$. If any resulting segment exceeds the maximum length constraint $K_{\max}$, additional boundaries are inserted to enforce it, yielding segment start positions $\{s_j\}$. Conditioned on this segmentation, we compute operation log-probabilities at each segment start from the predicted logits: `\setlength{\abovedisplayskip}{2pt}`{=latex} `\setlength{\belowdisplayskip}{2pt}`{=latex} $$\textstyle \log p(o_j \mid x, s_j) = \log \textsc{Softmax}\!\big(\boldsymbol{\ell}^{\text{op}}_{s_j}\big)[o_j].$$ Rather than selecting operations independently via local argmax, we apply a small beam search over segment-level operation choices to account for non-local interactions between segments and to ensure globally consistent execution programs. This search operates over a highly constrained space and produces a ranked set of candidate execution programs $\pi(x)$. Finally, each candidate program is mapped deterministically to a concrete executed program using the segment-to-path rules in Section `\ref{sec:prog-repr}`{=latex}. ```{=latex} \resizebox{0.5\textwidth}{!}{ \begin{tabular}{llccccc} \toprule \textbf{Method} & \textbf{p@} & \textbf{DM-1} & \textbf{DM-2} & \textbf{DM-3} & \textbf{DM-4} & \textbf{DM-5} \\ \midrule \rowcolor{gray!15}\textbf{LLaMA-3.2-3B-Instruct}\\ Base ($\tau$=0) & 1 & 51.1 & 30.6 & 29.9 & 27.1 & 26.1 \\ \midrule Base (sampling) & 1 & 48.9 & 29.4 & 30.4 & 28.6 & 26.6 \\ & 2 & 53.9 & 33.2 & 34.2 & 32.0 & 31.2 \\ & 3 & 59.6 & 35.9 & 36.5 & 34.6 & 33.2 \\ & 4 & 61.7 & 38.9 & 38.7 & 36.5 & 33.5 \\ & 5 & 65.2 & 41.2 & 39.5 & 38.5 & 34.5 \\ \midrule ShortGPT & 1 & 16.3 & 7.1 & 5.3 & 6.5 & 5.3 \\ & 2 & 21.3 & 13.4 & 11.1 & 10.9 & 9.9 \\ & 3 & 24.1 & 14.2 & 11.6 & 12.0 & 11.7 \\ & 4 & 25.5 & 15.1 & 12.9 & 12.2 & 12.2 \\ & 5 & 27.7 & 18.7 & 16.2 & 13.3 & 14.0 \\ \midrule MindSkip & 1 & 14.2 & 6.5 & 4.8 & 6.5 & 4.8 \\ & 2 & 24.1 & 12.8 & 9.9 & 12.5 & 7.1 \\ & 3 & 26.2 & 12.8 & 10.9 & 12.5 & 8.4 \\ & 4 & 29.1 & 13.6 & 12.7 & 12.8 & 8.9 \\ & 5 & 30.5 & 16.9 & 16.5 & 14.6 & 10.4 \\ \midrule FlexiDepth & 1 & 2.1 & 0.0 & 0.3 & 0.5 & 0.5 \\ & 2 & 16.3 & 7.1 & 5.6 & 6.5 & 5.6 \\ & 3 & 19.9 & 12.2 & 8.9 & 8.1 & 8.9 \\ & 4 & 20.6 & 12.2 & 8.9 & 8.1 & 8.9 \\ & 5 & 24.8 & 14.5 & 10.9 & 8.6 & 10.4 \\ \midrule DR.LLM & 1 & 51.1 & 23.1 & 26.8 & 18.8 & 18.0 \\ & 2 & 57.4 & 27.6 & 30.9 & 19.0 & 19.3 \\ & 3 & 64.5 & 31.5 & 37.2 & 22.1 & 24.6 \\ & 4 & 65.2 & 36.5 & 42.5 & 28.6 & 29.4 \\ & 5 & 68.8 & 39.8 & 45.3 & 30.2 & 31.7 \\ \midrule \textbf{\textsc{PoLar}} & 1 & 54.6 & 32.3 & 32.7 & 29.7 & 26.9 \\ & 2 & 65.2 & 38.0 & 37.0 & 37.0 & 32.2 \\ & 3 & 67.4 & 41.2 & 41.8 & 39.6 & 36.3 \\ & 4 & 69.5 & 44.5 & 45.3 & 42.7 & 39.1 \\ \rowcolor{blue!10} & 5 & \textbf{74.5} & \textbf{46.9} & \textbf{47.6} & \textbf{45.1} & \textbf{41.9} \\ \midrule $\Delta$ vs.\ Base (sampling) & 5 & \textbf{+9.3} & \textbf{+5.7} & \textbf{+8.1} & \textbf{+6.6} & \textbf{+7.4} \\ \bottomrule \end{tabular}} ``` ```{=latex} \begin{table*}[t] \caption{ \textbf{Out-of-distribution (OOD) performance at pass@1.} We report accuracy using Qwen1.5-MoE-A2.7B-Chat. } \label{tab:ood_pass1_qwen15} \resizebox{1.0\textwidth}{!}{ \begin{tabular}{lcc|ccccccccccccc} \toprule & & & \textbf{MMLU-Pro} \\ \cmidrule(lr){4-16} \textbf{Method} & \textbf{ASDiv} & \textbf{MAWPS} & \textbf{Math} & \textbf{Phys} & \textbf{Chem} & \textbf{Law} & \textbf{Eng} & \textbf{Other} & \textbf{Econ} & \textbf{Health} & \textbf{Psych} & \textbf{Bus} & \textbf{Bio} & \textbf{Phil} & \textbf{Hist} \\ \midrule Base ($\tau$=0) & 59.1 & 41.7 & 13.9 & 15.6 & 13.8 & 16.6 & 15.1 & 22.8 & 31.0 & 26.8 & 30.7 & 18.4 & 34.7 & 22.8 & 22.6 \\ ShortGPT & 2.3 & 0.6 & 3.8 & 0.4 & 4.2 & 3.3 & 3.8 & 2.6 & 2.4 & 2.2 & 2.8 & 4.9 & 2.9 & 1.8 & 4.7 \\ MindSkip & 0.0 & 0.0 & 0.9 & 1.2 & 1.0 & 1.9 & 1.3 & 1.3 & 0.6 & 0.7 & 1.3 & 0.8 & 1.7 & 0.6 & 1.8 \\ FlexiDepth & 0.0 & 0.0 & 2.1 & 3.5 & 2.5 & 4.4 & 3.3 & 2.5 & 1.7 & 3.5 & 2.0 & 2.9 & 3.8 & 3.0 & 3.7 \\ DR.LLM & 59.1 & 41.3 & 14.6 & 17.4 & 13.2 & 19.8 & 16.0 & 20.8 & 31.8 & 27.0 & 32.2 & 17.5 & 33.5 & 22.8 & 21.3 \\ \midrule \textsc{PoLar} & \textbf{63.8} & \textbf{46.7} & \textbf{18.5} & \textbf{20.3} & \textbf{18.3} & \textbf{20.4} & \textbf{19.9} & \textbf{26.6} & \textbf{34.6} & \textbf{29.5} & \textbf{35.3} & \textbf{20.9} & \textbf{36.9} & \textbf{25.9} & \textbf{23.5} \\ \bottomrule \end{tabular} } \end{table*} ``` # Experiments We evaluate PoLar across both in-distribution and out-of-distribution benchmarks to assess whether learning *latent execution programs* provides a practical and transferable alternative to search-based test-time computation. ## Experimental Setup **Models.** We evaluate PoLar on a diverse set of pretrained, instruction-tuned LLMs spanning different architectures and scales: *LLaMA-3.2-3B-Instruct*, *Qwen1.5-MoE-A2.7B-Chat*, *Qwen2.5-3B-Instruct*, and *Qwen3-8B*. All models are used in a fully frozen setting with no parameter updates. **Datasets.** We use **DART-Math** [@tong2024dart], a structured mathematical reasoning dataset with five difficulty levels (DM-1 to DM-5), as in-distribution benchmark. For out-of-distribution (OOD) evaluation, we use **ASDiv** [@miao2020diverse] and **MAWPS** [@kadlvcik2023calc], which focus on arithmetic word problems, as well as subject subsets from **MMLU-Pro** [@wang2024mmlu] spanning mathematics, natural sciences, social sciences, and humanities. These benchmarks differ substantially from DART-Math in both format and domain coverage. **In-distribution evaluation.** For DART-Math, we first deduplicate examples within each difficulty level by question text and then adopt a difficulty-wise train/test split: each difficulty level is split independently, and models are trained and evaluated within the same difficulty distribution. **Out-of-distribution (OOD) evaluation.** For OOD evaluation, PoLar is trained on the union of DART-Math training data across all difficulty levels and evaluated zero-shot. This setting directly tests whether PoLar learns transferable computation control strategies rather than heuristics specific to a dataset or difficulty level. **Metric.** We report **pass@$k$ accuracy**, defined as the probability that at least one of the top-$k$ candidates produces a correct answer. For PoLar, the $k$ candidates correspond to the top-$k$ predicted execution programs selected via beam search. For sampling-based baselines, $k$ corresponds to the number of stochastic decoding samples. Unless otherwise stated, OOD results are reported using pass@1. **Baselines.** We compare PoLar against standard inference and representative dynamic-computation methods. **Base ($\tau=0$)** uses greedy decoding with temperature $\tau=0$. **Base (sampling)** samples $k$ outputs using stochastic decoding with $\tau \in \{0.3, 0.7, 1.0\}$ and reports the best result across temperatures, increasing output diversity without altering internal execution. **DR.LLM** [@heakl2025dr] learns layer-routing policies from execution paths and applies them at inference time. **ShortGPT** [@men2025shortgpt] statically prunes layers based on estimated importance, yielding a reduced-depth model. **MindSkip** [@he2024router] and **FlexiDepth** [@luo2025adaptive] learn router-based dynamic-depth policies, primarily optimized for inference efficiency. Several approaches, such as Mixture-of-Depths [@raposo2024mixture], LaCo [@yang2024laco], and Mixture-of-Recursions [@bae2025mixture], require substantial additional training or architectural modification. In contrast, PoLar performs lightweight test-time program selection without modifying pretrained model parameters. More dataset and training details are in Appendix `\ref{app:emp_detail}`{=latex}. ## Main Results We evaluate in-distribution performance on DART-Math. Table `\ref{tab:passatk_llama32_3b}`{=latex} reports pass@$k$ results using *LLaMA-3.2-3B-Instruct*, with complete results provided in Appendix `\ref{app:id_pref}`{=latex}.
Pass@k accuracy and unique depth on LLaMA-3.2-3B-Instruct, averaged across DART-Math DM-1 to DM-5. (a) reports macro-averaged pass@k accuracy for Base (sampling) and PoLar. (b) reports the CDF of the shortest successful execution program among the top-5 candidates, measured by unique executed layers as a percentage of the original model depth.
**Accuracy gains arise from improved latent execution within the frozen model.** At pass@1, PoLar outperforms Base (sampling) across all difficulty levels. For example, on DM-1, accuracy improves from 48.9% to 54.6%, yielding an absolute pass@1 gain of +5.7 percentage points. Since pass@1 evaluates a single decoded output, this gain reflects more effective latent execution selection rather than output-space diversity. **Exploring the execution-program space enables effective test-time scaling.** Increasing the number of candidate execution programs ($k$) monotonically improves PoLar, evidencing effective test-time computation scaling through execution-program exploration. Figure `\ref{fig:passk_depth}`{=latex}(a) illustrates this behavior on LLaMA-3.2-3B-Instruct, macro-averaged across DART-Math DM-1 to DM-5: Base (sampling) improves from 32.8% at pass@1 to 43.8% at pass@5, while PoLar improves from 35.1% to 51.2%. Thus, at pass@5, PoLar achieves an average absolute gain of +7.4 percentage points over Base (sampling). Crucially, Figure `\ref{fig:passk_depth}`{=latex}(b) shows that among successful top-5 candidates, PoLar often finds execution programs that use fewer unique layers than a standard forward pass, indicating that the gains arise from better latent execution programs rather than simply executing the full depth. In contrast, Base (sampling) explores output-space diversity under a fixed computation graph and therefore always uses the original full-depth execution. These results show that structured execution-program exploration can outperform output sampling under the same frozen model. **Program-level execution exploration is more effective than local routing decisions.** Existing dynamic-depth methods primarily make local, layer-wise routing decisions, which restrict inference to a limited execution space and often degrade accuracy in our setting. DR.LLM supports both layer skipping and repetition but operates at the individual-layer level, limiting global coordination across depth. In contrast, PoLar formulates inference as *program-level* exploration over execution programs defined on packed contiguous segments, enabling coordinated skip and repeat patterns across depth. This design directly reflects the execution structures uncovered by MCTS, while replacing expensive search with a lightweight, learned predictor.`\looseness-1`{=latex} **PoLar incurs negligible inference overhead and reduces end-to-end latency.** Beyond counting executed layers, we measure wall-clock latency on *Qwen1.5-MoE-A2.7B-Chat* with 24 layers. As shown in Table `\ref{tab:latency_overhead}`{=latex}, the encoder, predictor head, and beam search introduce a total additional overhead of only 3.05 ms, corresponding to 0.8% of a standard forward pass and approximately 0.23 LLM layers. This overhead is small compared with the latency reduction achieved by executing fewer layers. Consequently, PoLar reduces end-to-end latency while improving accuracy: it achieves 0.83$\times$ the base runtime on easier inputs and 0.95$\times$ on harder inputs. The learned predictor is also lightweight in parameter count: across all evaluated backbones, it contains approximately 2.1M parameters, corresponding to only 0.01%--0.06% of the base LLM size. Full parameter counts are provided in Appendix `\ref{app:predictor_size}`{=latex}. ```{=latex} \resizebox{0.48\textwidth}{!}{ \begin{tabular}{lccc} \toprule \textbf{Component-wise overhead} \\ \midrule \textbf{Component} & \textbf{Latency (ms)} & \textbf{Equiv. layers} & \textbf{\% full forward} \\ \midrule One LLM layer & 13.23 & 1.00 & 3.5\% \\ \midrule Predictor head & 0.99 & 0.07 & 0.3\% \\ Beam search & 0.11 & 0.01 & 0.03\% \\ Encoder & 1.95 & 0.15 & 0.5\% \\ \textbf{Total additional overhead} & \textbf{3.05} & \textbf{0.23} & \textbf{0.8\%} \\ \midrule \textbf{End-to-end latency} \\ \midrule \textbf{Method} & \textbf{Avg. layers} & \textbf{Latency (ms)} & \textbf{Rel. / Acc. gain} \\ \midrule Base & 24.00 & 373.45 & 1.00$\times$ / -- \\ \textsc{PoLar} (DM-1) & 23.30 & 311.41 & 0.83$\times$ / +5.7 \\ \textsc{PoLar} (DM-5) & 23.76 & 353.31 & 0.95$\times$ / +9.4 \\ \bottomrule \end{tabular} } ``` ## Out-of-Distribution Performance We evaluate the OOD generalization of execution programs learned from in-distribution data. As shown in Table `\ref{tab:ood_pass1_qwen15}`{=latex}, PoLar consistently outperforms standard inference on all OOD benchmarks using *Qwen1.5-MoE-A2.7B-Chat*, with full results reported in Appendix `\ref{app:ood_pref}`{=latex}. **Execution programs learned on mathematical datasets transfer across domains.** On arithmetic word problem benchmarks such as ASDiv and MAWPS, PoLar achieves clear improvements over the standard forward pass. More notably, on MMLU-Pro, PoLar improves accuracy across diverse subject areas. We conjecture that this cross-domain transfer comes from two complementary sources. First, the external input representation maps examples from different domains into a shared semantic space, allowing the small PoLar prediction head trained on mathematics to generalize beyond its training distribution. Second, the predicted programs are constrained to simple structural patterns, namely contiguous segments with limited recurrence, which encourages reusable computation strategies rather than benchmark-specific execution heuristics. # Related Works Transformers process inputs through sequential layer stacks, making layer-level computation reduction a critical research direction. Early-exit and layer skipping methods [@liu2020fastbert; @xin2020deebert; @zhou2020bert; @liu2021faster] dynamically terminate computation at intermediate layers using auxiliary classifiers and confidence metrics, allowing easy inputs to exit early. LayerSkip [@elhoushi2024layerskip] shares classifiers across layers to reduce overhead. LayerDrop [@fan2019reducing] trains models so arbitrary layer subsets can be skipped during inference. ShortGPT [@men2025shortgpt] assesses layer importance based on input-output similarity and drops low-importance layers. LaCo [@yang2024laco] merges layers using weight arithmetic. Recent work introduces learned routing for adaptive skipping: FlexiDepth [@luo2025adaptive] and MindSkip [@he2024router] 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 [@dehghani2018universal] apply self-attention blocks recurrently with halting mechanisms to adapt depth per token. Recent looped transformers [@fan2024looped; @yang2023looped] repeatedly apply single blocks to achieve better length generalization on algorithmic tasks by adjusting loop counts during inference. The Inner Thinking Transformer [@chen2025inner] 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. @li2025skip 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 [@heakl2025dr] 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{\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. # 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 {#impact-statement .unnumbered} 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. # References
Andreas, Jacob, Marcus Rohrbach, Trevor Darrell, and Dan Klein. 2016. “Neural Module Networks.” In *Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition*, 39–48.
Bae, Sangmin, Yujin Kim, Reza Bayat, Sungnyun Kim, Jiyoun Ha, Tal Schuster, Adam Fisch, et al. 2025. “Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation.” *arXiv Preprint arXiv:2507.10524*.
Chen, Yilong, Junyuan Shang, Zhenyu Zhang, Yanxi Xie, Jiawei Sheng, Tingwen Liu, Shuohuan Wang, Yu Sun, Hua Wu, and Haifeng Wang. 2025. “Inner Thinking Transformer: Leveraging Dynamic Depth Scaling to Foster Adaptive Internal Thinking.” *arXiv Preprint arXiv:2502.13842*.
Dehghani, Mostafa, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Łukasz Kaiser. 2018. “Universal Transformers.” *arXiv Preprint arXiv:1807.03819*.
Elhoushi, Mostafa, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, et al. 2024. “Layerskip: Enabling Early Exit Inference and Self-Speculative Decoding.” In *Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)*, 12622–42.
Eyzaguirre, Cristóbal, Felipe Del Rio, Vladimir Araujo, and Alvaro Soto. 2022. “DACT-BERT: Differentiable Adaptive Computation Time for an Efficient BERT Inference.” In *Proceedings of NLP Power! The First Workshop on Efficient Benchmarking in NLP*, 93–99.
Fan, Angela, Edouard Grave, and Armand Joulin. 2019. “Reducing Transformer Depth on Demand with Structured Dropout.” *arXiv Preprint arXiv:1909.11556*.
Fan, Ying, Yilun Du, Kannan Ramchandran, and Kangwook Lee. 2024. “Looped Transformers for Length Generalization.” *arXiv Preprint arXiv:2409.15647*.
He, Shwai, Tao Ge, Guoheng Sun, Bowei Tian, Xiaoyang Wang, and Dong Yu. 2024. “Router-Tuning: A Simple and Effective Approach for Enabling Dynamic-Depth in Transformers.” *arXiv Preprint arXiv:2410.13184*.
Heakl, Ahmed, Martin Gubri, Salman Khan, Sangdoo Yun, and Seong Joon Oh. 2025. “Dr. LLM: Dynamic Layer Routing in LLMs.” *arXiv Preprint arXiv:2510.12773*.
Kadlčı́k, Marek, Michal Štefánik, Ondrej Sotolár, and Vlastimil Martinek. 2023. “Calc-x and Calcformers: Empowering Arithmetical Chain-of-Thought Through Interaction with Symbolic Systems.” In *Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing*, 12101–8.
Li, Ziyue, Yang Li, and Tianyi Zhou. 2025. “Skip a Layer or Loop It? Test-Time Depth Adaptation of Pretrained Llms.” *arXiv Preprint arXiv:2507.07996*.
Liu, Weijie, Peng Zhou, Zhiruo Wang, Zhe Zhao, Haotang Deng, and Qi Ju. 2020. “Fastbert: A Self-Distilling Bert with Adaptive Inference Time.” In *Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics*, 6035–44.
Liu, Yijin, Fandong Meng, Jie Zhou, Yufeng Chen, and Jinan Xu. 2021. “Faster Depth-Adaptive Transformers.” In *Proceedings of the AAAI Conference on Artificial Intelligence*, 35:13424–32. 15.
Luo, Xuan, Weizhi Wang, and Xifeng Yan. 2025. “Adaptive Layer-Skipping in Pre-Trained Llms.” *arXiv Preprint arXiv:2503.23798*.
Men, Xin, Mingyu Xu, Qingyu Zhang, Qianhao Yuan, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, and Weipeng Chen. 2025. “Shortgpt: Layers in Large Language Models Are More Redundant Than You Expect.” In *Findings of the Association for Computational Linguistics: ACL 2025*, 20192–204.
Miao, Shen-Yun, Chao-Chun Liang, and Keh-Yih Su. 2020. “A Diverse Corpus for Evaluating and Developing English Math Word Problem Solvers.” In *Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics*, 975–84.
Raposo, David, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro. 2024. “Mixture-of-Depths: Dynamically Allocating Compute in Transformer-Based Language Models.” *arXiv Preprint arXiv:2404.02258*.
Tong, Yuxuan, Xiwen Zhang, Rui Wang, Ruidong Wu, and Junxian He. 2024. “Dart-Math: Difficulty-Aware Rejection Tuning for Mathematical Problem-Solving.” *Advances in Neural Information Processing Systems* 37: 7821–46.
Wang, Yubo, Xueguang Ma, Ge Zhang, Yuansheng Ni, Abhranil Chandra, Shiguang Guo, Weiming Ren, et al. 2024. “Mmlu-Pro: A More Robust and Challenging Multi-Task Language Understanding Benchmark.” *Advances in Neural Information Processing Systems* 37: 95266–90.
Xin, Ji, Raphael Tang, Jaejun Lee, Yaoliang Yu, and Jimmy Lin. 2020. “DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference.” *arXiv Preprint arXiv:2004.12993*.
Yang, Liu, Kangwook Lee, Robert Nowak, and Dimitris Papailiopoulos. 2023. “Looped Transformers Are Better at Learning Learning Algorithms.” *arXiv Preprint arXiv:2311.12424*.
Yang, Yifei, Zouying Cao, and Hai Zhao. 2024. “Laco: Large Language Model Pruning via Layer Collapse.” *arXiv Preprint arXiv:2402.11187*.
Zhou, Wangchunshu, Canwen Xu, Tao Ge, Julian McAuley, Ke Xu, and Furu Wei. 2020. “Bert Loses Patience: Fast and Robust Inference with Early Exit.” *Advances in Neural Information Processing Systems* 33: 18330–41.
```{=latex} \newpage ``` ```{=latex} \clearpage ``` ```{=latex} \appendix ``` ```{=latex} \onecolumn ``` # 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, @liu2021ebert demonstrate that a significant fraction of BERT's attention heads can be dropped with negligible performance loss, and @gordon2020compressing investigate fine-grained weight pruning in BERT. [@fan2019reducing] 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.`\looseness-1`{=latex} By contrast, early-exit or input-adaptive methods add auxiliary classifiers at intermediate layers so that \"easy\" inputs exit early. Notable examples include FastBERT [@liu2020fastbert] and DeeBERT [@xin2020deebert], which insert classifiers after each block and use confidence or entropy metrics to decide when to stop. PABEE [@zhou2020bert] employs a patience criterion to halt when predictions stabilize. DACT-BERT [@eyzaguirre2022dact] adopts a differentiable Adaptive Computation Time mechanism to learn how many Transformer layers to run for each example. @liu2021faster 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. @xu2023lgvit propose LGViT, which adds heterogeneous exit heads (local and global) to ViT so that vision transformers can terminate early with minimal feature loss. @tang2023you 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.`\looseness-1`{=latex} #### Looped Transformer and Recurrent Depth Another line of research makes Transformer depth adaptive by looping or repeating layers. The Universal Transformer [@dehghani2018universal] 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. @fan2024looped 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, @yang2023looped 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 @chen2025inner 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. @wu2024routing 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. @jain2024mixture 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 [@andreas2016neural] 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. # Searching the Execution Program Space {#app:mcts} This appendix provides full details of the execution-program search procedure used to test the conjecture in Section `\ref{sec:analysis}`{=latex}. 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. ## 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$ transformer layers, where each layer defines a fixed computation function $$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 $$\pi = (i_1, i_2, \ldots, i_K), \quad i_k \in \{0,\ldots,D-1\},$$ which induces the composed computation $$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. ## 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$ indices from the program; - **Repeat**: duplicate a contiguous block of $k$ indices for $r$ repetitions. In all experiments, block size $k$ and repetition count $r$ are bounded by small constants ($k,r \leq 4$). These constraints significantly reduce the branching factor while retaining the ability to realize layer skipping, recurrence, and emergent reordering patterns. ## 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$ with ground-truth answer $y$, we define a binary reward $$r(\pi, x) = \mathbf{1}\{F_{\pi}(x) = y\},$$ where $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: $$\mathrm{UCB}(\pi) = \frac{R(\pi)}{v(\pi)} + c \sqrt{\frac{\ln V}{v(\pi)}} - \lambda \frac{|\pi|}{D},$$ where $R(\pi)$ is the cumulative reward, $v(\pi)$ is the visit count, $V$ is the total number of simulations, and $\lambda$ penalizes long programs to encourage efficiency. ## Search Algorithm Algorithm `\ref{alg:mcts}`{=latex} summarizes the MCTS procedure. We initialize the root node with the standard forward execution $\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. ```{=latex} \begin{algorithm}[t]\caption{Monte Carlo Tree Search for Execution Programs} \label{alg:mcts} \begin{algorithmic}[1] \REQUIRE Input $x$, number of simulations $N_{\text{sim}}$ \STATE Initialize root program $\pi_0 = (0,1,\ldots,D-1)$ \FOR{$n = 1$ to $N_{\text{sim}}$} \STATE \textbf{Selection:} traverse tree using UCB to reach a leaf program \STATE \textbf{Expansion:} generate child programs via valid skip/repeat actions \STATE \textbf{Simulation:} execute $F_{\pi}(x)$ and compute reward $r(\pi, x)$ \STATE \textbf{Backpropagation:} update visit counts and cumulative rewards \ENDFOR \STATE \textbf{return} explored programs and associated statistics \end{algorithmic} \end{algorithm} ``` ```{=latex} \clearpage ``` # Experimental Results {#app:experiment_result} 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. ## In-Distribution Performance {#app:id_pref} 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 `\ref{tab:passatk_qwen}`{=latex}, `\ref{tab:passatk_qwen25}`{=latex}, and `\ref{tab:passatk_qwen3}`{=latex} present pass@k accuracy under different inference strategies for Qwen1.5-MoE-A2.7B-Chat, Qwen2.5-3B-Instruct, and Qwen3-8B, respectively. Across all models and difficulty levels, PoLar achieves strong in-distribution performance and generally outperforms standard inference and prior dynamic-depth baselines. In particular, increasing $p@k$ leads to monotonic accuracy improvements for PoLar, demonstrating effective test-time computation scaling through execution-program exploration. At $p@5$, PoLar achieves substantial absolute gains over Base (sampling), with improvements of +15.6 / +17.5 / +15.0 / +10.1 / +12.2 on Qwen1.5-MoE-A2.7B-Chat, +40.4 / +19.2 / +17.5 / +14.8 / +13.5 on Qwen2.5-3B-Instruct, and +14.1 / +23.4 / +15.5 / +10.6 / +9.9 on Qwen3-8B for DM-1 to DM-5, respectively. Notably, methods that rely solely on layer skipping (e.g., ShortGPT, MindSkip, FlexiDepth) often suffer accuracy degradation, especially on harder difficulty levels. In contrast, PoLar jointly supports layer skipping and recurrence, allowing it to retain or improve accuracy while exploring diverse execution programs. Compared to DR.LLM, which performs layer-level routing, PoLar achieves stronger performance in most settings, particularly at larger $p@k$, indicating the benefit of structured, program-level execution prediction. ```{=latex} \resizebox{0.6\textwidth}{!}{ \begin{tabular}{llccccc} \toprule \textbf{Method} & \textbf{p@} & \textbf{DM-1} & \textbf{DM-2} & \textbf{DM-3} & \textbf{DM-4} & \textbf{DM-5} \\ \midrule \rowcolor{gray!15}\textbf{Qwen1.5-MoE-A2.7B-Chat}\\ Base ($\tau$=0) & 1 & 48.2 & 25.5 & 24.1 & 16.7 & 10.9 \\ \midrule Base (sampling) & 1 & 41.8 & 24.9 & 21.3 & 18.2 & 12.4 \\ & 2 & 46.8 & 28.2 & 23.3 & 20.3 & 14.2 \\ & 3 & 49.6 & 30.0 & 25.8 & 21.6 & 14.7 \\ & 4 & 50.4 & 31.8 & 26.8 & 22.7 & 15.5 \\ & 5 & 50.4 & 33.5 & 27.3 & 22.7 & 16.5 \\ \midrule ShortGPT & 1 & 22.7 & 12.8 & 8.6 & 8.3 & 3.0 \\ & 2 & 45.4 & 23.7 & 15.4 & 13.5 & 7.1 \\ & 3 & 54.6 & 30.9 & 18.2 & 16.9 & 10.4 \\ & 4 & 56.0 & 33.2 & 21.3 & 17.2 & 10.4 \\ & 5 & 59.6 & 35.0 & 23.0 & 18.2 & 11.2 \\ \midrule MindSkip & 1 & 22.0 & 12.2 & 8.9 & 8.1 & 2.8 \\ & 2 & 36.9 & 22.0 & 13.4 & 12.8 & 5.1 \\ & 3 & 48.9 & 27.6 & 17.2 & 15.6 & 6.9 \\ & 4 & 51.1 & 30.9 & 19.7 & 17.4 & 8.6 \\ & 5 & 54.6 & 31.5 & 22.0 & 20.3 & 10.4 \\ \midrule FlexiDepth & 1 & 15.6 & 7.7 & 5.6 & 3.1 & 2.8 \\ & 2 & 27.0 & 13.4 & 7.8 & 6.2 & 4.6 \\ & 3 & 44.0 & 19.9 & 16.2 & 12.2 & 6.9 \\ & 4 & 44.7 & 21.7 & 18.7 & 13.5 & 6.9 \\ & 5 & 47.5 & 23.7 & 19.0 & 14.6 & 6.9 \\ \midrule DR.LLM & 1 & 35.5 & 23.4 & 20.3 & 16.9 & 11.4 \\ & 2 & 54.6 & 26.7 & 25.1 & 23.2 & 18.8 \\ & 3 & 56.0 & 33.2 & 27.8 & 24.5 & 19.3 \\ & 4 & 56.7 & 37.1 & 31.9 & 24.5 & 20.1 \\ & 5 & 63.1 & 40.9 & 34.7 & 28.6 & 23.9 \\ \midrule \textbf{\textsc{PoLar}} & 1 & 53.9 & 36.8 & 28.4 & 21.1 & 20.3 \\ & 2 & 58.2 & 39.8 & 32.2 & 25.3 & 22.8 \\ & 3 & 61.7 & 43.0 & 34.9 & 27.9 & 25.1 \\ & 4 & 64.5 & 47.2 & 40.0 & 29.4 & 27.9 \\ \rowcolor{blue!10} & 5 & \textbf{66.0} & \textbf{51.0} & \textbf{42.3} & \textbf{32.8} & \textbf{28.7} \\ \midrule $\Delta$ vs.\ Base (sampling) & 5 & \textbf{+15.6} & \textbf{+17.5} & \textbf{+15.0} & \textbf{+10.1} & \textbf{+12.2} \\ \bottomrule \end{tabular}} ``` ```{=latex} \resizebox{0.6\textwidth}{!}{ \begin{tabular}{llccccc} \toprule \textbf{Method} & \textbf{p@} & \textbf{DM-1} & \textbf{DM-2} & \textbf{DM-3} & \textbf{DM-4} & \textbf{DM-5} \\ \midrule \rowcolor{gray!15}\textbf{Qwen2.5-3B-Instruct}\\ Base ($\tau$=0) & 1 & 17.7 & 7.1 & 3.8 & 3.4 & 1.5 \\ \midrule Base (sampling) & 1 & 17.7 & 9.2 & 4.3 & 2.3 & 1.5 \\ & 2 & 22.0 & 10.7 & 5.6 & 3.6 & 2.3 \\ & 3 & 24.8 & 12.5 & 6.3 & 3.9 & 2.3 \\ & 4 & 27.0 & 13.1 & 6.3 & 4.4 & 2.5 \\ & 5 & 28.4 & 13.4 & 6.8 & 5.5 & 2.5 \\ \midrule ShortGPT & 1 & 12.1 & 3.0 & 1.5 & 1.3 & 0.3 \\ & 2 & 35.5 & 10.7 & 5.6 & 2.3 & 2.3 \\ & 3 & 35.5 & 12.5 & 7.3 & 5.2 & 3.0 \\ & 4 & 41.1 & 13.9 & 8.6 & 8.6 & 5.8 \\ & 5 & 53.9 & 21.7 & 13.9 & 12.8 & 9.1 \\ \midrule MindSkip & 1 & 12.8 & 3.9 & 1.5 & 1.3 & 1.3 \\ & 2 & 24.8 & 10.4 & 5.8 & 5.2 & 3.6 \\ & 3 & 35.5 & 14.2 & 9.1 & 7.6 & 4.3 \\ & 4 & 39.0 & 18.4 & 12.4 & 10.4 & 6.3 \\ & 5 & 46.1 & 26.1 & 16.5 & 13.5 & 9.1 \\ \midrule FlexiDepth & 1 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 \\ & 2 & 0.7 & 0.3 & 0.3 & 0.0 & 0.5 \\ & 3 & 7.8 & 2.4 & 0.3 & 1.8 & 0.5 \\ & 4 & 21.3 & 6.8 & 3.5 & 5.2 & 2.5 \\ & 5 & 25.5 & 8.3 & 4.3 & 5.5 & 2.5 \\ \midrule DR.LLM & 1 & 1.4 & 6.5 & 3.5 & 2.6 & 0.8 \\ & 2 & 4.3 & 17.5 & 12.7 & 6.2 & 1.8 \\ & 3 & 7.1 & 22.3 & 15.2 & 7.3 & 2.5 \\ & 4 & 8.5 & 23.7 & 15.7 & 9.6 & 2.8 \\ & 5 & 9.9 & 27.6 & 18.7 & 11.5 & 3.8 \\ \midrule \textbf{\textsc{PoLar}} & 1 & 46.1 & 19.9 & 14.9 & 12.0 & 7.1 \\ & 2 & 57.4 & 25.5 & 19.0 & 15.4 & 11.4 \\ & 3 & 63.1 & 28.2 & 21.3 & 17.7 & 13.2 \\ & 4 & 66.0 & 31.5 & 22.5 & 18.2 & 15.2 \\ \rowcolor{blue!10} & 5 & \textbf{68.8} & \textbf{32.6} & \textbf{24.3} & \textbf{20.3} & \textbf{16.0} \\ \midrule $\Delta$ vs.\ Base (sampling) & 5 & \textbf{+40.4} & \textbf{+19.2} & \textbf{+17.5} & \textbf{+14.8} & \textbf{+13.5} \\ \bottomrule \end{tabular}} ``` ```{=latex} \resizebox{0.6\textwidth}{!}{ \begin{tabular}{llccccc} \toprule \textbf{Method} & \textbf{p@} & \textbf{DM-1} & \textbf{DM-2} & \textbf{DM-3} & \textbf{DM-4} & \textbf{DM-5} \\ \midrule \rowcolor{gray!15}\textbf{Qwen3-8B}\\ Base ($\tau$=0) & 1 & 57.4 & 30.6 & 17.5 & 16.7 & 11.7 \\ \midrule Base (sampling) & 1 & 56.7 & 30.0 & 19.0 & 16.9 & 11.9 \\ & 2 & 62.4 & 34.1 & 21.5 & 19.3 & 12.4 \\ & 3 & 64.5 & 35.6 & 22.3 & 19.5 & 13.5 \\ & 4 & 66.0 & 35.9 & 22.5 & 19.8 & 14.5 \\ & 5 & 66.0 & 37.1 & 23.0 & 20.1 & 15.2 \\ \midrule ShortGPT & 1 & 19.1 & 7.4 & 4.6 & 5.2 & 3.3 \\ & 2 & 39.0 & 11.6 & 5.6 & 7.3 & 5.1 \\ & 3 & 51.1 & 15.1 & 8.4 & 8.9 & 6.3 \\ & 4 & 57.4 & 22.0 & 12.2 & 12.5 & 7.4 \\ & 5 & 60.3 & 23.7 & 14.4 & 13.5 & 10.2 \\ \midrule MindSkip & 1 & 47.5 & 19.0 & 10.6 & 12.5 & 7.1 \\ & 2 & 54.6 & 22.8 & 14.4 & 15.6 & 8.6 \\ & 3 & 60.3 & 28.8 & 17.7 & 17.4 & 10.2 \\ & 4 & 67.4 & 31.5 & 19.2 & 19.0 & 11.2 \\ & 5 & 68.8 & 34.7 & 20.0 & 19.8 & 13.2 \\ \midrule FlexiDepth & 1 & 1.4 & 0.3 & 0.3 & 0.3 & 0.3 \\ & 2 & 27.7 & 9.2 & 5.1 & 3.1 & 2.0 \\ & 3 & 34.8 & 13.1 & 8.4 & 6.5 & 2.8 \\ & 4 & 44.0 & 18.1 & 11.4 & 10.7 & 4.1 \\ & 5 & 55.3 & 23.4 & 13.2 & 14.1 & 6.3 \\ \midrule DR.LLM & 1 & 5.0 & 26.1 & 15.7 & 6.8 & 1.3 \\ & 2 & 12.8 & 27.3 & 19.5 & 8.6 & 2.0 \\ & 3 & 17.7 & 40.9 & 25.8 & 10.9 & 3.0 \\ & 4 & 22.0 & 49.6 & 30.9 & 14.1 & 4.1 \\ & 5 & 26.2 & 53.4 & 32.4 & 15.4 & 5.1 \\ \midrule \textbf{\textsc{PoLar}} & 1 & 61.0 & 42.7 & 21.3 & 21.4 & 16.0 \\ & 2 & 69.5 & 47.8 & 27.6 & 26.6 & 19.5 \\ & 3 & 74.5 & 54.9 & 33.2 & 28.6 & 21.6 \\ & 4 & 78.0 & 57.6 & 35.7 & 29.4 & 23.4 \\ \rowcolor{blue!10} & 5 & \textbf{80.1} & \textbf{60.5} & \textbf{38.5} & \textbf{30.7} & \textbf{25.1} \\ \midrule $\Delta$ vs.\ Base (sampling) & 5 & \textbf{+14.1} & \textbf{+23.4} & \textbf{+15.5} & \textbf{+10.6} & \textbf{+9.9} \\ \bottomrule \end{tabular}} ``` ## Out-of-Distribution Generalization {#app:ood_pref} We further evaluate the out-of-distribution (OOD) generalization of PoLar on benchmarks that differ substantially from DART-Math in both format and domain. Tables `\ref{tab:ood_pass1_llama}`{=latex}, `\ref{tab:ood_pass1_qwen2.5}`{=latex}, and `\ref{tab:ood_pass1_qwen8}`{=latex} report pass@1 accuracy on ASDiv, MAWPS, and subject-wise subsets of MMLU-Pro using LLaMA-3.2-3B-Instruct, Qwen2.5-3B-Instruct, and Qwen3-8B, respectively. Across all evaluated models, PoLar consistently improves over standard inference on arithmetic word problem benchmarks (ASDiv and MAWPS), indicating strong transfer from structured mathematical reasoning to natural language problem settings. On MMLU-Pro, which spans diverse domains including mathematics, natural sciences, social sciences, and humanities, PoLar achieves broad and consistent gains across most subject areas. These results suggest that the execution programs learned by PoLar capture general, transferable computation control strategies rather than dataset-specific heuristics. Despite being trained on mathematical reasoning data, PoLar generalizes effectively to heterogeneous domains, highlighting the robustness of program-of-layers inference and its applicability beyond the original training distribution. ```{=latex} \begin{table*}[htbp] \captionsetup{skip=-0pt} \caption{ \textbf{Out-of-distribution (OOD) performance at pass@1.} We report accuracy on OOD benchmarks using LLaMA-3.2-3B-Instruct. } \label{tab:ood_pass1_llama} \resizebox{1.0\textwidth}{!}{ \begin{tabular}{lcc|ccccccccccccc} \toprule & & & \textbf{MMLU-Pro} \\ \cmidrule(lr){4-16} \textbf{Method} & \textbf{ASDiv} & \textbf{MAWPS} & \textbf{Math} & \textbf{Phys} & \textbf{Chem} & \textbf{Law} & \textbf{Eng} & \textbf{Other} & \textbf{Econ} & \textbf{Health} & \textbf{Psych} & \textbf{Bus} & \textbf{Bio} & \textbf{Phil} & \textbf{Hist} \\ \midrule Base ($\tau$=0) & 78.4 & 71.5 & 19.5 & 19.0 & 18.6 & 17.1 & 19.8 & 22.4 & 36.1 & 28.5 & 33.0 & 22.2 & 44.9 & 23.8 & 28.3 \\ ShortGPT & 2.3 & 4.6 & 29.9 & 26.9 & 26.9 & 23.3 & 20.1 & 33.0 & 36.7 & 35.3 & 38.2 & 28.6 & 36.0 & 41.7 & 40.9 \\ MindSkip & 9.0 & 7.5 & \textbf{40.4} & \textbf{47.7} & \textbf{43.3} & 31.8 & \textbf{43.7} & 38.5 & 47.6 & 47.2 & 51.5 & 31.7 & 52.3 & \textbf{49.7} & \textbf{55.6} \\ FlexiDepth & 4.7 & 1.3 & 7.8 & 5.9 & 4.8 & 6.4 & 6.3 & 4.7 & 5.9 & 5.9 & 3.6 & 6.0 & 4.9 & 7.4 & 4.7 \\ DR.LLM & 75.4 & 61.3 & 13.6 & 14.4 & 16.2 & 10.0 & 16.4 & 9.0 & 13.6 & 10.6 & 15.2 & 13.0 & 19.8 & 8.2 & 23.1 \\ \midrule \textsc{PoLar} & \textbf{81.4} & \textbf{73.7} & 40.1 & 39.0 & 41.0 & \textbf{34.3} & 43.2 & \textbf{42.9} & \textbf{50.1} & \textbf{52.4} & \textbf{52.4} & \textbf{32.7} & \textbf{56.6} & 44.5 & 48.8 \\ \bottomrule \end{tabular} } \end{table*} ``` ```{=latex} \begin{table*}[htbp] \captionsetup{skip=-0pt} \caption{ \textbf{Out-of-distribution (OOD) performance at pass@1.} We report accuracy on OOD benchmarks using Qwen2.5-3B-Instruct. } \label{tab:ood_pass1_qwen2.5} \resizebox{1.0\textwidth}{!}{ \begin{tabular}{lcc|cccccccccccccc} \toprule & & & \textbf{MMLU-Pro} \\ \cmidrule(lr){4-17} \textbf{Method} & \textbf{ASDiv} & \textbf{MAWPS} & \textbf{Math} & \textbf{Phys} & \textbf{Chem} & \textbf{Law} & \textbf{Eng} & \textbf{Other} & \textbf{Econ} & \textbf{Health} & \textbf{Psych} & \textbf{Bus} & \textbf{Bio} & \textbf{Phil} & \textbf{Hist} \\ \midrule Base ($\tau$=0) & 49.5 & 36.2 & 26.3 & 28.7 & 27.7 & 24.4 & 37.0 & 33.9 & 47.5 & 40.3 & 52.6 & 31.3 & 62.2 & 33.7 & 34.6 \\ ShortGPT & 1.3 & 1.2 & 8.9 & 9.3 & 10.0 & 12.1 & 12.1 & 12.8 & 12.6 & 13.4 & 8.9 & 11.5 & 10.3 & 9.8 & 12.1 \\ MindSkip & 8.0 & 3.5 & 7.6 & 7.5 & 8.3 & 6.0 & 9.5 & 7.4 & 8.6 & 9.8 & 7.5 & 8.5 & 7.5 & 9.4 & 6.8 \\ FlexiDepth & 2.7 & 0.6 & 9.8 & 9.3 & 10.5 & 10.4 & 11.4 & 11.5 & 10.8 & 11.5 & 9.9 & 8.2 & 10.6 & 11.4 & 11.0 \\ DR.LLM & 0.0 & 0.0 & 4.0 & 4.8 & 2.8 & 3.2 & 3.0 & 3.2 & 2.8 & 2.6 & 1.4 & 4.6 & 5.8 & 3.0 & 4.2 \\ \midrule \textsc{PoLar} & \textbf{78.1} & \textbf{57.7} & \textbf{32.1} & \textbf{35.6} & \textbf{33.0} & \textbf{27.6} & \textbf{41.5} & \textbf{38.9} & \textbf{53.0} & \textbf{46.5} & \textbf{57.4} & \textbf{33.4} & \textbf{66.0} & \textbf{38.9} & \textbf{35.4} \\ \bottomrule \end{tabular} } \end{table*} ``` ```{=latex} \begin{table*}[htbp] \captionsetup{skip=-0pt} \caption{ \textbf{Out-of-distribution (OOD) performance at pass@1.} We report accuracy on OOD benchmarks using Qwen3-8B. } \label{tab:ood_pass1_qwen8} \resizebox{1.0\textwidth}{!}{ \begin{tabular}{lcc|cccccccccccccc} \toprule & & & \textbf{MMLU-Pro} \\ \cmidrule(lr){4-17} \textbf{Method} & \textbf{ASDiv} & \textbf{MAWPS} & \textbf{Math} & \textbf{Phys} & \textbf{Chem} & \textbf{Law} & \textbf{Eng} & \textbf{Other} & \textbf{Econ} & \textbf{Health} & \textbf{Psych} & \textbf{Bus} & \textbf{Bio} & \textbf{Phil} & \textbf{Hist} \\ \midrule Base ($\tau$=0) & 67.1 & 52.1 & 21.4 & 24.8 & 19.2 & 32.6 & 23.4 & 43.0 & 60.4 & 56.2 & 62.8 & 26.6 & 72.2 & 42.3 & 50.1 \\ ShortGPT & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 \\ MindSkip & 20.9 & 13.1 & 26.1 & 33.9 & 31.5 & \textbf{58.0} & 32.1 & \textbf{67.0} & 64.7 & \textbf{68.0} & 71.9 & \textbf{36.6} & 68.8 & \textbf{68.5} & \textbf{65.9} \\ FlexiDepth & 1.0 & 0.4 & 0.9 & 1.0 & 0.6 & 0.6 & 0.5 & 0.8 & 0.9 & 1.0 & 0.1 & 0.8 & 1.0 & 0.6 & 1.3 \\ DR.LLM & 69.4 & 0.0 & 0.6 & 0.2 & 1.8 & 1.8 & 0.8 & 0.6 & 1.2 & 1.4 & 2.0 & 1.2 & 0.8 & 0.6 & 0.8 \\ \midrule \textsc{PoLar} & \textbf{72.8} & \textbf{66.9} & \textbf{26.8} & \textbf{34.6} & \textbf{23.4} & 48.7 & \textbf{35.8} & 57.1 & \textbf{65.2} & 67.4 & \textbf{72.4} & 34.1 & \textbf{75.4} & 56.7 & 62.0 \\ \bottomrule \end{tabular} } \end{table*} ``` ```{=latex} \clearpage ``` # Empirical Details {#app:emp_detail} ## Dataset Details All in-distribution experiments are conducted on **DART-Math**, a structured mathematical reasoning benchmark consisting of five difficulty levels, denoted as **DM-1** to **DM-5**. We construct the in-distribution data from the DART-Math MATH pool. Since the released pool may contain multiple response records associated with the same underlying question, we deduplicate examples within each difficulty level by question text before splitting. We then adopt a difficulty-wise data split, where each difficulty level is split independently into training, validation, and test sets. After deduplication, the five difficulty levels contain 565, 1,349, 1,579, 1,537, and 1,577 examples for DM-1 to DM-5, respectively. For each difficulty level, we use approximately 62.5% of the examples for training, 12.5% for validation, and 25% for testing. Overall, this results in 4,130 training examples, 826 validation examples, and 1,651 test examples. All in-distribution results are reported on the held-out test sets corresponding to the same difficulty level used for training. #### Revision note. This arXiv version updates the DART-Math in-distribution evaluation protocol by removing duplicate questions prior to data splitting. We have rerun the affected in-distribution analyses, and the qualitative conclusions remain unchanged. The out-of-distribution (OOD) results are unaffected. ## Training Configuration We train the PoLar prediction network using supervised learning on the training splits described above. All hyperparameters are selected via validation tuning. #### Optimization. We use the AdamW optimizer and tune the learning rate from {1e-4, 3e-4, 5e-4, 8e-4, 1e-3, 3e-3}, the batch size from {32, 128, 256}, and the number of training epochs from {3, 10} based on validation performance. We adopt a cosine learning rate schedule with linear warmup, and tune the number of warmup steps on the validation split together with the other hyperparameters. ## Predictor Size {#app:predictor_size} The learned PoLar predictor is lightweight compared with the frozen base LLM. Across the evaluated models, the predictor contains approximately 2.1M parameters, corresponding to only 0.01%--0.06% of the base model size. This small footprint makes training and inference inexpensive relative to standard LLM fine-tuning or full-model execution. ```{=latex} \resizebox{0.75\textwidth}{!}{ \begin{tabular}{lccc} \toprule \textbf{Model} & \textbf{Base LLM params} & \textbf{Predictor params} & \textbf{Predictor / Base} \\ \midrule Qwen1.5-MoE-A2.7B-Chat & 14.32B & 2.11M & 0.0148\% \\ Qwen2.5-3B-Instruct & 3.40B & 2.12M & 0.0623\% \\ Qwen3-8B & 8.19B & 2.12M & 0.0258\% \\ LLaMA-3.2-3B-Instruct & 3.61B & 2.11M & 0.0586\% \\ \bottomrule \end{tabular} } ``` ## Direct Prompting We adopt a direct prompting strategy throughout all experiments, without eliciting chain-of-thought or intermediate reasoning. The model is explicitly instructed to output only the final answer in a strictly formatted form. Given a math problem instance `question`, the input prompt is constructed as follows: Solve the following math problem and output ONLY the final answer directly, formatted strictly as \boxed{ANSWER}. ### Problem Start {question} ### Problem End Answer: This prompt design enforces concise answer generation and isolates the effect of latent execution programs from token-level reasoning strategies.