Library of Layers (LoL): Adaptively Ordering Transformer Layers at Runtime

Abdellah Rahmani; Sean McLeish; Ang Li; Sanae Lotfi; Pascal Frossard; Micah Goldblum

cs.LGNov 2, 2026

Library of Layers (LoL): Adaptively Ordering Transformer Layers at Runtime

Authors: Abdellah Rahmani, Sean McLeish, Ang Li, Sanae Lotfi, Pascal Frossard, Micah Goldblum

Organizations: EPFL · University of Maryland · New York University · Columbia University

Abstract

Computer programmers chain together the same reusable computational primitives, for example, conditional statements or loops, but in different orders to implement different algorithms. In contrast, existing language model architectures always execute their layers in the same order. We introduce Library of Layers (LoL), an architecture that dynamically orders its layers during inference to solve the problem at hand. LoL has a single hyperparameter, the size of its global library of transformer blocks, that allows it to transition seamlessly between parameter reuse and sparsity, generalizing the parameter-efficiency and FLOPs-efficiency tradeoffs of dense, depth-recurrent, and mixture of experts architectures in one unified family. We show that LoL achieves competitive performance with all three above baselines when pretrained on text data. On synthetic sudoku and maze data, we observe that the model employs distinct layer orderings to implement different solution strategies, and we can intervene on routing to force a model to apply a specific solution strategy to a problem instance. In a multilingual natural language setting, we find that LoL chooses different layer order patterns based on language and script. Similarly, we train LoL on safety data and fit probes to its router outputs, finding that these probes distinguish safe from unsafe text with high AUROC. LoL provides a path toward a future in which models dynamically compose their own computational graph, enabling more efficient and controllable reasoning.

Figures & tables

Figure 1: We introduce Library of Layers (LoL), a novel architecture where the next executed transformer layer is chosen at runtime. We use strided window attention within the routed core to allow information to flow through the sequence dimension while using the same attention and MLP parameters for all tokens within each window.

Figure 2: Library of Layers uses strided window attention and a top-1 router to allow for efficient training. LoL is made up of a prelude, a routed core, and a coda. The layers in the prelude and coda can attend over the whole sequence. In the core, the sequence is split into windows of

W

tokens, and at every routed depth a router selects one layer from the shared library per window (top-left). Attention is restricted to each window, but the window boundaries shift by

S

tokens between successive routed depths, allowing information to flow across windows (Strided Window Attention, top-right). Window

0

has no history, so its router uses the first token and a learned vector

e^{(\ell)}

for routed depth

\ell

(bottom-left). For window

i>0

, the router averages the mean hidden states of earlier windows

0,\dots,i-1

and selects one layer (bottom-right). With

W=4

and

S=2

, Layers 1 and 3 share the same boundaries while Layer 2 is offset, so information can pass between neighboring windows. Token 5 is highlighted in red to show this: in Layer 1 it is the first token in a new window, while in Layer 2 its window includes tokens 3 and 4.

Figure 3: Library of Layers learns distinct paths for different languages and distinct strategies when solving sudoku puzzles and mazes. Left: PCA projection of the routing paths of a $360M$ parameter LoL model trained on an equal mixture of $10$ languages. We see the paths taken are distinct per language and Latin-script languages use similar paths. Middle and Right: We train $360M$ parameter LoL models to solve sudoku (middle) and mazes (right), using $5$ different strategies. We then force each strategy’s mode path at inference time and see the models use the correct strategy (algorithm) with close to $100\%$ accuracy.

Figure 4: Library of Layers is competitive with vanilla dense models. Left: Cross-entropy training loss over training FLOPs for

1B

models. We see growing gives LoL an early loss advantage and all models finish with comparable training loss. Right: Heatmap of expert utilization of one LoL model. We see experts mainly tend to be used in groups of neighboring depths, but some are used throughout the model.

Model	ARC-E	ARC-C	HellaSwag	PIQA	GSM8K
Random	25.0	25.0	25.0	50.0	0.0
Vanilla	$40.8\pm 1.0$	$28.2\pm 1.3$	$31.4\pm 0.5$	$56.6\pm 1.2$	$29.7\pm 1.3$
Vanilla (sliding window)	$49.5\pm 1.0$	$28.5\pm 1.3$	$31.3\pm 0.5$	$58.8\pm 1.1$	$27.5\pm 1.2$
LoL	$50.3\pm 1.0$	$29.1\pm 1.3$	$31.2\pm 0.5$	$59.7\pm 1.1$	$27.9\pm 1.2$

