Nested Sparsity for Scalable Pretraining and Continual Learning

• MoE conditional computation sparsifies MLP weights, but every token still uses the same attention parameters, which account for substantial active parameters and long-context computation.Source: §1 · pp. 1–2
• Continual pretraining on new domains or longer contexts must retain performance on old data, which neural networks tend to lose on a new distribution.Source: §1 · p. 2
• NESTOR hierarchically combines attention and expert sparsity, generalizing MoE and dense transformers across 100M–1B active parameters.Its three sparsity axes are attention weight, attention map, and expert sparsity.Source: §1 · p. 3
• 1.13–1.29× as compute-efficient as the MoE transformer and 3.1–4.1× as compute-efficient as a dense transformer, with neither factor shrinking as scale grows.Source: §1 · p. 3
• NESTOR matches MoE adaptation gains across specialized domains while forgetting less, and preserves or grows its advantage when extending context from 4k to as long as 64k.At equal continual-pretraining FLOPs, its advantage over MoE grows with context length.Source: §1 · p. 3
• The paper’s hierarchical-latent-structure model reproduces NESTOR’s pretraining compute efficiency, its growth with context length, and reduced forgetting.Source: §1 · p. 3
• The results suggest conditional computation need not be confined to the MLP, and composing new sparsity mechanisms is a promising architecture-design direction.Source: §1 · p. 3

• A column router chooses attention operators for each token and restricts attention to tokens routed to the same column, sparsifying both parameters and attention pattern.Source: §2 · pp. 3–4
• Each active column has its own expert router, so column selection also determines which experts and multi-layer module a token uses.Source: §2 · p. 4
• Routing is computed once per block, keeping each token in the same columns across M layers; setting C = 1 recovers MoE, and also setting E = 1 recovers dense transformers.Source: §2 · p. 4
• NESTOR is presented as a meta-architecture that can compose parallel deep networks, although the experiments use full causal attention within each column.Source: §2 · pp. 3–4
• Unlike sparse-attention methods that route attention separately in every layer, NESTOR’s token-level choice selects both the attention context and the module spanning M layers.Source: §2 · p. 4
• Experiments increase active columns with depth, from Kc = 3 to all C = 8 in the final block, and use learned attention residuals and balanced top-K routing.Source: §2 · p. 5

• The comparison spans 100M, 200M, 500M, and 1B active parameters, with NESTOR using no more training FLOPs than either baseline and approximately matching MoE stored parameters.Models train on DCLM-baseline at context length 4096 for 50 tokens per active parameter.Source: §3 · p. 5
• Under the common approximation that counts only linear-layer FLOPs, NESTOR uses 4–8% more FLOPs than the baselines at context length 4096.Source: §3 · p. 5
• 1.22× as much compute for the MoE transformer and 4.1× for the dense transformer to match NESTOR’s loss at the largest scale.Across all four scales, NESTOR achieves lower held-out loss than both baselines; its fitted exponent matches MoE’s to three digits, with a smaller coefficient and irreducible loss.Source: §3 · pp. 5–6, Figure 2 · p. 6

4.1 Domain Adaptation

• Domain adaptation is tested on books, mathematics, code, and German without replay, measuring target-corpus loss reduction and pretraining-corpus loss increase.The models first undergo context extension to 32k, then train on 0.5B tokens per target domain.Source: §4.1 · p. 6
• NESTOR’s average gains across four target domains are within 1% of the 500M MoE’s, while it forgets 3% less on average.With 25% replay of pretraining data, NESTOR forgets 5% less than MoE on average, with gains within 2%.Source: §4.1 · pp. 6–7, Figure 3 · p. 7

4.2 Long Context Extension

• Context extension is evaluated from 4k to 8k, 16k, 32k, and 64k using long documents, with NESTOR tested at equal tokens and equal FLOPs.Source: §4.2 · p. 7
• At equal tokens, NESTOR achieves context-extension gains comparable to MoE and broadly preserves its pre-extension advantage; at equal FLOPs, its advantage grows with context length.Lower FLOPs per token allow more training; spending savings on fresh data yields the largest improvements.Source: §4.2 · p. 7

