PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2611.03526v1cs.LGcs.AIcs.CL

TL;DR

Transformer layers usually execute in a fixed order, limiting input-dependent composition of reusable computations. Library of Layers (LoL) routes inputs through an ordered sequence of blocks from a shared library, with library size spanning reuse and sparsity. The paper reports competitive performance across these regimes and tasks, and finds that routing patterns can reflect and control solution strategies.

  • Problem

    Existing transformer architectures execute layers in a fixed order, unlike computational systems that compose reusable primitives in different orders for different tasks.

    Source: §1 · p. 1

  • Method

    LoL uses a router to select and order transformer blocks from a global library conditional on the input, with library size controlling reuse and sparsity.

    Source: §1 · pp. 1–2

  • Results

    LoL is competitive across reuse and sparse regimes and on FLOPs-matched math data, while learned routing patterns reflect task strategies and can be forced to apply them.

    Source: §1 · pp. 2–3, §4.1.3 · p. 7, §4.1.4 · p. 8

  • Takeaways & Limitations

    Layer ordering can become learned computation, and extracted routing paths provide a mechanism to study and control latent computation.

    Source: §1 · p. 2

Abstract

from arXiv · show

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.

1 INTRODUCTION

Library of Layers (LoL) lets a router choose and order transformer blocks from a shared library based on the input, addressing fixed layer order while spanning parameter reuse and sparsity through library size. The paper reports competitive performance across these regimes and tasks, and shows that learned routing patterns can be interpreted and controlled.

1 INTRODUCTION

  • Fixed layer order may require useful computational primitives to be represented repeatedly at different depths, motivating reusable blocks composed in different orders.Source: §1 · p. 1
  • A router selects and orders blocks from a global transformer-block library conditional on the input, making layer choice, order, and reuse count adaptive.The router is trained end-to-end without explicit supervision.Source: §1 · pp. 1, 3
  • Windowed attention routes different windows while preserving causality and shifting boundaries across depths so information can flow between neighboring windows.Source: §1 · p. 2
  • LoL remains competitive across reuse and sparse regimes, including against looped-transformer and MoE-like compute tradeoffs, and is competitive with dense transformers on FLOPs-matched math data.The experiments vary library size on web text and also evaluate 1B-parameter models on math data.Source: §1 · pp. 2–3
  • On sudoku and maze tasks, distinct learned routing patterns correspond to solution strategies, and forcing a routing path makes the model follow the selected strategy.In multilingual modeling, routing patterns vary by language and script; probes on router outputs distinguish safe from unsafe text with high AUROC.Source: §1 · p. 2
  • The authors present layer ordering as learned computation and routing patterns as a means to study and control latent computation.Source: §1 · p. 2

2 RELATED WORK

Prior work studies layer reuse, sparsity, and dynamic skipping, but generally retains a fixed execution order. LoL combines a shared global library with input-dependent selection and ordering of complete transformer blocks.

2 RELATED WORK

  • Looped transformers reuse blocks across depth, while mixture-of-experts models increase sparsity; LoL spans both regimes by varying the size of its shared block library.Prior work also explores recurrence in synthetic tasks and parameter sharing in other transformer components.Source: §2 · p. 3
  • Gong et al. (2024) add reuse after pretraining, use top-k routing with k greater than one, and combine selected experts' attention weights and FFN outputs.LoL instead uses top-1 routing and dynamically selects each layer in the final model.Source: §2 · p. 3
  • Prior layer reuse and dynamic-skipping approaches still enforce execution order, whereas LoL makes block ordering and reuse input-dependent.Source: §2 · p. 3

3 HOW TO LEARN THE ORDER OF A FORWARD PASS

LoL places an adaptive routed core between fixed-order transformer stacks and uses windowed routing to select blocks while preserving causal information flow. A router uses prior-window information, and training gradually increases routed depth.

