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

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

We set K max = 4 K_{\max}=4 based on empirical evidence. Finding 4 in Section 2 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 ) [s_{j},s_{j+1}) , the execution program assigns one of three operations { skip , keep , repeat } \{\textsf{skip},\textsf{keep},\textsf{repeat}\} , which determines how the segment is executed: skip \displaystyle\textsf{skip} : ∅ , \displaystyle:\emptyset, keep \displaystyle\textsf{keep} : [ s j , … , s j + 1 − 1 ] , \displaystyle:[s_{j},\ldots,s_{j+1}-1], repeat \displaystyle\textsf{repeat} : [ s j , … , s j + 1 − 1 , s j , … , s j + 1 − 1 ] . \displaystyle:[s_{j},\ldots,s_{j+1}-1,\ s_{j},\ldots,s_{j+1}-1]. The skip operator omits a segment to reduce computation, while repeat applies a single additional pass. Operations are represented by a categorical label vector 𝐳 op ​ ( x ) ∈ { skip , keep , repeat } D , \textstyle\mathbf{z}^{\text{op}}(x)\in\{\textsf{skip},\textsf{keep},\textsf{repeat}\}^{D}, where 𝐳 i op \mathbf{z}^{\text{op}}_{i} is defined only when 𝐳 i seg = 1 \mathbf{z}^{\text{seg}}_{i}=1 (i.e., at segment start positions); labels at all other positions are ignored.