5 Theoretical Analysis

• The model assumes coarse labels determine which context tokens matter, their weighting pattern, and the computed statistic, while fine labels determine how that statistic is used.For each coarse label, fine-label maps share a pattern and row space, which may differ across coarse labels.Source: §5 · p. 8
• As T grows, both predictors’ excess risks decay at the same rate, σ2d/T, while the nested predictor uses fewer dimensions when r ≪ D and E ≫ 1.The dimension count includes patterns, row spaces, and maps; the nested predictor shares a pattern and row space across the experts of each coarse label.Source: §5 · pp. 9–10
• Both predictors can represent the target, and NESTOR’s routing fixes cross-label pattern blocks to zero rather than estimating them from data.The comparison is a stylized one-layer model with one head, linear attention and maps, and squared loss.Source: §5 · pp. 8–10
• The flat predictor needs (S0 + S −1)/(S0 + (S −1)/C) times as many FLOPs as the nested predictor for equal expected excess risk, with the factor growing from dF/dN to CdF/dN as S grows.The multiplier combines the cost of estimating pattern blocks and row spaces with the cost of scoring context tokens that are irrelevant to a given coarse label.Source: §5 · pp. 9–10
• Under continual training on a new coarse label, NESTOR’s other columns never change, so it has no forgetting on other coarse labels.In the flat predictor, forgetting on other labels comes from changes to its shared value map and grows with overlap between old and new row spaces.Source: §5 · p. 10
• The theoretical efficiency and forgetting results depend on how relevant tokens, patterns, and row spaces vary across coarse labels.The analysis states that nesting improves sample and compute efficiency when those structures differ between labels.Source: §5 · p. 10

6 Conclusion

• At 100M–1B active parameters, the MoE transformer needs 1.13–1.29× NESTOR’s compute to reach the same loss at equal stored parameters, with no downward trend as scale grows.The conclusion also reports that NESTOR adapts to new domains as well as MoE while forgetting less, and its long-context advantage grows with context length at equal FLOPs.Source: §6 · p. 10
• NESTOR’s router sends each token to several columns, each with its own attention operators and experts, combining their sparsity in one architecture.The conclusion describes NESTOR as generalizing both MoE and dense transformers.Source: §6 · p. 10
• NESTOR’s results are evaluated only by held-out loss, and its largest model has 1B active parameters; downstream benchmarks and larger scales remain to be tested.The authors also identify sparse or linear attention within columns as a future direction.Source: §6 · p. 10

A Related Work

• Unlike hierarchical mixtures that route MLPs at both levels, NESTOR’s first routing level selects an attention operator and context, fixed over several layers.NESTOR applies token-dependent modules at every layer, rather than routing some tokens around a layer’s module.Source: A Related Work · p. 15
• Column-routed attention sparsifies both attention parameters and attention patterns through one routing decision shared with experts and held fixed over a block of layers.The related work contrasts this with sparse attention methods that sparsify the attention map or attention parameters separately.Source: A Related Work · p. 15
• Unlike fixed sliding windows, NESTOR’s routed subsequence can contain tokens from anywhere in the sequence, and its active columns increase with depth until the last block uses full attention.The related work distinguishes this from interleaved global and local attention layers.Source: A Related Work · pp. 15–16
• NESTOR applies to each token only its column’s attention operators and experts, plus shared experts, without task labels or changes to training.The related-work discussion contrasts this with continual-learning methods that add or freeze task-specific parameters, and with shared MoE attention parameters that can change under new-domain gradients.Source: A Related Work · p. 16

B Mathematical Description of the Architecture

