Nested Sparsity for Scalable Pretraining and Continual Learning

Nested Sparsity for Scalable Pretraining and Continual Learning. By Tianzhe Chu and two others, from the University of Hong Kong and UC Berkeley. Published on the twenty-second of October. Its code repository is worked out: about one hundred and eighty source files, with a readme.

Mixture-of-experts transformers are the standard backbone for modern large-scale deep learning, and they rely on conditional computation: each forward pass uses only a sparse subset of the model's weights. In these models the sparsity arises in one particular place, the MoE operator, which generalizes the multi-layer perceptron of the classical transformer. That leaves an open question about the other half of the block. Could the attention operator be generalized in the same spirit?

The paper explores exactly that: an alternative mechanism for conditional computation that generalizes the attention operator, by analogy with what MoE does for the perceptron. The sparsity this mechanism induces nests with the sparsity of the MoE operator. Nesting here means the two compose hierarchically, one inside the other, rather than sitting side by side. A transformer built on this nested sparsity is called Nestor.

The abstract does not spell out the internals of the new operator, so neither will this. What it does fix is the terms of the comparison. Nestor is measured against a conventional MoE transformer at equal compute and at equal sparsity. Under those conditions, a difference in results cannot come from using more computation or more active weights. It has to come from where the sparsity sits.

In language model pretraining under those conditions, Nestor performs better than the conventional MoE transformer. It is also well suited to continual pretraining, in two ways. It can be adapted to longer contexts more effectively than conventional MoEs. And it adapts more efficiently to new domains, while forgetting less on the old ones. These benefits persist as model size, dataset size and context length grow. Code and checkpoints are available on GitHub.

Sparsity in the standard design lives in the perceptron. Nestor extends it into attention as well, nested inside the same model, and gains on pretraining, context extension and domain adaptation without spending more compute. For the authors, this points to alternative conditional computation mechanisms as a direction for deep neural network architectures in their own right.

Nested Sparsity for Scalable Pretraining and Continual Learning. By Tianzhe Chu and two others, from the University of Hong Kong and UC Berkeley. Published on the twenty-second of October. This is the long version: the whole paper rather than the abstract.

Mixture-of-experts transformers are now the standard backbone for large language models, and their trick is conditional computation: each token uses only a small, input-dependent subset of the weights. But in every one of these models the trick stops at the feed-forward layer. Every token is processed by the same attention parameters in every layer. That matters, because attention accounts for a considerable fraction of the active parameters a token uses, and at long context lengths it accounts for most of the computation. The question the authors ask is whether conditional computation can be applied to attention as well, and how it should interact with the experts.

Start with the operator being generalized. A mixture-of-experts layer replaces one feed-forward network with many expert networks. A learned router scores each token, the top few experts are selected, and their outputs are combined with weights from a softmax over those scores. Stored parameters grow with the number of experts; compute per token grows only with the number selected. Attention has been made conditional before, in two separate ways. Sparse attention restricts which tokens each token attends to, but every token still shares the same attention parameters. Mixtures of attention heads route each token to a subset of the parameters, but every token still attends to the whole context.

The architecture is called Nestor. Its attention is replaced by several attention operators with separate parameters, which the authors call columns. A column router picks a few columns for each token. A column then computes ordinary causal attention, but only among the tokens routed to it, so one routing decision fixes both the weights a token uses and the tokens it can see. Each column also owns its own set of experts and its own expert router, so an expert is indexed by a pair: which column, then which expert inside it. That is the nesting. The coarse choice of column comes first, and the fine choice of expert happens within it.

Layers are grouped into blocks, and the column routing is computed once per block, from the block's input. A token stays in the same columns for every layer of the block, so within a block each column behaves like a small mixture-of-experts transformer running on its own subsequence. Set the number of columns to one and the ordinary mixture-of-experts transformer comes back exactly; set the experts to one as well and the dense transformer comes back. In all experiments there are eight columns, and the number active rises with depth, from three in the first block to all eight in the last, which computes full attention over the whole sequence. Both routers are balanced without an auxiliary loss. Heads per column are set so that projection costs stay close to the baselines, which means the saving comes only from the attention products themselves: a token in a sparse block pays for attention over just the fraction of the sequence its columns see.

Pretraining compares three architectures at four scales, from one hundred million to one billion active parameters, with Nestor using no more training compute than either baseline and storing about as many parameters as the mixture-of-experts model. Nestor has the lowest held-out loss at every scale. From fitted scaling laws, the mixture-of-experts transformer needs between one point one three and one point two nine times the compute to match it, and the dense transformer between three point one and four point one times. The factor does not shrink with scale.

The comparison that most directly tests the thesis is an ablation in the appendix, at the two hundred million scale. Smaller, more numerous experts alone change nothing. Nestor's expert design with a single column matches the baseline. Column-routed attention with the baseline's experts tied across columns lands within the variation between seeds. Only the full model, with experts nested inside columns, is lower: by about two hundredths of a nat on each of two seeds, and by a little more at five hundred million. Neither half carries the gain; the composition does. In continual pretraining on books, mathematics, code and German, Nestor gains within one percent of the baseline while forgetting three percent less, and the dense model forgets twenty-one percent more. Under context extension to sixty-four thousand tokens, Nestor's cost per token falls to roughly four fifths of the baseline's, and at equal compute its advantage grows with length.

The conclusion names two limits. The largest model has one billion active parameters, and every model is judged by held-out loss, so whether the multiplier holds at larger scales or carries over to downstream benchmarks after post-training is untested. And only columns with full causal attention were studied. Beyond that, each domain adaptation run was done once, with one seed. The seed-to-seed difference at two hundred million is about one hundredth of a nat, so the ablation margins are only two or so seed gaps. The column schedule was chosen at two hundred million and reused at other scales without further search. Sliding window baselines, run with a single seed, are comparable to Nestor at some scales, though their efficiency relative to it declines as scale grows. And counting only linear layers, Nestor uses four to eight percent more compute than the baselines.

The idea to carry away is that one routing decision can choose three things at once: which attention weights a token uses, which other tokens it reads, and which family of experts it can reach. The ablation is the part to remember, because it says the gain comes not from routing attention, nor from finer experts, but from putting the experts inside the attention route.