Mobius Learning: Cyclic Depth Folding in Transformers

2026-07-20

AI Summary

This work addresses the limitation in conventional Transformers where shallow and deep layer roles are rigidly tied to fixed positional assignments, constraining model expressivity and training efficiency. To overcome this, the authors propose Möbius Learning, a novel architecture employing a cyclic depth-folding mechanism that enables a shared set of modules to process distinct data streams in a rotating order during training. This design allows individual modules to simultaneously assume both shallow and deep functional roles, thereby breaking the strict positional binding. Integrated with module role superposition, distributed memory optimization, and the Muon optimizer, the approach achieves lower validation loss than fixed-order recurrent Transformers under the same computational budget—demonstrated on 2.5B FineWeb tokens across four nodes—offering a new paradigm for efficient training in memory-constrained settings.

Abstract

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.

Problem

• Transformer
• depth-role superposition
• cyclic depth folding
• distributed training
• memory-constrained

Innovation

• Mobius Learning
• cyclic depth folding
• depth-role superposition
• distributed training
• Transformer architecture