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

AI summary

Overview

Research area: Machine learning / Transformer architecture design, specifically input-adaptive computation and conditional layer routing.

Technical level: Intermediate. The paper is readable with a basic understanding of transformer blocks and mixture-of-experts routing, but the windowed attention and routing scheme require some architectural familiarity.

Scope (one sentence): The paper introduces and evaluates Library of Layers (LoL), a transformer architecture whose layers are selected and ordered at runtime from a shared library by a learned router, and studies both its competitiveness against dense, looped, and mixture-of-experts baselines and the interpretability of the resulting routing paths.

What This Paper Is About

Standard transformer language models execute their layers in one fixed order for every input, even though the same reusable computational pieces could in principle be recombined differently for different problems. LoL replaces that fixed order with a small router that picks, at each step of the forward pass, which transformer block to apply next from a global library of blocks. The goal is to make layer selection, ordering, and reuse depend on the input, while keeping performance competitive with existing efficient architectures.

Key Contributions

1. A runtime-ordered architecture. The authors propose Library of Layers, in which a router selects blocks from a global library of transformer layers, making layer selection, ordering, and the number of times each layer is used conditional on the input. The library size is the single hyperparameter that moves the model between parameter reuse and sparsity.

2. Strided window attention for stable and causal routing. The paper introduces a windowed attention scheme that lets different windows of tokens follow different routing paths while preventing attention-key/value mismatches across differently routed positions, preserving causality, and allowing information to flow between neighboring windows as depth increases.

3. Empirical evidence that routing paths encode interpretable strategies. Using maze, sudoku, multilingual, and safety data, the authors show the router learns input-adaptive patterns, that these patterns can be used to control the model at test time by forcing a layer execution order, and that probes on router outputs distinguish safe from unsafe text.

4. A unified comparison across sparsity regimes. LoL is benchmarked against looped transformers (small library) and sparse mixture-of-experts models (large library), and at 1B parameters against dense transformers, showing competitive performance across the spectrum.

Main Findings

• Competitive with dense baselines at 1B parameters. On a mixture of Nemotron and OpenMath data, LoL models (16 layers: 4 prelude, 4 coda, routed core depth 8) trained for approximately 90 ExaFLOPs (roughly 20B tokens for the most efficient model) match dense baselines in training loss. On ARC-Easy, ARC-Challenge, HellaSwag, PIQA, and GSM8K, LoL scores 50.3 ± 1.0, 29.1 ± 1.3, 31.2 ± 0.5, 59.7 ± 1.1, and 27.9 ± 1.2 respectively; nearly all evaluations stay within 2 standard errors of the alternatives, with the exception of the full-attention vanilla model, which scores substantially lower on ARC-Easy (40.8 ± 1.0).

• Competitive with looped transformers. With the library size fixed to match the recurrent block size of Huginn-style looped models and 10B tokens of FineWeb, LoL is competitive in loss, though it shows slightly higher variance across three seeds, which the authors attribute to router initialization leading to different explored paths.

• Competitive with sparse mixture of experts. At a 3:2 sparsity comparison (MoE with three experts per layer and two active; LoL with 24 experts over 16 routed depths) trained for approximately 17 ExaFLOPs (about 10B tokens for the most efficient model), the MoE baseline achieves a slightly lower training loss, a gap of approximately 0.05 (2.75 versus 2.8).

• Distinct routing per language. A 360M model trained on an equal mixture of 10 languages from FineWeb-2 for 10B tokens develops routing paths that are reasonably distinct per language, with Latin-script languages sharing similar paths; the paper reports that the same layers are used across languages but executed in different orders.

• Routing features encode safety labels. Training on an equal mixture of safe (score-0) and unsafe (score-4 and score-5) FineWeb sequences for 10B tokens, a ridge regression on router probabilities reaches a held-out test AUROC of 0.92, and a sequence-level probe averaging router probabilities across windows reaches 0.95. Development-set cross-validation results include 0.8754 for top-1 ridge, 0.8244 for single winning-expert probability, 0.8300 for expert-frequency log-likelihood ratio, and 0.5014 mean / 0.5231 max for shuffled-label controls over 99 permutations.