Table 1: Library of Layers is competitive across benchmarks. We benchmark the same models shown in Figure 4 on ARC, HellaSwag, PIQA, and GSM8K. Performance is broadly comparable, except that the full-attention vanilla model scores substantially lower on ARC-Easy.

Figure 5: Library of Layers is competitive with looped transformer models. Left: Cross-entropy evaluation loss, matching the number of unique layers in the core block of the LoL and looped transformer. We see LoL performs competitively with looped transformer models and offers more flexibility in layer ordering. Error bars are min/max loss over $3$ random seeds. We note larger variance for the LoL models as changing the random seed changes the paths explored during training. Right: Cross-entropy training loss over training FLOPs, averaged over $3$ random seeds per model.

Figure 6: Library of Layers is competitive with mixture of experts models. Left: Cross-entropy training loss over training FLOPs, averaged over

3

random seeds. We see growing gives LoL an early loss advantage and both models finish training with similar training losses. Right: Heatmap of expert utilization of one LoL model. We see some experts are heavily used by specific routed depths or neighboring groups of depths, while others are used diversely over depths, showing LoL learns to use the sparsity effectively.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Increasing depth over training saves FLOPs with only a minor increase in final loss.

Figure 8: Loss-free load balancing ( Wang et al., 2024 ) achieves a lower final cross-entropy loss when compared to auxiliary loss load balancing methods.

Figure 9: Sharing router weights across routed depths causes minimal change in cross-entropy loss.

Figure 10: Varying the attention pattern for dense fixed-depth models, we see full attention achieves the lowest loss.

Figure 11: Evolution of multilingual routing during training.

Readout or control	AUROC
Top-1 ridge	0.8754
Single winning-expert probability	0.8244
Expert-frequency log-likelihood ratio	0.8300
Shuffled labels (99 permutations, mean / max)	0.5014 / 0.5231

Appendix

Table 2: Development-set cross-validation AUROCs for safety readouts and the document-level label-shuffling control.

Strategy	Solving Step Format
Repeated row-major naked-single sweeps.	r5c5 = { 5 }
Hidden singles, scanning rows, columns, then boxes.	row3 5 -> r3c7
Digit-major crosshatching over boxes.	3 box4 -> r6c3
First-in-first-out propagation of forced cells, recording newly forced peers.	set r5c5 = 5 : r5c2 -> { 2 }
Deterministic DFS backtracking: first empty cell, digits ascending.	r1c3 = 1 x

Appendix

Table 3: Sudoku data strategies and data formatting.

Figure 12: Routing agreement on sudoku. Each panel shows one strategy, with window positions as rows and routed depths

k

as columns. Each cell gives the percentage of observed routing decisions that agree with the mode-path expert at that position. Gray cells indicate positions with no observations.

Figure 13: Mode routing paths on sudoku. Each panel shows one strategy, with window positions as rows and routed depths

k

as columns. Cell labels and colors identify the expert assigned to each (strategy, window position, routed depth) tuple.

Figure 14: Strategy adherence and sudoku solution accuracy under forced routing.

Figure 15: Example of sudoku data. We input the grid into the model as a sequence, as shown below.

Strategy	Example Solving Steps
Breadth-first search in discovery order.	r1c1 r1c2 r2c1 r1c3 r2c2 r3c1
Depth-first traversal with backtracking.	r1c1 r1c2 r1c3 r1c4 r1c3 r1c2
A* search using Manhattan distance.	r1c1 r1c2 r2c2 r2c3 r3c3 r3c4
Alternating BFS from the start and goal.	r1c1 r8c8 r1c2 r8c7 r1c3 r8c6
Repeated row-major dead-end filling.	r8c1 r8c2 r8c3 r1c8 r2c8 r3c8

Appendix

Table 4: Maze data strategies and data formatting.

Figure 16: Routing agreement on mazes. Each panel shows one strategy, with window positions as rows and routed depths

k

as columns. Each cell gives the percentage of observed routing decisions that agree with the mode-path expert at that position. The gray row marks window W3 for dead-end filling; no validation observations are available for that window-strategy combination.