• The column router selects columns and weights once per block, and the same selections and weights are used across all M layers; the last block activates every column.The router’s balancing bias is updated based on column load, without an auxiliary balancing loss.Source: B Mathematical Description of the… · pp. 16–17
• Each column’s attention and column-specific expert updates affect only that column’s stream, while shared-expert updates affect every stream.All columns receive inputs derived from the mean of the streams, and the block output is their final mean.Source: B Mathematical Description of the… · p. 17
• With balanced routing, assigning each column H/Kc heads keeps projection FLOPs matched while reducing attention’s quadratic term; the saving grows with context length.Each column gathers its routed tokens, applies causal attention to that shorter sequence, then scatters the result back.Source: B Mathematical Description of the… · p. 18
• Each layer has shared experts and each column has its own experts; their widths are set so each token uses total expert hidden width 4D per block.The experiments use eight shared experts with top-2 routing, and column experts use top-2 where width permits, otherwise top-1.Source: B Mathematical Description of the… · p. 18
• The output head mixes the embedding, block outputs, and per-column deviations from each block’s mean before predicting the next token.The deviations sum to zero and are used because normalized streams are close to their mean; the loss is next-token cross-entropy.Source: B Mathematical Description of the… · pp. 18–19
• The architecture generalizes the MoE baseline when there is one column, while the dense baseline further uses a single SwiGLU MLP per layer.Both baselines retain the same block attention residual described for NESTOR.Source: B Mathematical Description of the… · p. 19

Sizes

• Across four scales, NESTOR uses seven blocks with Kc = 3, 4, 5, 6, 6, 7, 8, while block depth grows with scale and experts per column increase at 500M and 1B to match stored parameters.Baselines share width, depth, head width, and training budget; the MoE uses 24 experts with 3 active.Source: §C.1 · p. 19, Table 1 · p. 19

Derived quantities

• Only the column count and its depth schedule are design choices; other FLOP-affecting quantities follow rules that keep active MLP width at 4D and attention projections within 10% of baseline cost.The chosen layer splits use 1–2% fewer FLOPs than the budget.Source: §C.1 · pp. 19–20

The schedule of active columns

• Gradual schedules generally achieve lower held-out loss than constant schedules, and most varying schedules outperform every constant schedule, including full attention.The sparse-to-dense step is the exception among varying schedules.Source: §C.1 · p. 20, Figure 6 · p. 21
• With layer composition, FLOPs, and stored parameters fixed, dense-to-sparse ordering beats its reverse, while a dense middle beats dense ends; long initial sparsity performs worst among varying schedules.Gradual schedules also beat two-level schedules at similar mean Kc.Source: §C.1 · p. 20, Figure 6 · p. 21
• At 200M, NESTOR’s ramp and the symmetric schedule with a dense middle have the lowest losses, with no clear difference across two seeds; at 500M, the ramp wins in a single-seed comparison.The symmetric schedule uses fewer attention FLOPs at long contexts, making it an alternative there.Source: §C.1 · p. 20

Pretraining results

• Table 3 reports training tokens, compute, stored parameters, and held-out loss for pretraining runs at four model scales.Its final column estimates baseline compute needed to reach NESTOR’s held-out loss, normalized by NESTOR’s training compute.Source: §C.2 · p. 21, Table 3 · p. 22

Pretraining comparison

• The three sliding-window MoE baselines vary how attention FLOPs are distributed over depth, with one matched to NESTOR’s depth and approximate compute and stored parameters.The other baselines follow MoE depth and expert configurations, distributing windowed attention either near the top or uniformly over depth.Source: §C.2 · p. 21

Sliding window attention and NESTOR

• Sliding-window baselines are more compute-efficient than the MoE and at some scales comparable to NESTOR, but their compute efficiency relative to NESTOR decreases with scale.Fitted scaling laws give the sliding-window baselines a higher irreducible loss than NESTOR.Source: §C.2 · pp. 21–22

C.3 Architecture Ablations

• The full nested model beats the baseline by 0.016 and 0.025 nats on two 200M seeds and by 0.035 nats at 500M, while column-routed attention alone matches baseline loss.Doubling the baseline’s expert count while halving expert width does not explain the gain; one-column variants also match the MoE transformer.Source: §C.3 · p. 22

C.4 Domain Adaptation Settings

• With replay, NESTOR and the MoE transformer have equal mean domain gains of 0.687, while NESTOR has lower mean forgetting, 0.073 versus 0.076, and forgets less on all four corpora.Replay reduces forgetting by about 70% for each model while reducing adaptation gain by about 6%.Source: §C.4 · p. 26

C.5 Long Context Extension Settings

