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

Summary (Overview)

• Novel Architecture: The paper introduces Library of Layers (LoL), a transformer architecture where a router dynamically selects the next transformer block from a global library at each routed depth during inference, making layer ordering input-dependent rather than fixed.
• Unified Framework: By varying a single hyperparameter (library size L), LoL generalizes across three architectural regimes: dense transformers (L ≈ D), depth-recurrent/looped models (small L), and mixture of experts (large L), spanning a continuum of parameter-reuse and FLOPs-efficiency tradeoffs.
• Strided Window Attention: The paper introduces a windowed attention masking strategy that allows different token windows to adaptively select different layer orderings while preserving causality and preventing Q-K mismatch issues.
• Interpretability and Control: LoL learns distinct routing paths for different languages, solution strategies (sudoku, mazes), and safety labels. Forcing a strategy's learned routing path at inference time causes the model to follow that strategy with ~100% accuracy.
• Competitive Performance: LoL matches or exceeds baselines including vanilla dense transformers, looped transformers (Huginn), and sparse MoE models across benchmarks including ARC, HellaSwag, PIQA, and GSM8K.

Introduction and Theoretical Foundation

Background and Motivation

Traditional transformer language models execute their layers in a fixed, predefined order for every input. This rigidity contrasts sharply with modular programming, where reusable computational primitives are dynamically composed in different orders to implement different algorithms. The authors argue this inflexibility makes models less efficient, as useful computational primitives may need to be represented repeatedly at different depths.

Prior Work and Limitations

The paper situates itself relative to three research directions:

1. Parameter Sharing: Looped transformers (Dehghani et al., 2018; Lan et al., 2019; Geiping et al., 2026) repeatedly reuse the same transformer blocks, achieving parameter efficiency and even Turing completeness (Giannou et al., 2023).
2. Mixture of Experts (MoE): MoE layers (Shazeer et al., 2017) increase sparsity by routing tokens to subsets of parameters, as used in models like GPT-4 and other large-scale systems.
3. Dynamic Layer Skipping: Some work (Elhoushi et al., 2024; Raposo et al., 2024; Bae et al., 2025) investigates dynamically skipping layers during the forward pass.

However, all prior approaches still enforce a fixed order in which layers must be executed. The key insight is that layer ordering itself can be a learned, input-dependent computation.

Key Challenges

The authors identify a fundamental challenge with dynamic layer selection: allowing different token positions to select different layer orderings creates a Q-K mismatch problem. If queries and keys are computed by different attention layers at different token positions, the attention computation becomes unstable and harms performance. Conversely, choosing a single layer ordering for the entire context risks breaking causality or making decisions with insufficient information.

Methodology

Architecture Overview

LoL is built from decoder-only transformer blocks organized into three components:

• Prelude (P): Fixed-order stack of decoder layers with full causal attention at the beginning
• Routed Core: Applies blocks selected from a library $\{E_1, \dots, E_L\}$
• Coda (C): Fixed-order stack of decoder layers with full causal attention at the end

Key parameters:

• Library size (L): Number of distinct blocks stored
• Core depth (D): Number of block applications executed
• Each library block is called an expert; each application is a routed depth

Strided Window Attention

To enable different token positions to follow different routing paths while preventing Q-K mismatch:

1. The sequence is chunked into windows of size W at each routed depth
2. All tokens in a window use the same expert with causal attention restricted to that window
3. Window boundaries shift by S tokens between successive depths

The window assignment at depth ℓ is given by:

$$w_i^{(\ell)} = \left\lfloor \frac{i + \delta_\ell}{W} \right\rfloor$$

where the boundary shift is:

$$\delta_\ell = (\ell - 1)S \mod W$$

For example, with n = 8, W = 4, and S = 2:

• Depth 1: windows {0,...,3} and {4,...,7}
• Depth 2: windows {0,1}, {2,...,5}, and {6,7}
• Depth 3: returns to the first split

This striding allows information to flow between adjacent windows while preserving causality.

Router Design

The router is a linear layer that selects experts per window:

• For window k > 0: Input is the mean over prior windows of the mean hidden state per window
• For the first window (k = 0): Input is the mean of a learned depth-specific vector $\mathbf{e}^{(\ell)}$ and the first token's hidden state

Training techniques:

• Loss-free load balancing (Wang et al., 2024; Liu et al., 2024)
• Gaussian noise added to router logits before softmax for exploration (following Shazeer et al., 2017)
• FiLM conditioning (Perez et al., 2018): learned depth-dependent scale and bias after each normalization
• Top-1 routing (unlike MoE's top-k with k > 1)

Depth Growing

Training begins with a shallow routed core and gradually increases D over training using a 1−sqrt schedule for 75% of training. This saves compute while ultimately constructing a deep, expressive model. The growing schedule follows:

$$D(t) = D_{\text{max}} \cdot (1 - \sqrt{1 - t/T})$$

for training progress t over total steps T.

Empirical Validation / Results

1. Multilingual Routing Patterns

Trained 360M parameter models on an equal mix of 10 languages from FineWeb-2 for 10B tokens. PCA projection of routing paths (one-hot encoded expert selections per depth) reveals:

• Routing paths are distinct per language
• Latin-script languages share similar paths
• The model gradually learns to distinguish languages over training

2. Safety Detection

Trained on an equal mixture of safe and unsafe text (using harmfulness scores from Maini et al., 2026):

• Window-level probe: Ridge regression on router probabilities achieves AUROC of 0.92
• Sequence-level probe: Router probabilities averaged across windows achieve AUROC of 0.95
• Label shuffling controls yield chance performance

3. Sudoku Strategy Control

Five sudoku solving strategies with distinct step syntaxes were trained. Strategy descriptions were masked from all parts of the network except the router. Results:

• Forcing each strategy's mode path at inference produces the expected strategy with close to 100% accuracy
• LoL solves sudoku puzzles with close to 100% accuracy when given the strategy token

4. Maze Exploration Order

Five maze-solving strategies (BFS, DFS, A*, bidirectional search, dead-end filling) share the same output syntax and differ only in cell exploration order. Forcing LoL to follow a particular computational path produces the requested algorithm's trace with nearly 100% accuracy, demonstrating that layer ordering controls exploration order itself.

5. Mathematical Reasoning at 1B Scale

Trained on Nemotron and OpenMath data for ~90 ExaFLOPs (~20B tokens):

Table 1: Benchmark Results (1B models)

LoL is competitive with baselines across all benchmarks, with the only notable exception being the full-attention vanilla model scoring substantially lower on ARC-Easy.

6. Comparison to Looped Transformers

Compared against Huginn (Geiping et al., 2026) with matched library/recurrent block sizes:

• LoL is competitive in loss with parameter-matched looped models
• LoL shows slightly higher variance across seeds (attributed to router initialization sensitivity)
• Unlike Huginn, LoL can offer sparse behavior by scaling library size

7. Comparison to Sparse Mixture of Experts

Compared at 3:2 sparsity ratio (MoE: 3 experts per layer, 2 active; LoL: 24 experts over 16 routed depths):

• MoE baseline achieves slightly lower training loss (2.75 vs. 2.8)
• Growing gives LoL an early loss advantage
• Expert utilization heatmaps show some experts specialize by depth while others are used diversely

Theoretical and Practical Implications

Unified Architecture Family

LoL provides a single model family that spans:

• Dense transformers (L ≈ D): minimal reuse, no sparsity
• Looped transformers (small L): maximal reuse, dense computation
• Mixture of experts (large L): minimal reuse, high sparsity

This unification suggests that parameter efficiency and FLOPs efficiency are not separate design axes but endpoints of a continuum controlled by library size.

Interpretability and Control

The routing patterns provide a window into the model's "latent computation":

• Language-specific paths: Different languages use different computational routes
• Strategy encoding: The model learns to implement different algorithms (BFS, DFS, A*) by reordering its own layers
• Safety signal: Router probabilities contain linearly accessible information about text harmfulness

Test-Time Intervention

A key practical contribution is the ability to force routing paths at inference time, enabling:

• Strategy control without retraining
• Verification that specific computational paths correspond to specific behaviors
• A mechanism for studying latent computation in transformer models

Test-Time Scaling Potential

LoL can in principle employ larger computational depth at inference than at training time, offering a mechanism for test-time scaling and increased latent reasoning capacity.

Conclusion

The paper introduces Library of Layers (LoL), an architecture that makes layer ordering a learned, input-dependent computation. By varying only the library size, LoL generalizes the parameter-reuse and sparsity tradeoffs of dense, depth-recurrent, and mixture-of-experts architectures.

Key takeaways:

1. Dynamic layer ordering is feasible and competitive: LoL matches or exceeds strong baselines across diverse benchmarks at 360M and 1B scales.
2. Routing paths are interpretable and controllable: Distinct paths emerge for languages, algorithms, and safety labels, and these paths can be intervened upon at test time.
3. Unified design space: A single hyperparameter (library size) spans the design space from dense to sparse computation.

Future directions identified by the authors:

• Safety interventions: Whether routing features can support interventions in chat settings (lying, scheming, safety behaviors)
• Test-time scaling: Employing larger computational depth at inference than training
• Test-time learning: Combining top-1 routing with test-time learnable adapters to create a growing library of parametric memory available to all layers

The paper positions LoL as a path toward models that dynamically compose their own computational graphs, enabling more efficient and controllable reasoning.