Figure 17: Mode routing paths on mazes. Each panel shows one strategy, with window positions as rows and routed depths

k

as columns. Cell labels and colors identify the expert with the highest mean router probability for each (strategy, window position, routed depth) tuple during validation. For dead-end filling, W3 has no validation observations and reuses the W2 expert sequence if reached.

Figure 18: Strategy and maze solution generation under forced routing.

Figure 19: Example of maze data. We input the maze into the model as a sequence, as shown below.

Figure 20: Extended version of Figure 4 with “Vanilla (strided window)” and “LoL (untied router).”

Model	ARC-E	ARC-C	HellaSwag	PIQA	GSM8K
Random	25.0	25.0	25.0	50.0	0.0
Vanilla	$40.8\pm 1.0$	$28.2\pm 1.3$	$31.4\pm 0.5$	$56.6\pm 1.2$	$29.7\pm 1.3$
Vanilla (sliding window)	$49.5\pm 1.0$	$28.5\pm 1.3$	$31.3\pm 0.5$	$58.8\pm 1.1$	$27.5\pm 1.2$
LoL	$50.3\pm 1.0$	$29.1\pm 1.3$	$31.2\pm 0.5$	$59.7\pm 1.1$	$27.9\pm 1.2$
Vanilla (strided window)	$48.5\pm 1.0$	$28.0\pm 1.3$	$30.8\pm 0.5$	$59.0\pm 1.1$	$25.9\pm 1.2$
LoL (untied router)	$50.1\pm 1.0$	$30.0\pm 1.3$	$31.4\pm 0.5$	$57.1\pm 1.2$	$27.7\pm 1.2$

Appendix

Table 5: Extended version of Table 1 with “Vanilla (strided window)” and “LoL (untied router).”

Figure 21: Heatmap of expert utilization of LoL models shown in Figure 5 .

Figure 22: Heatmap of expert utilization of LoL models shown in Figure 6 .

Transformer

Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.

See all 558 papers

Skip a Layer or Loop It? Learning Program-of-Layers in LLMs
Jun 4, 2026Ziyue Li, Yang Li, Tianyi ZhouLLM Inference LLM Inference Acceleration
University of Maryland, College Park, MD, USA · 2MBZUAI, Abu Dhabi, UAE.
Aletheia: Gradient-Guided Layer Selection for Efficient LoRA Fine-Tuning Across Architectures
Apr 4, 2026Abdulmalek SaketLLM Fine-Tuning Efficient Language Model Training
Royal Fenice Kft / ALETHEIA PROTOCOL research, Budapest, Hungary
MACRO: Markov Chain Routing of Transformer Layers
Aug 6, 2026Paweł Batorski, Abtin Pourhadi, Akylgali Aitaza +2Markov Models LLM Inference Acceleration
Heinrich Heine University Düsseldorf · Jagiellonian University · IDEAS Research Institute

Jun 4, 2026J/K move · Enter open · S save

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

Ziyue Li, Yang Li, Tianyi Zhou

University of Maryland, College Park, MD, USA · 2MBZUAI, Abu Dhabi, UAE.

LLM Inference LLM Inference Acceleration LLM Reasoning

Open paper AskPDF arXiv

Large language models (LLMs) perform inference by following a fixed depth and order, non-recurrent execution of all layers. We reveal the wide existence of training-free, flexible, dynamic program-of-layers (PoLar), where pretrained layers can be packed as modules and then skipped or looped to form a customized program for each input. For most inputs, substantially shorter program executions can achieve the same or better accuracy, while incorrect predictions of the original LLM can be corrected by alternative programs with fewer layers. These observations indicate that inference admits multiple valid latent computations beyond the standard forward pass. To efficiently achieve PoLar in practice, we propose a lightweight PoLar prediction network, which learns to generate execution programs that dynamically skip or repeat pretrained layers for each input. Experiments on mathematical reasoning benchmarks demonstrate that PoLar consistently improves accuracy over standard inference and prior dynamic-depth methods, often while executing fewer layers, and that these gains persist under out-of-distribution evaluation. Our results suggest that fixed-depth execution captures only a narrow subset of an LLM's latent reasoning capacity.

Abstract

Transformer

Explore similar work

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

Aletheia: Gradient-Guided Layer Selection for Efficient LoRA Fine-Tuning Across Architectures

MACRO: Markov Chain Routing of Transformer Layers