• At equal FLOPs, NESTOR’s margin over the MoE transformer increases with context length, reaching −0.047 to −0.061 nats with new documents and −0.023 to −0.047 with repeated documents at 64k.Across 100M, 200M, 500M, and 1B scales, equal-token margins remain similar to pretraining; at 32k and 64k, new-document losses are 0.004 to 0.025 nats lower than repeated-document losses.Source: §C.5 · p. 29
• At equal tokens, NESTOR’s margin over the MoE transformer remains similar to pretraining at every context length, ranging from −0.012 to −0.035 nats across scales.The exact ranges are −0.028 to −0.030 at 100M, −0.012 to −0.015 at 200M, −0.031 to −0.035 at 500M, and −0.017 to −0.025 at 1B.Source: §C.5 · p. 29

Balance and entropy

• Column-router softmax entropy is 0.95–0.997 of log C, yet noise-free routing agrees with training-time-noise routing on 85–98% of tokens per block.Only 1–16% of tokens receive the same selection as bias-only routing; expert-router entropy is lower, at 0.4–0.9 of log E.Source: §C.6 · p. 31, Figure 14 · p. 33

Information in each level

• Column assignments contain 0.05 bits of corpus information out of 2.32, while expert choices contain 0.28 bits within NESTOR columns and 0.25 bits in the MoE transformer.With frozen routers, 4.7% of column assignments change during adaptation, versus 26% of within-column expert assignments and 22% of MoE expert assignments.Source: Table 9 · p. 31, §C.6 · p. 32

Routing counterfactuals

• Replacing column assignments raises held-out loss by 0.40–0.48 nats at 200M, 0.74 at 500M, and 0.62 at 1B; shuffled and random conditions agree.The increase is 18–28 times the MoE–NESTOR loss gap at each scale, and is similar in the column-routed model with shared experts, indicating the effect comes from column-routed attention rather than nesting.Source: §C.6 · pp. 33–34, Table 10 · p. 33
• At 200M, replacing both NESTOR expert levels raises loss by 3.8–4.0 nats, similar to 4.0–4.2 for the MoE expert level; at 1B, increases are 3.6 versus 4.5.At larger scales, the column level accounts for a larger share of routing’s effect; within NESTOR’s expert levels, loss depends most on shared experts.Source: §C.6 · pp. 32, 34

D Optimization and Initialization

• Attention and expert matrices use Muon, while other parameters use AdamW; the column router’s last layer receives a learning rate scaled by 1/M.Muon’s base learning rate was selected at 200M and reused at other scales; AdamW’s was fixed without a separate sweep.Source: D Optimization and Initialization · p. 34
• The architectures share an initialization scheme, with the final column-router layer initialized so routing is close to uniform and attention-residual queries initialized at zero.The router layer’s standard deviation is 1/Din; attention and expert output matrices use a standard deviation divided by √(2L).Source: D Optimization and Initialization · p. 34

E Compute Multipliers under Scaling Laws

• Under the stated scaling-law hypothesis, the compute multiplier between full architectures is constant across loss levels and grows with context length because NESTOR uses fewer attention FLOPs.The section gives full-architecture compute accounting analogous to the fixed-size result for Model 1.Source: E Compute Multipliers under Scaling Laws · p. 35

E.1 The Compute Multiplier

• The fitted compute multiplier is 1.13–1.29× across four scales with no trend, consistent with the proposition’s prediction that it is constant across losses under its hypotheses.The empirical fits have separate offsets, which the proposition does not allow; the appendix discusses their effect.Source: §E.1 · pp. 35–36

E.2 Long Contexts

• For equal shapes and balanced routing, the compute multiplier R(S) grows with context length and converges to 1/ρ, reflecting NESTOR’s lower attention-FLOP ratio.At equal FLOPs, NESTOR is trained on R(S) times as many tokens; at fixed context length in pretraining, R(S) tends to 1 as width grows.Source: §E.1 · p. 35, §E.2 · p. 36
• The equal-FLOP margin grows with context length when the learning-curve parameters change little with context, consistent with the observed Figure 4 pattern.The loss difference at equal FLOPs exceeds the equal-token difference by a positive term that increases with R(S).Source: §E.2 · p. 36