• Forced routing paths control sudoku strategy. Using five sudoku solving strategies, each with its own step syntax and an [Sx] identifier masked from all parts of the network except the router, the model solves with the expected strategy at close to 100% accuracy when the strategy's mode path is forced at inference.

• Forced routing paths control maze exploration order. On 8×8 loop-free mazes with BFS, DFS, A*, bidirectional search, and dead-end filling, all five strategies share the same output syntax and differ only in cell visitation order; forcing a strategy's path still produces the requested algorithm's trace with nearly 100% accuracy.

• Growing depth saves compute. Starting training with a shallow routed core and increasing depth over training (using a 1-sqrt schedule for 75% of training in the 1B experiments) gives LoL an early loss advantage over fixed-depth dense baselines and saves FLOPs with only a minor increase in final loss.

• Expert usage is depth-structured. Heatmaps of expert utilization show experts are often used in groups of neighboring depths, while some experts are used across a wider range of depths.

Methodology in Plain English

The model is split into three parts: a fixed "prelude" of standard layers at the start, a fixed "coda" at the end, and a "core" in the middle where routing happens. The core has a library of distinct transformer blocks (called experts) and a number of applications (routed depths). At each routed depth, a small linear router picks exactly one expert per window of tokens, and the chosen expert's output is blended in with a weight equal to its router probability.

The key design problem is causality plus stability. If every token routed independently, neighboring tokens could be processed by experts with different attention projections, and the authors found empirically that this leads to training instability and poor performance. If instead one ordering were chosen for the whole sequence, the model could only use the prompt to decide, or it would break causality. The solution is to split the sequence into windows of W tokens, route each window as a unit, and shift window boundaries by S positions at each depth. Successive depths therefore overlap, letting information pass between neighboring windows in a way analogous to the growing field of view of sliding window attention.

The router itself may only look at earlier windows: for window 0 it uses a learned depth-specific vector combined with the first token's hidden state, and for later windows it averages the mean hidden states of all preceding windows. Training uses loss-free load balancing, Gaussian noise on router logits to encourage exploration, and FiLM-style depth-dependent scale and bias after each normalization.

To interpret the model, the authors extract a "mode path" per task strategy — the highest-mean-probability expert at each window and depth, with residual weights fixed to those means — and then force that path at inference while removing the strategy hint from the input.

Why This Matters

The paper reframes layer order as something a model can learn rather than something an architect fixes, and it shows this yields a practical knob: changing the library size moves the model continuously between heavy parameter reuse (looped-style) and heavy sparsity (MoE-style). Because the execution path is explicit and intervenable, the architecture doubles as a handle for interpretability and control.

Real-world applications:

• Steering reasoning style at inference time (for example, choosing between search strategies) without retraining separate models.
• Safety monitoring, using router probabilities as a cheap internal signal that distinguishes safe from unsafe text (AUROC of 0.92 to 0.95 in the reported setting).
• Multilingual serving, where one set of parameters serves 10 languages through different execution orders.
• Compute-adaptive deployment, where the same architecture family spans dense and sparse operating points depending on hardware or latency budgets.

Industry relevance: Serving one parameter pool across many behaviors is attractive for memory-constrained deployment, and the reported competitiveness against MoE at matched active parameters makes the approach directly comparable to the dominant sparse-scaling paradigm.

Future Directions

• Extending routing-based control to chat behaviors. The authors ask whether routing features can support interventions in chat settings, including behaviors related to lying, scheming, and safety.
• Test-time depth scaling. LoL can in principle use a larger computational depth at inference than at training, offering a mechanism for test-time scaling and greater capacity for latent reasoning.
• Combining routing with test-time learning. Instead of routing over full transformer blocks, the router could select among a growing library of test-time learnable adapters, making parametric memory available to all layers.
• Improving on the initial design. The authors describe the paper as an initial exploration and expect future work to substantially improve the architecture and training scheme.

Target Audience

Machine learning researchers and engineers working on efficient transformer architectures, conditional computation, mixture-of-experts systems, and mechanistic interpretability. It is also relevant to practitioners who need one model to cover multiple languages or behaviors, and to readers interested in controlling model reasoning strategies at inference time rather than only through prompting.

Authors’ 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.