3.1 THE LIBRARY

  • A prelude and coda of fixed-order decoder layers surround a routed core that applies blocks from a shared library; library size counts distinct blocks and core depth counts applications.A sequence of block applications is called a routing path, and different token windows may follow different paths.Source: §3.1 · pp. 3–4

3.2 ROUTING

  • Routing windows separately avoids the attention mismatch and instability of assigning neighboring tokens to different experts, while prior-window information preserves causality.Window boundaries shift between depths to let information flow between adjacent chunks.Source: §3.2 · p. 4

3.2.1 STRIDED WINDOW ATTENTION

  • Window boundaries shift by stride S between routed depths, so overlapping windows pass information between neighboring regions while attention remains causal within each window.For W = 4 and S = 2, depth 1 uses windows 0–3 and 4–7, while depth 2 uses 0–1, 2–5, and 6–7.Source: §3.2.1 · pp. 4–5, Figure 2 · p. 4

3.2.2 THE ROUTER

  • For each window after the first, a linear router uses the mean hidden states of earlier windows; the first window uses its first token and a learned depth-specific vector.Source: §3.2.2 · p. 5

3.3 INCREASING DEPTH DURING TRAINING

  • Training gradually increases routed core depth, reducing compute early while ultimately producing a deeper model for deployment.The authors call this procedure growing.Source: §3.3 · p. 5

4 EXPERIMENTS

Experiments trace LoL’s routing on multilingual, safety, and synthetic tasks, then evaluate math training and compare LoL with dense, looped, and sparse mixture-of-experts models.

4.1.1 MULTILINGUAL

  • Routing paths are reasonably distinct by language, while Latin-script languages share similar paths; the same layers can be used in different orders.Source: §4.1.1 · p. 6

4.1.2 SAFETY

  • 0.92 per-window AUROC and 0.95 after averaging windows: router-probability ridge probes distinguish score-0 safe from score-4/5 unsafe sequences.Source: §4.1.2 · p. 6

4.1.3 SUDOKU

  • Close to 100% accuracy: forcing each sudoku strategy’s mode path makes the model use the expected strategy on new grids.Source: §4.1.3 · p. 7

4.2 MATHEMATICAL REASONING AT THE 1B SCALE

  • Within 2 standard errors, LoL evaluations are competitive across ARC-Easy, ARC-Challenge, HellaSwag, PIQA, and GSM8K; the vanilla model scores poorly on ARC-Easy.Source: Table 1 · p. 8, §4.2 · p. 8

4.1.4 MAZES

  • Nearly 100% accuracy: forcing a maze strategy’s mode path produces its requested trace, even though strategies share syntax and differ in cell-exploration order.Source: §4.1.4 · p. 8

4.2 MATHEMATICAL REASONING AT THE 1B SCALE

  • LoL finishes with training loss competitive with parameter-equivalent dense baselines on synthetic math data, while experts cluster across neighboring depths and some span wider ranges.Source: Figure 4 · p. 7, §4.2 · p. 8

4.3.1 COMPARISON TO LOOPED TRANSFORMERS

  • LoL is competitive in loss with parameter-matched looped models, with slightly higher final-loss variance across three seeds; enlarging its library permits sparse behavior.Source: §4.3.1 · p. 9, Figure 5 · p. 9

4.3.2 COMPARISON TO SPARSE MIXTURE OF EXPERTS

  • 0.05 (2.75 vs. 2.8): the MoE baseline has slightly lower training loss than LoL, whose experts specialize at some routed depths and span others.Source: §4.3.2 · p. 9, Figure 6 · p. 10

5 DISCUSSION

LoL unifies parameter reuse and sparsity, with library size controlling the shift among reused, dense, and sparse computation. The discussion highlights routing-based control and identifies several directions for future work.

