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

Library of Layers (LoL). Adaptively Ordering Transformer Layers at Runtime. By Abdellah Rahmani and five others, from EPFL, the University of Maryland, and two other institutions. Published on the second of November.

Programmers build very different algorithms out of the same small set of reusable primitives. A conditional, a loop: the pieces stay the same, and the order in which they are chained is what changes. Language models do not work like that. Existing architectures run their layers in one fixed order, every time, whatever the problem in front of them.

Library of Layers, or LoL, lets the model pick the order. It is an architecture that dynamically orders its layers during inference, choosing a sequence suited to the problem at hand. The layers are drawn from a global library of transformer blocks, and the size of that library is the architecture's single hyperparameter. The analogy to programming is direct. The blocks play the part of reusable primitives, and the order the model assembles them in plays the part of the algorithm. Two different inputs can travel through the same library along two different routes.

That one setting does a lot of work. Turning it moves the model smoothly between reusing parameters and being sparse. Seen that way, LoL generalizes the tradeoffs in parameter efficiency and in FLOPs efficiency that belong to three familiar designs: dense models, depth-recurrent models and mixture of experts. All three become members of one unified family, rather than separate choices an architect has to make up front. Instead of picking between reuse and sparsity, a designer picks a library size, and the position on that spectrum follows from it.

Pretrained on text, LoL is competitive with all three of those baselines. The more distinctive results are about what the ordering reveals. On synthetic sudoku and maze data, the model uses distinct layer orderings to carry out different solution strategies. The authors can intervene on the routing and force the model to apply one specific strategy to a given problem instance. The ordering is not just a side effect, then. It corresponds to a strategy, and changing it changes how the problem gets solved.

The same pattern appears outside puzzles. In a multilingual natural language setting, LoL chooses different layer order patterns depending on the language and the script. And when LoL is trained on safety data, probes fitted to its router outputs separate safe text from unsafe text with high AUROC. In other words, the route a model takes through its library carries readable information about the input.

The authors frame LoL as a step toward models that compose their own computational graph on the fly. The appeal is twofold. Efficiency, because one library can stand in for dense, recurrent and sparse designs. And control, because a routing choice that can be read and overridden is a handle on how the model reasons.

Library of Layers (LoL). Adaptively Ordering Transformer Layers at Runtime. By Abdellah Rahmani and five others, from EPFL, the University of Maryland, and two other institutions. Published on the second of November. This is the long version: the whole paper rather than the abstract.

A programmer builds very different algorithms out of the same few primitives. A database engine filters, joins, sorts and aggregates, in whatever order the query needs. A transformer language model cannot do that. Every input passes through the same layers in the same fixed order. The authors argue this rigidity may cost efficiency, because a useful computational step may have to be represented again and again at different depths. Mixture of experts, looped models and layer skipping all loosen the architecture, but all still fix the order in which layers run. Library of Layers, from EPFL, the University of Maryland, New York University and Columbia University, makes that order something the model chooses at runtime.

The model has three parts. A prelude of ordinary layers at the start and a coda at the end, both in fixed order with full causal attention. Between them sits the routed core. The core owns no layers of its own. It draws from a shared library of transformer blocks, each called an expert, and at every routed depth a small router picks one block to apply. Library size counts the distinct blocks stored; core depth counts the block applications executed. That single dial is the claim to generality. A library smaller than the depth forces reuse, which behaves like a looped transformer. A library much larger than the depth leaves most blocks unused on any path, which behaves like a sparse mixture of experts.

The fix is to route windows rather than tokens. At each routed depth the sequence is cut into windows, every token in a window uses the same block, and causal attention is restricted to that window. Queries and keys therefore always come from the same block. For natural language the windows are a hundred and twenty-eight tokens long. Then the boundaries move. At each successive depth they shift by a stride, sixty-four tokens for text, so a window at one depth straddles the seam between two windows at the depth before. Information crosses the seam, and the field of view widens with depth, much as it does under sliding window attention.

The router has to respect causality too. For any window after the first, it sees only earlier windows: it averages their mean hidden states, and a single linear layer turns that average into a choice. The first window has no history, so its router input is the mean of the first token's state and a learned vector specific to that depth. Selection is top one. The chosen block's update is scaled by its router probability, which is what lets gradient reach the router at all. Each depth also gets its own learned scale and bias after every normalization, so one block can behave a little differently depending on where it runs. Last, cost. Training spend is driven mainly by core depth, so the model starts shallow and grows deeper during training, which the appendix shows saves compute for only a minor increase in final loss.

The sharpest test of the thesis is not a loss curve. It is whether the chosen order actually carries the algorithm. On mazes, five strategies, among them breadth-first search, depth-first search and dead-end filling, all write the same output: a list of visited cells. They differ only in the order of visits. The strategy tag is shown to the router and masked from the rest of the network. After training, the authors extract each strategy's mode path, swap the tag for neutral filler, and force that path on new mazes. The model produces the requested algorithm's trace with close to one hundred percent accuracy, the weakest strategy still at ninety-nine point eight percent. Because the output syntax is identical across strategies, the ordering of blocks is what decides which algorithm runs. Sudoku tells the same story, at ninety-nine point five percent or better.

Routes also carry meaning nobody supervised. Trained on ten languages, paths cluster by language, and the Latin-script languages share similar paths. Trained on safe and unsafe web text without labels, a ridge probe on router probabilities separates the two with a test AUROC of zero point nine two per window and zero point nine five per sequence, while shuffled labels fall to chance. At one billion parameters on math data, LoL matches dense baselines in training loss and on five benchmarks. It is competitive with Huginn looped models, and trails a mixture of experts baseline slightly, two point eight against two point seven five in training loss.

There is no limitations section; the discussion calls the work an initial exploration. Several limits are visible in what is reported. LoL's final losses vary more across three seeds than the looped baselines, which the authors hypothesise comes from router initialization steering which paths get explored. The mixture of experts baseline ends lower in training loss, and the reply, that higher loss often coexists with strong downstream performance in latent reasoning models, is a citation rather than a downstream test of these models. The safety probe measures agreement with labels from another classifier, which the authors call noisy, and it detects rather than intervenes: steering behaviour in chat settings is left open. The clean control results come from synthetic tasks with scripted algorithms, and the largest model has one billion parameters.

The idea to carry away is that layer order can be learned rather than fixed. Keep a library of blocks, let a router pick the next one for each window, and one dial, library size, slides between looped reuse, dense depth and expert sparsity. The bonus is visibility. The path a window took is a readable record of its computation, and forcing a path is a lever on it. On mazes, that lever chooses the algorithm.