E.3 Proof of Proposition 3 and Extensions

• For a fixed-ratio allocation, the scaling-law loss has a unique minimizer because its excess-loss terms are opposing exponentials and the objective diverges at both ends.At the optimum, αa_sN^-α = βb_sT^-β.Source: §E.3 · pp. 36–37
• For loss L > L∞, solving the optimized scaling law gives compute F = k_0q_s(H/(L − L∞))^(1/p), proportional to q_s.Source: §E.3 · p. 37
• When a_N ≤ a_F, b_N ≤ b_F, and ν_N(N) ≤ ν_F(N) for all N, NESTOR has the lower loss frontier at every compute budget.At the same compute and model size, it can train on at least as many tokens as the MoE model.Source: §E.3 · p. 37
• With fixed tokens per parameter, a lower NESTOR loss limit makes the compute multiplier grow with scale.When the attention ratio varies with model size, the multiplier depends on the loss; if α ≠ β, it also depends on L.Source: §E.3 · p. 37

F.1 Sample and Compute Efficiency

• At equal compute, the flat predictor needs κ = (d_F/d_N)(S_0 + S − 1)/(S_0 + (S − 1)/C) times as many samples as the nested predictor.The context-cost factor increases with context length, from d_F/d_N at S = 1 toward C·d_F/d_N as S grows.Source: §F.1 · p. 38
• The flat predictor’s extra dimensions include zero off-diagonal pattern blocks that must be estimated, whereas routing removes other-label tokens before NESTOR applies its pattern.Source: §F.1 · p. 38
• For C ≥ 2, the flat predictor has C(C − 1)D² + C(E − 1)r(D − r) more local dimensions than the nested predictor.The two terms arise from routed attention and nesting experts within columns, respectively.Source: §F.1 · p. 38
• The limiting risk dimension is independent of label probabilities and context length; each dimension contributes σ² to the mean limiting excess risk.Rarely used modules have larger weight error but are used less often, and those factors cancel in excess risk.Source: §F.1 · p. 39
• Nesting’s statistical advantage requires both column-routed attention and experts within columns; neither component alone improves the loss.Source: §F.1 · p. 39
• The dimension comparison assumes specialization matters: if all tokens are relevant and patterns and row spaces do not vary by column, a shared attention can represent the target with fewer separately estimated components.In this case, NESTOR estimates a separate pattern and row space for each column.Source: §F.1 · p. 39

F.2 Forgetting

• During adaptation to a new coarse label, NESTOR changes only that label’s column, preserving predictions on other coarse labels and yielding zero forgetting there.Source: §F.2 · p. 39
• For the flat predictor, changes to a value map shared across labels are the source of forgetting on other coarse labels.Holding V = I fixed preserves the function class and removes this forgetting in the stated comparison.Source: §F.2 · p. 40

F.3 Proof of Theorem 4

• Least-squares excess risk asymptotically scales with σ² times the local dimension of the predictor’s function class at the target.The proof treats estimation as projection of an unconstrained least-squares estimator onto the class’s tangent space.Source: §F.3 · pp. 40–41
• The predictor-class dimensions are dim D_N = Cr(ED + D − r) + CD² and dim D_F = CEr(2D − r) + C²D².The proof subtracts C reparameterization directions to obtain the local function-class dimensions.Source: §F.3 · p. 42
• For C ≥ 2 and 1 ≤ r < D, the flat predictor’s dimension exceeds the nested predictor’s by a positive amount.Source: §F.3 · p. 42

F.5 Further Comparisons

• Wider flat experts do not reduce the limiting excess risk in the comparison; when Cr < D and the column row spaces are independent, a shared value map can represent the target but has extra map dimensions.Those extra terms count expert parameters used to select each column’s r coordinates from the shared map’s outputs.Source: §F.5 · p. 43
• An MoE transformer can simulate a NESTOR column with the same stored parameters, but requires a residual stream wider by a constant multiple of D.The authors therefore expect constant-factor differences between architectures at equal resources.Source: §F.5 · pp. 43–44