5 DISCUSSION

  • Changing library size lets LoL span parameter reuse and sparsity, including regimes associated with looped models, feedforward transformers, and MoEs.The authors describe this as one model family, with library size as the control.Source: §5 · p. 10
  • LoL supports routing interventions for language and algorithm patterns, while whether router features enable chat-setting interventions remains open.The discussion also notes that probes distinguish safe from unsafe text, but does not establish chat interventions.Source: §5 · p. 10
  • Future directions include increasing computational depth at inference and routing over growing libraries of test-time learnable adapters.The authors present both as promising possibilities, not demonstrated results.Source: §5 · p. 10

A ADDITIONAL ARCHITECTURE AND TRAINING DETAILS

This section specifies LoL's input-output setup, model and tokenizer choices, router training, normalization, and general training conventions.

A ADDITIONAL ARCHITECTURE AND TRAINING DETAILS

  • Given core depth D, library size L, and input sequence x ∈V n, the architecture groups produce output probabilities p ∈Rn×|V | as specified in Algorithm 1.The passage states the interface and directs readers to Algorithm 1 for the group definitions.Source: A ADDITIONAL ARCHITECTURE AND TRAINING… · p. 16
  • Main experiments use 360M-parameter GPT-2 Medium-based models, with a Llama-3.2-1b-based shape for the 1B model and task-specific tokenizers and window settings.Natural-language experiments use context length 1024; tokenizer choice varies between predominantly English and multilingual experiments.Source: A ADDITIONAL ARCHITECTURE AND TRAINING… · p. 16
  • Router training uses loss-free load balancing: a per-expert bias affects selection, while the selected expert's softmax probability preserves a differentiable path from loss to router.Selection itself is discrete and is not differentiated through.Source: A ADDITIONAL ARCHITECTURE AND TRAINING… · p. 16
  • LoL blocks use learned depth-dependent scale and bias after normalization, and experiments use Muon with AdamW plus document-blocked attention.Tokens in one document never attend to tokens in another document.Source: A ADDITIONAL ARCHITECTURE AND TRAINING… · pp. 16–17

B ADDITIONAL ARCHITECTURE ABLATIONS

The ablations examine growing LoL during training, router load balancing and weight sharing, and attention patterns in dense baselines.

B.1 LIBRARY OF LAYERS

  • Increasing depth over training saves FLOPs with only a minor increase in final loss.Figure 7 reports the compute-loss tradeoff for growing LoL during training.Source: Figure 7 · p. 17, §B.1 · p. 17
  • Loss-free load balancing achieves lower final cross-entropy loss than auxiliary-loss load balancing methods.Figure 8 is the comparison described in the ablation overview.Source: Figure 8 · p. 17, §B.1 · p. 17
  • Sharing router weights across routed depths has minimal impact on final training loss.This is the Figure 9 ablation finding.Source: §B.1 · p. 17, Figure 9 · p. 18

B.2 BASELINES

  • Among varied attention patterns in dense fixed-depth models, full attention achieves the lowest loss, while pattern changes have a small impact on training loss.The Figure 10 caption identifies full attention as lowest-loss; the text characterizes the overall effect of varying patterns as small.Source: §B.2 · p. 17, Figure 10 · p. 18

C ADDITIONAL EXPERIMENTAL INFORMATION

Figure 11 concerns how multilingual routing changes during training.

C ADDITIONAL EXPERIMENTAL INFORMATION

  • Figure 11 tracks the evolution of multilingual routing during training.The supplied caption does not state a particular routing outcome.Source: Figure 11 · p. 19

C.2 SAFETY

The safety analysis compares router-based readouts and a label-shuffling control using development-set cross-validation AUROC. Simpler readouts recover much of the ridge probe’s performance, but the labels are noisy classifier-generated labels.

C.2 SAFETY

  • The evaluation reports mean AUROC over five development-set cross-validation folds and includes a document-level label-shuffling control.The readouts include ridge regression on top-1 expert probabilities, a single winning-expert probability, and an expert-frequency log-likelihood ratio.Source: §C.2 · p. 20
  • Simpler safety readouts recover much of the top-1 ridge probe’s performance, while document-level label shuffling reduces ridge performance to chance.The labels come from a classifier trained by Maini et al. (2026), so the reported results measure agreement with noisy labels.Source: §C.2 · p. 20
  • Table 2 covers development-set cross-validation AUROCs for safety readouts and the document-level label-shuffling control.Source: Table 2 · p. 20