This operator set is intentionally minimal and empirically grounded. Finding 4 in Section 2 shows valid execution programs rarely require more than a single re-execution within a segment, and that skip and repeat account for the most effective execution patterns, offering strong performance–efficiency trade-offs. The keep operator preserves the original computation when no modification is needed. Although our implementation allows at most one additional execution through repeat , the representation is not fundamentally limited to a single recurrence. The operation vocabulary can be extended to { repeat - ​ 2 , … , repeat - ​ k } \{\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 2 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.

### 3.2 Program-of-Layers ( PoLar ) Prediction Network

We train a lightweight predictor to output logits for the program representation defined in Section 3.1 .

Architecture. Given an input x x , we first encode it using a frozen embedding model ( Qwen3-Embedding-0.6B ), as token-level representations 𝐇 = E ⁡ ( x ) ∈ ℝ T × d q , \textstyle\mathbf{H}=E(x)\in\mathbb{R}^{T\times d_{q}}, where T T is the token length and d q d_{q} is the hidden size of the embedding model. We project token representations to a working dimension d d : 𝐇 ~ = 𝐇𝐖 h ∈ ℝ T × 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 ∈ { 0 , … , D − 1 } i\in\{0,\ldots,D-1\} with a learnable embedding 𝐞 i ∈ ℝ d \mathbf{e}_{i}\in\mathbb{R}^{d} , and stack them as 𝐄 ∈ ℝ D × d \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: 𝐗 = MHA ​ ( 𝐐 , 𝐊 , 𝐕 ) , 𝐐 = 𝐄 , 𝐊 = 𝐇 ~ , 𝐕 = 𝐇 ~ , \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 𝐗 ∈ ℝ D × d \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: 𝐗 ′ = Enc layer ​ ( 𝐗 ) ∈ ℝ D × d . \textstyle\mathbf{X}^{\prime}=\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: ℓ seg = 𝐗 ′ ​ 𝐖 seg + 𝐛 seg ∈ ℝ D , ℓ op = 𝐗 ′ ​ 𝐖 op + 𝐛 op ∈ ℝ D × 3 . \bm{\ell}^{\text{seg}}=\mathbf{X}^{\prime}\mathbf{W}_{\text{seg}}+\mathbf{b}_{\text{seg}}\in\mathbb{R}^{D},\bm{\ell}^{\text{op}}=\mathbf{X}^{\prime}\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 2 ). Each program is deterministically parsed into program representation, producing ground-truth segmentation and operation labels 𝐳 seg ​ ( x ) \mathbf{z}^{\text{seg}}(x) and 𝐳 op ​ ( x ) \mathbf{z}^{\text{op}}(x) in the format defined in Section 3.1 . 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 ( 𝐳 seg ∗ ( x ) , 𝐳 op ∗ ( x ) ) \big(\mathbf{z}^{\text{seg}*}(x),\mathbf{z}^{\text{op}*}(x)\big) . Let p i seg = σ ⁡ ( ℓ i seg ) p^{\text{seg}}_{i}=\sigma(\ell^{\text{seg}}_{i}) and 𝐩 i op = Softmax ​ ( ℓ i op ) \mathbf{p}^{\text{op}}_{i}=\textsc{Softmax}(\bm{\ell}^{\text{op}}_{i}) . Segmentation is supervised with binary cross-entropy over boundary indicators: ℒ seg = − ∑ i = 0 D − 1 [ 𝐳 i seg ∗ log p i seg + ( 1 − 𝐳 i seg ∗ ) log ( 1 − p i seg ) ] . \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 = 𝐳 seg ∗ i m_{i}=\mathbf{z}^{\text{seg}*}_{i} , we compute ℒ op = − ∑ i = 0 D − 1 m i ⋅ log 𝐩 i op [ 𝐳 i op ∗ ] . \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 ℒ = ℒ seg + ℒ op . \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 ℓ seg \bm{\ell}^{\text{seg}} . If any resulting segment exceeds the maximum length constraint K max K_{\max} , additional boundaries are inserted to enforce it, yielding segment start positions { s j } \{s_{j}\} . Conditioned on this segmentation, we compute operation log-probabilities at each segment start from the predicted logits: log ⁡ p ⁡ ( o j ∣ x , s j ) = log ⁡ Softmax ​ ( ℓ s j op ) ​ [ o j ] . \textstyle\log p(o_{j}\mid x,s_{j})=\log\textsc{Softmax}\!\big(\bm{\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 π ⁡ ( x ) \pi(x) . Finally, each candidate program is mapped deterministically to a concrete executed program using the segment-to-path rules in Section 3.1 .

## 4 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.

### 4.1 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 ( Tong et al., 2024 ) , 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 ( Miao et al., 2020 ) and MAWPS ( Kadlčík et al., 2023 ) , which focus on arithmetic word problems, as well as subject subsets from MMLU-Pro ( Wang et al., 2024 ) 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 k accuracy , defined as the probability that at least one of the top- k k candidates produces a correct answer. For PoLar , the k k candidates correspond to the top- k k predicted execution programs selected via beam search. For sampling-based baselines, k 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 ( τ = 0 \tau=0 ) uses greedy decoding with temperature τ = 0 \tau=0 . Base (sampling) samples k k outputs using stochastic decoding with τ ∈ { 0.3 , 0.7 , 1.0 } \tau\in\{0.3,0.7,1.0\} and reports the best result across temperatures, increasing output diversity without altering internal execution. DR.LLM ( Heakl et al., 2025 ) learns layer-routing policies from execution paths and applies them at inference time. ShortGPT ( Men et al., 2025 ) statically prunes layers based on estimated importance, yielding a reduced-depth model. MindSkip ( He et al., 2024 ) and FlexiDepth ( Luo et al., 2025 ) learn router-based dynamic-depth policies, primarily optimized for inference efficiency. Several approaches, such as Mixture-of-Depths ( Raposo et al., 2024 ) , LaCo ( Yang et al., 2024 ) , and Mixture-of-Recursions ( Bae et al., 2025 ) , 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 D .

### 4.2 Main Results

We evaluate in-distribution performance on DART-Math. Table 2 reports pass@ k k results using LLaMA-3.2-3B-Instruct , with complete results provided in Appendix C.1 .

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 k ) monotonically improves PoLar , evidencing effective test-time computation scaling through execution-program exploration. Figure 8 (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 8 (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.

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 4 , 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 D.3 .

### 4.3 Out-of-Distribution Performance

We evaluate the OOD generalization of execution programs learned from in-distribution data. As shown in Table 3 , PoLar consistently outperforms standard inference on all OOD benchmarks using Qwen1.5-MoE-A2.7B-Chat , with full results reported in Appendix C.2 .

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.

## 5 Related Works