Research map / Large language models / Reasoning and chain of thought

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

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Routing, pruning and distillation for reasoning · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows that frozen pretrained LLMs admit shorter or better layer programs (skip/repeat contiguous segments) per input, and trains a lightweight predictor to output input-specific layer programs in one shot, improving math reasoning accuracy over standard inference and dynamic-depth baselines.

Problem

Standard LLM inference uses a fixed depth and order over all layers for every input, though some inputs need less computation and others might benefit from different layer execution; searching for such programs per input (MCTS) is too expensive for inference.

Approach

Treat pretrained layers as atomic functions packed into contiguous segments, each assigned skip/keep/repeat, and train a lightweight cross-attention predictor over frozen embeddings to output an input-specific program in a single shot, supervised by MCTS-discovered valid programs.

Key results

• On LLaMA-3.2-3B-Instruct DART-Math, PoLar pass@5 averages 51.2% vs 43.8% for Base (sampling), a +7.4 point gain.
• PoLar pass@1 on DM-1 for LLaMA-3.2-3B rises from 48.9% (Base sampling) to 54.6%.
• PoLar predictor adds about 0.23 layer-equivalents of latency (0.8% of forward pass) and reaches 0.83x base runtime on easy DM-1 inputs.

Surprising findings

• Skip-only and loop-only MCTS searches are substantially weaker than joint skip-and-loop search, with recurrence-only outperforming skipping-only.
• Shorter valid programs exist for 71.9% of initially correct inputs, suggesting standard inference over-computes.

Limitations

• Segments capped at length 4 and repeat limited to a single extra pass, restricting the program space.
• Evaluated only on mathematical reasoning benchmarks plus arithmetic and MMLU-Pro subsets; OOD results depend on the predictor generalizing from math data.
• MCTS supervision is expensive offline and used only as a diagnostic on a small in-distribution set (about 4k examples).

Experimental setup

Datasets: DART-Math ASDiv MAWPS MMLU-Pro

Benchmarks: DART-Math ASDiv MAWPS MMLU-Pro

Metrics: pass@k accuracy accuracy latency average executed layers

Baselines: Base (greedy) Base (sampling) ShortGPT MindSkip FlexiDepth DR.LLM

Builds on: MCTS DR.LLM ShortGPT LayerSkip

Closest papers by meaning

• AdaSkip: Adaptive Sublayer Skipping for Accelerating Long-Context LLM Inference (AAAI 2025 · similarity 0.89)
• SkipGPT: Each Token is One of a Kind (ICML 2025 · similarity 0.89)
• Polar Sparsity: High Throughput Batched LLM Inferencing with Scalable Contextual Sparsity (NeurIPS 2025 · similarity 0.89)
• EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism (ICML 2024 · similarity 0.89)
• MISA: Memory-Efficient LLMs Optimization with Module-wise Importance Sampling (NeurIPS 2025 · similarity 0.89)
• Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA (ICLR 2025 · similarity 0.89)