Fully Nested Transformers
Abstract
Matryoshka Representation Learning produces representations that can be truncated at different granularities to adapt to diverse downstream requirements, but computing them still requires a full forward pass through the model. Elastic architectures enable adaptive inference budgets, but their smaller models do not generally produce representations that are prefixes of those produced by the larger models. To unify these two forms of adaptivity within a single language model, we design StairFormer: a Transformer architecture that nests a hierarchy of prefix submodels by preserving prefix structure at each layer. Running a larger model produces, as prefixes, the same representations that would have been produced by independently running each smaller model. Prefix representations are efficiently computed using the computational paths of smaller submodels, which can then be progressively refined by activating larger portions of the model. This full nesting enables cascading inference and reuse of intermediate computations across model scales because submodels share the same forward pass. Empirically, we show that StairFormer maintains competitive language modeling performance at 1.13B parameters, achieving a CORE metric within an 11.1\% relative gap of a standard Transformer baseline while satisfying these full nesting constraints.