EPFL and NYU's 'Library of Layers' Reorders Blocks at Runtime

A 360M-parameter transformer learns to pick the order of its own layers, and the researchers can override that order to make it solve a sudoku puzzle a chosen way. The paper from Sean McLeish and co-authors at EPFL, NYU, Columbia and the University of Maryland was posted to Hugging Face this week.

They call it Library of Layers, or LoL. "We introduce Library of Layers (LoL), an architecture that dynamically orders its layers during inference to solve the problem at hand," the authors write. A single hyperparameter, the size of the shared block library, slides the model between dense, depth-recurrent, and mixture-of-experts behavior.

The interventions are the striking part. The team trains LoL on five sudoku-solving strategies, each with its own trace syntax, and extracts a mode routing path through the router for each. Forcing that path at inference produces the requested strategy's trace with close to 100% accuracy. The same holds on 8x8 mazes solved with BFS, DFS, A*, bidirectional search and dead-end filling, where every strategy shares the same output syntax and only the cell exploration order differs.

On safety data, a ridge regression trained on router probabilities over a 50M-token held-out set separates safe from unsafe text with an AUROC of 0.92 at the window level and 0.95 at the sequence level. At 1B parameters on math data, LoL stays "within 2 standard errors" of a vanilla dense baseline on ARC, HellaSwag, PIQA and GSM8K; the vanilla full-attention model is the one that scores substantially lower on ARC-Easy. Code is at github.com/arahmani1/library-of-layers, placing it among the open-source AI research we have tracked at a steady pace this quarter.