C.3 SUDOKU

The sudoku experiments use puzzles solved under five strategies and assess whether LoL follows a requested strategy under forced routing. The reported results indicate close to 100% accuracy on strategy use and puzzle solving.

C.3 SUDOKU

  • The dataset contains 20M sequences, about 10B training tokens after deduplication, generated from puzzles with 30–36 digits and solved by five strategies.Each puzzle contributes one sequence for each strategy; Table 3 covers the strategies and data formatting.Source: §C.3 · p. 20, Table 3 · p. 20
  • The evaluation reports strategy adherence, trace legality, rule-faithfulness, answer correctness, and the fraction of padding tokens before end of sequence.Source: §C.3 · p. 23
  • LoL uses the expected sudoku strategy with legal traces and solves the grid with close to 100% accuracy, without adding unneeded padding.Source: Figure 14 · p. 22, §C.3 · p. 23

C.3.1 SUDOKU DATA EXAMPLE

The appendix illustrates sudoku and maze examples as sequences of inputs, strategy labels, reasoning traces, and answers, alongside routing visualizations for the maze strategies.

C.3.1 SUDOKU DATA EXAMPLE

  • Sudoku examples are formatted as an input grid, a strategy label, solving steps, a final grid, and an end-of-sequence token.Figure 15 visualizes an example grid and sequences for solving it under each strategy.Source: §C.3.1 · p. 23, Figure 15 · p. 23

S1: Naked singles

  • The example presents a Sudoku grid labeled Q, an S1 solving trace labeled R with listed cell assignments, and a completed answer grid labeled A.Source: S1: Naked singles · p. 24

S4: Propagation

  • The propagation example presents a sudoku trace with explicit placements and propagated candidate values, ending in the completed grid.Source: S4: Propagation · p. 25

S5: Backtracking

  • The backtracking example records trial placements, rejected branches, and reversals before reaching the completed grid.Source: S5: Backtracking · p. 25

C.4 MAZE

  • Maze data use loop-free 8×8 mazes with distinct start and goal cells, strategy-specific traces, and approximately 10B training tokens across 36M sequences.The construction requires unique solving traces, paths of at least 12 cells, at least one DFS backtrack, and at least three BFS expansions outside the solution path.Source: §C.4 · pp. 25–26, Table 4 · p. 26

S1 BFS

  • Maze routing visualizations show agreement with mode-path experts and identify the highest-mean-probability expert by strategy, window position, and routed depth.For dead-end filling, W3 has no validation observations and reuses the W2 expert sequence if reached.Source: Figure 16 · p. 26, Figure 17 · p. 27

S1 BFS S2 DFS walk

C.4.1 MAZE DATA EXAMPLE

The appendix provides a maze input-and-trace example, examples of several search strategies, and extended benchmark and expert-utilization visualizations.

C.4.1 MAZE DATA EXAMPLE

  • Maze examples are formatted as maze, start and goal cells, a strategy label, search trace, solution path, and end-of-sequence token.Figure 19 shows an example maze with sequences for each strategy.Source: §C.4.1 · pp. 27–28, Figure 19 · p. 28

S1: BFS

  • The BFS example lists the cells explored in its search trace and the resulting path from the start cell to the goal.Source: S1: BFS · p. 28

S2: DFS walk

S4: Bidirectional BFS

S5: Dead-end filling

C.5 MATH

  • The math appendix extends Figure 4 and Table 1 with strided-window vanilla and untied-router LoL variants, using five-shot evaluation except three-shot GSM8K.The three-shot GSM8K setting fits within the 1024-token context.Source: §C.5 · p. 29, Table 5 · p. 29, Figure 20 · p. 30

C.6 LOOPED TRANSFORMERS

C.7 MOE

Loading 2611.03526v1…