Transformer layers adaptively reordered to solve tasks efficiently
Library of Layers (LoL): Adaptively Ordering Transformer Layers at Runtime
Machine LearningArtificial IntelligenceComputation and Language
Summary
Most language models always process information through their layers in the same order, but this paper introduces a new approach where the order of layers can change depending on the problem. The researchers created a system called Library of Layers (LoL) that picks different layer sequences to solve different tasks, like solving Sudoku or understanding multiple languages. This adaptability lets LoL reuse parts or use fewer computing steps, balancing speed and accuracy. The model can also be guided to apply specific strategies, making it more controllable for different situations.
What this means in practice
- •For software engineers: Build AI systems that adapt their internal computations dynamically to balance speed and accuracy for diverse language tasks and problem-solving.
- •For multilingual ai developers: Create language models that automatically choose different processing strategies based on language or script to improve multilingual understanding.
Authors
Abdellah Rahmani, Sean McLeish, Ang Li, Sanae Lotfi, Pascal Frossard, Micah Goldblum
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.