Search papers, labs, and topics across Lattice.
This paper introduces Mobius Learning, a novel training architecture for Transformers that employs cyclic depth folding, allowing block groups to be optimized in both shallow and deep roles across different data streams. By shifting the order in which blocks are applied, the model achieves depth-role superposition, leading to improved performance. In experiments with a modified GPT-2 model, Mobius Learning demonstrated lower validation loss compared to traditional fixed-order architectures, highlighting its potential for more efficient training.
Cyclic depth folding enables Transformer blocks to excel in both shallow and deep roles, achieving better performance than traditional fixed-order training.
Transformer-based language models organize computation along an ordered depth axis, where shallow and deep blocks often develop distinct representational roles. We challenge the conventional view that these roles must remain tied to a block's position in the ordered sequence. We introduce Mobius Learning, a training architecture based on cyclic depth folding, in which different data streams follow cyclically shifted block orders. The same block group is therefore applied early in the block sequence for some data streams and late for others, so it is optimized in both shallow and deep roles, a phenomenon we call depth-role superposition. Surprisingly, in four-worker experiments with a modded GPT-2 small (124M) model trained on 2.5B FineWeb tokens using Muon, Mobius Learning achieves lower validation loss than a fixed-order looped Transformer at larger numbers of Transformer block-sequence passes. This counterintuitive result shows that a block group need not remain confined to one fixed shallow or deep role within the block sequence and opens a new design space based on cyclic depth folding. Crucially, this structure makes Mobius Learning particularly well suited to memory-constrained distributed training: raw training data remain local, while each worker stores one block group rather than the complete Transformer block stack.
