Nested Sparsity for Scalable Pretraining and Continual Learning
AuthorsTianzhe Chu, Druv Pai, Yi Ma
AffiliationsUniversity of Hong Kong · UC Berkeley
Nestor combines nested attention and MoE sparsity to make language-model pretraining and continual adaptation more efficient.
Key results
The MoE baseline requires this many times NESTOR’s compute to reach the same loss across tested scales.
Active parameters; trained on DCLM-baseline.
Average reduction versus the MoE baseline across four domains without replay.
Nats lower than the MoE baseline for the 500M model at equal compute.
What the paper found
Mixture-of-experts models, including DeepSeek-V3 and OpenAI’s gpt-oss, make computation conditional by routing tokens to selected MLP experts, while standard attention remains shared. NESTOR extends routing to attention: a token selects columns that have their own attention modules and experts, nesting attention-weight, attention-map, and expert sparsity. Because each column attends only to tokens routed to it, the same decision selects both the attention parameters and the context; routing is held constant across a block of layers. On DCLM-baseline pretraining from 100M to 1B active parameters, NESTOR reaches the same loss as an MoE using 1.13–1.29× less compute, with the advantage persisting across scale. In continued training on books, mathematics, code, and German, it achieves similar adaptation gains while forgetting less: on average, 3% less than the MoE without replay. For long-context extension, the attention savings become more valuable as sequences grow: at 64k context, the 500M NESTOR model beats its MoE counterpart by 0.061 nats at equal compute. The results suggest that conditional computation in attention can complement MoE routing, improving pretraining efficiency and reducing interference during adaptation.
Original abstract
Mixture-of-experts (MoE) transformers, the standard backbone for modern large-scale deep learning, employ conditional computation: only a sparse subset of model weights is used in each forward pass. In MoE transformers, this sparsity arises in the MoE operator, which generalizes the multi-layer perceptron in the classical transformer. In this work, we explore an alternative mechanism for conditional computation that analogously generalizes the attention operator. The sparsity induced by this mechanism naturally nests (i.e., hierarchically composes) with the sparsity induced by the MoE operator. We show that a transformer employing this nested sparsity mechanism, which we call Nestor, performs better than a conventional MoE transformer at equal compute and equal sparsity in language model pretraining. We also show that Nestor is well-suited for continual pretraining: it can be adapted to longer contexts more effectively than conventional MoEs, as well as adapt more efficiently to new domains while forgetting less on old domains. Furthermore, we demonstrate that these benefits persist as model size, dataset size, and context length grow. Our results demonstrate the potential of alternative conditional computation mechanisms for deep neural network architectures. Code and checkpoints are available on GitHub at https://github.com/Ma-Lab-Berkeley/Nestor .
Read the original paperMore in Efficient AI
Browse all 70 papers →AB-Attention: Accurate Block Attention for Accelerating Video Diffusion Transformers
Bin Xu, Mo Zhang, Zhengfeiyang Zhang, Xinfeng Yang, Yixing Ma, Wei Wang, HeXin Wang, Pengfei Zhou, Wangbo Zhao
AB-Attention speeds up video generation by skipping most attention computations while more accurately estimating the attention it omits.
Extreme Compression of Per-Layer Embeddings via Cross-Layer Latent Vector Quantization
Jiacheng Liu, Yaxin Luo, Xiaohan Zhao, Peng Sun, Xinyi Shang, Jiacheng Cui, Yonina C. Eldar, Zhiqiang Shen
A new way to compress Transformer embeddings dramatically reduces memory use while largely preserving model performance.
Rethinking Sequence Compression with Test-time Regression
Sihan Yang, Jie Huang, Hanxuan Li, Yifan Guo, Weigao Sun, Yiran Zhong, Weiyang Liu
A new regression-based compressor aims to keep language and vision models efficient while preserving more useful information in fewer tokens.