PaperlayerView exact revision on arXiv

Source-linked AI summary

X-MoD: Practical Scaling Laws for Sparse-Depth Routing Beyond Mixture-of-Depths

Bowen Dong, Yilong Fan, Tengyu Pan, Yike Zhang, Zhenyu Li, Zijian Zhang, Xuewei Li, Mei Yu, Jianyong Wang

arXiv:2609.34212v1cs.LGcs.CL

TL;DR

MoD’s fixed alternation limits the separation between total and active-equivalent capacity, leaving sparse-depth architecture design without a practical scaling rule. X-MoD decouples token sparsity from anchor stride, and a FLOP-matched scaling law models capacity, context, and stride effects. The paper reports improvements over Dense and MoD, competitiveness with representative MoE baselines, and held-out loss-prediction support.

  • Problem

    MoD’s one-sparse–one-dense alternation tightly couples total and active-equivalent capacity, while practical guidance for sparse-depth routing remains an open design question.

    Source: §1 · p. 1

  • Method

    X-MoD decouples token sparsity from anchor stride, and a FLOP-matched residual over dense baselines yields a scaling law for choosing routing configurations.

    Source: §1 · p. 2

  • Results

    X-MoD improves over Dense and MoD, is competitive with representative MoE baselines, and is supported by held-out loss prediction and ablations.

    Source: §4.3 · p. 7, §7 · p. 10

  • Takeaways & Limitations

    The results suggest sparse-depth routing is a viable and quantitatively analyzable axis for scaling conditional-computation Transformers.

    Source: §7 · p. 10

    Systems measurements cover the evaluated implementation; efficient kernels, scheduling, and parallelization remain directions for improving practicality at larger scales.

    Source: §5 · p. 9

Abstract

from arXiv · show

Mixture-of-Depths (MoD) enables conditional computation across Transformer depth by routing only a subset of tokens through selected layers, but its original one-sparse--one-dense alternation tightly couples total capacity to active capacity and limits sparse-depth scaling. We introduce X-MoD, a scalable sparse-depth architecture that decouples token sparsity from anchor stride, allowing total parameter count to grow while keeping active-equivalent capacity nearly fixed. To make deep sparse routing trainable, X-MoD combines dense anchors with variance-scaled layer-wise gating and depth-wise token balancing. To make this regime analyzable and usable, we formulate sparse-depth routing as a conditional architecture-design problem: given compute, context length, and active-equivalent backbone size, how should the routing configuration be chosen? We develop a practical scaling-law framework by fitting X-MoD relative to FLOP-matched dense baselines, yielding an interpretable law that decomposes performance into sparse-capacity gain, sparse-context correction, and anchor-stride interaction. The law predicts validation loss across routing configurations and reveals how context length, model scale, and anchor stride shape sparse-depth performance. We validate the architecture and law through pretraining sweeps, held-out scaling-law prediction, ablations, downstream evaluations, and comparisons with Dense, MoD, and representative MoE baselines.

1 INTRODUCTION

Sparse-depth routing raises whether token routing can decouple total model capacity from active computation, as MoE does, and whether practical scaling laws can guide its design. X-MoD separates token sparsity from anchor stride, while a FLOP-matched framework models sparse-depth configuration.

Motivation

  • MoD’s one-sparse–one-dense alternation tightly couples total and active-equivalent capacity, while sparse-depth scaling also faces concentrated updates and reduced attention context.The introduction frames a need for both a trainable architecture and a rule balancing capacity against context.Source: §1 · p. 1

Architecture

  • X-MoD independently controls token sparsity K and anchor stride A, expanding total parameter count while keeping active-equivalent capacity nearly fixed.Dense anchors, variance-scaled layer-wise gating, depth-wise token balancing, and an asymmetric hierarchy support this regime.Source: §1 · p. 2, Figure 1 · p. 2

Scaling law and validation

  • A FLOP-matched residual over dense baselines yields an interpretable law for how total versus active capacity, sparse-layer context, and anchor stride shape performance.The paper validates the architecture and law through held-out prediction, ablations, downstream evaluations, and comparisons with Dense, MoD, and representative MoE baselines.Source: §1 · p. 2

2 X-MOD: A SCALABLE SPARSE-DEPTH ARCHITECTURE

X-MoD removes MoD’s fixed alternation by independently controlling token sparsity and sparse-depth density. Dense anchors and three stabilizing mechanisms support deeper sparse routing while allowing total capacity to grow relative to active-equivalent capacity.

2.1 REVISITING MIXTURE-OF-DEPTHS

  • MoD’s dense–sparse alternation stabilizes training but constrains total-to-active capacity, preventing substantial decoupling and continuous total-capacity growth at roughly fixed active capacity.The original design’s reported optimum is around K = 8 under this restricted regime.Source: §2.1 · p. 3

2.2 X-MOD

  • X-MoD places AK sparse layers between dense anchors, giving each token an average of A sparse updates and making the total-to-active ratio grow approximately linearly with K when the sparse backbone dominates.The architecture independently controls token sparsity and anchor stride.Source: §2.2 · p. 3

2.3 STABILIZATION AND TOKEN BALANCING

  • X-MoD supports deeper sparse stacks with whole-residual gating and variance scaling, cross-layer token-choice bias, and a dense initial prefix before sparse routing.The bias discourages repeatedly selecting the same tokens without enforcing uniform routing; the stabilizers make regimes beyond MoD’s restrictive alternation trainable.Source: §2.3 · pp. 3–4

3 PRACTICAL SCALING LAWS FOR X-MOD

The paper treats sparse-depth routing as a conditional design problem and fits a FLOP-matched residual to guide choices of sparsity and anchor stride. Its empirical regularities motivate a law that captures the capacity–context trade-off and anchor-stride interaction.

3.1 A CONDITIONAL DESIGN PROBLEM

  • Given compute, context length, and active-equivalent backbone size, the design problem is to choose K and A to minimize loss, using a matched-FLOPs residual to isolate sparse-depth effects from the dense reference.The dense reference is interpolated for the same sequence length and backbone family.Source: §3.1 · pp. 4–5

3.2 MECHANISTIC CONSTRAINTS AND EMPIRICAL REGULARITIES

  • Sweeps show U-shaped loss in K, an optimum that shifts strongly with context length but weakly with active-equivalent size, and an anchor-stride effect that varies with K.The 64k-context and 1.65B active-equivalent settings extend beyond the law’s fitting ranges and follow the same qualitative trends.Source: §3.2 · p. 5

3.3 A PRACTICAL SCALING LAW

  • With A = 1, increasing K initially improves performance through greater sparse capacity, but excessive sparsity hurts as each sparse layer receives less context; the approximation predicts strong context-length and weak active-scale dependence for optimal K.The anchor-stride term further favors larger sparsity as more active-equivalent budget goes to sparse refinements, with saturating gains.Source: §3.3 · p. 6

4 EXPERIMENTS

Across evaluated scales, X-MoD improves over Dense and original MoD and is competitive with or better than representative MoE baselines; held-out tests also show close agreement between scaling-law predictions and measurements.

4.1 EXPERIMENTAL SETTINGS

  • Models use FineWeb-Edu, are compared at matched training FLOPs, and are evaluated with causal-threshold routing for validation loss.Sweeps vary sparsity, context length, active-equivalent size, and anchor stride; prespecified L = 64k and Nact = 1.65B extensions are excluded from fitting.Source: §4.1 · p. 6

4.2 OUT-OF-SAMPLE VALIDATION

  • Close agreement between predictions and measurements persists in grouped tests that withhold context-length, anchor-stride, or model-scale groups.Accuracy also holds at longer-context and larger-model extensions.Source: §4.2 · pp. 6–7

4.3 MAIN RESULTS AGAINST DENSE, MOD, AND MOE

  • X-MoD consistently improves over Dense and original MoD, and is competitive with or better than representative MoE baselines at similar total parameter scale and training compute, with substantially lower per-token FLOPs.Downstream averages generally follow the validation-loss trend.Source: §4.3 · p. 7, Table 1 · p. 7

4.4 ABLATIONS

  • Routed-FFN controls improve over Dense at matched active-equivalent size, per-token FLOPs, and training compute, with further gains from the larger-capacity variant.These controls support a contribution from conditional parameter capacity even without sparse-layer attention-context selection.Source: §4.4 · pp. 7–8, Table 2 · p. 8
  • Under a fixed-compute objective, full X-MoD allocates computation more effectively than expanding sparse-layer key/value context.The Selected-Q/full-KV diagnostic retains X-MoD’s router, selected queries, total depth, and total parameters, while added K/V projections increase active computation.Source: §4.4 · p. 8, Table 2 · p. 8
  • Removing any stabilization component increases validation loss and lowers downstream average; dense anchors and depth-wise token-choice bias have the largest stated roles.Dense anchors support full-token synchronization, while token-choice bias addresses token concentration; gated residual scaling gives a smaller consistent improvement and the dense prefix improves early representation learning.Source: §4.4 · p. 8, Table 3 · p. 8
  • For A1K16, removing token-choice bias concentrates routing on a persistent subset, whereas full A4K8 and A3K16 combine near-complete interval coverage with low consecutive-layer overlap.Threshold and top-k masks agree highly when computed from the same routing scores across all three full-model configurations.Source: §4.4 · p. 8
  • Stable training through 412 total layers demonstrates that the architectural changes support substantial sparse-depth expansion.Source: §4.4 · p. 8

5 LIMITATIONS AND FUTURE WORK

The study evaluates X-MoD in isolation and its fitted law within the sizes, contexts, sparsities, and anchor strides studied. Broader model scales, combined routing mechanisms, and more efficient sparse-depth implementations remain future work.

5 LIMITATIONS AND FUTURE WORK

  • The study treats X-MoD as a sparse-depth architecture in isolation, leaving combinations with sparse-width routing, recurrent depth, or inference-time token pruning for future investigation.Whether these mechanisms’ scaling laws compose additively or interact nonlinearly is an open question.Source: §5 · pp. 8–9
  • The fitted law is evaluated on the studied model sizes, context lengths, sparsity levels, and anchor strides, including held-out groups and out-of-range configurations.Future tests should extend to production-scale models and substantially deeper sparse stacks.Source: §5 · p. 9
  • Systems measurements quantify throughput only for the evaluated implementation; efficient kernels, scheduling, and parallelization remain directions for larger-scale practicality.Source: §5 · p. 9

6 RELATED WORKS

Prior work studies conditional computation across model width, depth, and other sparse-compute mechanisms. The paper positions its focus as pretraining-time sparse-depth scalability and design under fixed compute, an area it describes as underexplored.

6 RELATED WORKS

  • MoE routes tokens across feed-forward experts to expand capacity under a nearly fixed active budget, while MoD routes selected tokens through Transformer layers.Related work also applies conditional depth to multimodal tokens and uses token skipping or layer pruning for downstream compute reduction; this paper instead studies pretraining-time sparse-depth scalability under fixed compute.Source: §6 · p. 9
  • Sparse-depth scaling remains underexplored relative to scaling-law work on dense models, context, data mixtures, and sparsely activated architectures, especially MoE.Source: §6 · p. 9

7 CONCLUSION

X-MoD decouples token sparsity from anchor stride, expanding total capacity while controlling active-equivalent capacity. Its scaling law and empirical evaluations support sparse-depth routing as a viable, quantitatively analyzable scaling axis.

7 CONCLUSION

  • X-MoD extends Mixture-of-Depths beyond one-sparse–one-dense alternation, increasing total capacity while keeping active-equivalent capacity controlled.Source: §7 · p. 9
  • The law decomposes X-MoD behavior into sparse-capacity reward, sparse-context correction, and anchor-stride interaction, providing quantitative configuration guidance.The authors report improvements over Dense and MoD, competitiveness with representative MoE baselines, and support from held-out loss prediction and ablations.Source: §7 · p. 10

DISCLOSURE OF GENERATIVE AI USE

The authors report using generative AI for translation and readability editing, and state that they reviewed the assisted work and take responsibility for the final content.

DISCLOSURE OF GENERATIVE AI USE

  • Generative AI assisted with translation and readability editing, but not with the listed research, methods, experiments, or results interpretation.The authors state that they reviewed all AI-assisted work and take responsibility for the final content.Source: DISCLOSURE OF GENERATIVE AI USE · p. 10

A SCALING-LAW DERIVATIONS

The appendix derives an anchor-stride measure of sparse-equivalent computation, motivates exposure and compute corrections for the scaling law, and gives an asymptotic intuition for optimal sparsity.

A.1 ANCHOR-STRIDE SPARSE-EQUIVALENT SHARE

  • The monotone, saturating quantity ω(A) captures the fraction of active-equivalent computation allocated to sparse refinements, equals 1 at A = 1, and approaches 2 as A grows.In the controlled sweep, changing A reallocates the non-prefix active path between sparse refinements and dense anchors without changing the total active-equivalent budget.Source: §A.1 · p. 15

A.2 CANDIDATE EXPOSURE CORRECTION AND COMPUTE-TO-TOKEN RELATION

  • The candidate exposure term tests whether sparse-versus-dense exposure differences add information beyond capacity and context terms.Source: §A.2 · p. 15
  • Sparse-parameter exposure is proportional to D_sparse/K because each sparse layer processes only a 1/K fraction of tokens.The analytic exposure correction motivates the candidate law, while empirical fits use logged token counts and FLOPs; the compute-to-token derivation separates token-linear work from attention work.Source: §A.2 · pp. 15–16

A.3 APPROXIMATE OPTIMUM SPARSITY

  • The approximate optimum-sparsity expression is an asymptotic design intuition; representative comparison configurations instead come from the matched-budget pilot grid.Source: §A.3 · p. 16

B.1 TRAINING AND EVALUATION

The appendix documents training and evaluation choices, comparison and pilot protocols, scaling-law sweeps, and grouped held-out validation. The procedures align comparisons by logged FLOPs and distinguish fitting-domain observations from held-out or out-of-range tests.

B.1 TRAINING AND EVALUATION

  • Models use a warmup-then-cosine schedule, matched active-equivalent backbones across architectures, and logged-FLOP accounting that includes routed-token fractions.Muon at peak learning rate 3.2 × 10^-3 gave the best dense validation loss and was used in the main runs.Source: §B.1 · pp. 17–18

B.2 MAIN-COMPARISON PROTOCOL

  • Main comparisons match Dense, MoD, MoE, and X-MoD at each scale's training-compute budget and L = 32768, then evaluate validation loss and six zero-shot tasks.Source: §B.2 · p. 18

B.3 PILOT CONFIGURATION-SELECTION PROTOCOL

  • A pilot sweep spans 24 X-MoD and 12 MoE configurations and selects representative settings before larger-scale comparisons.Source: §B.3 · pp. 18–19

B.4 SCALING-LAW SWEEPS AND OUT-OF-RANGE PROTOCOL

  • Scaling-law sweeps vary sparsity, context, active-equivalent size, and anchor stride; L = 64k and Nact = 1.65B are prespecified extensions excluded from fitting.At L = 64k, K = 1 is the dense reference; the best tested anchor stride is A = 5, not a claimed minimum beyond the evaluated range.Source: §B.4 · p. 19

B.5 FITTING AND OUT-OF-SAMPLE VALIDATION PROTOCOL

  • Grouped held-out splits withhold complete context, scale, or anchor-stride groups to test transfer beyond correlated points within observed sweeps.Out-of-range predictions freeze the law fitted on in-range observations; held-out dense observations define the matched-FLOPs reference but do not fit the sparse-depth residual law.Source: §B.5 · pp. 19–20

B.6 ABLATION AND CONTROLLED-SCALING PROTOCOLS

The protocols define controlled comparisons for X-MoD, align training-time top-k routing with causal threshold evaluation, and specify how routing diagnostics and systems efficiency are measured.

B.6 ABLATION AND CONTROLLED-SCALING PROTOCOLS

  • The mechanism controls compare routed FFN refinements with selected-query/full-KV attention; the latter is not active-equivalent matched and processes fewer training tokens at fixed compute.Selected-Q/full-KV processes 46.8% and 42.0% of full X-MoD’s training tokens at A4K8 and A3K16, respectively.Source: §B.6 · p. 20
  • Component removals use X-MoD A1K16 as the reference at Nact = 936M, L = 32768, and C = 1.5 × 1020 FLOPs, measuring ΔLoss against the full model.The same six downstream tasks are evaluated for the full model and each removal.Source: §B.6 · p. 20
  • The deep sweep compares Deep Dense with A1K4, A1K8, A1K12, and A1K16 at matched active capacity and compute, but varies K, total parameters, per-token FLOPs, and token exposure jointly.All Deep X-MoD variants use dense anchors, gated residual scaling, and token-choice bias.Source: §B.6 · p. 20

B.7 ROUTING DIAGNOSTIC PROTOCOL

  • A binary cross-entropy auxiliary loss aligns top-k training selections with the causal evaluation threshold of 0.5; its coefficient is set to 10−4 to balance routing alignment against language-modeling loss.A coefficient that is too large degrades language-modeling loss, while one that is too small leaves a larger train–evaluation routing gap.Source: §B.7 · pp. 20, 22
  • Routing diagnostics measure interval coverage, adjacent-layer Jaccard overlap, pooled update counts, and top-k/threshold mask agreement; repeated-subset and independent-subset schemes provide analytical references.Coverage and overlap are averaged over sequence–interval observations; update-count distributions pool token–interval observations.Source: §B.7 · p. 22

B.8 SYSTEMS MEASUREMENT PROTOCOL

  • Systems efficiency is measured on eight A100 80GB GPUs at L = 32k, with training throughput averaged over 2,000 steps and validation-forward throughput over four full validation evaluations.The validation-forward benchmark includes routing, token movement, full-vocabulary projection, and loss computation; it is not a KV-cache prefill or decoding benchmark.Source: §B.8 · p. 22

C.1 SCALING-LAW FITS AND EXPOSURE COMPARISON

The scaling-law fit captures matched-dense residual variation, while the experiments compare pilot configurations, total-depth scaling, routing diagnostics, and measured systems efficiency. X-MoD improves over Deep Dense across the tested deep configurations, and its pilot variants achieve the lowest loss among tested X-MoD configurations within the specified MoE parameter budgets.

C.1 SCALING-LAW FITS AND EXPOSURE COMPARISON

  • 98.5% of in-range matched-dense residual variance is explained by the fitted law across 109 observations; the prespecified L = 64k and Nact = 1.65B observations are excluded.The reduced law has lower combined in-range fitting RMSE than the full candidate law, 0.008758 versus 0.01349, supporting omission of PD.Source: §C.1 · p. 23

C.2 PILOT CONFIGURATION COMPARISON

  • Within K = 8 and K = 16 groups, X-MoD A4K8 and A3K16 achieve the lowest loss among tested X-MoD configurations without exceeding the respective MoE total-parameter budgets.The corresponding best tested MoE variants are MoE 8:64 and MoE 8:128.Source: §C.2 · p. 23

C.3 TOTAL-DEPTH SCALING

  • All evaluated Deep X-MoD configurations improve on the Deep Dense reference, with validation loss decreasing monotonically across the tested configurations.The sweep uses matched active-equivalent capacity and training compute.Source: §C.3 · p. 23, Table 9 · p. 24

C.4 ROUTING DIAGNOSTICS

  • Routing diagnostics report coverage, consecutive-layer overlap, and mask agreement; Fig. 5 plots token-wise sparse-update distributions for A1K16, A4K8, and A3K16.The figure distinguishes repeated-subset and independent-subset references; integer heights, not filled areas, represent probability.Source: §C.4 · p. 23, Table 10 · p. 24, Figure 5 · p. 25

C.5 SYSTEMS EFFICIENCY

  • Systems efficiency measurements are reported at the 936M active-equivalent scale, with speedups defined relative to adjacent, approximately total-parameter-matched MoE baselines.The table caption identifies scale, counted compute, and peak training memory as reported measurements.Source: Table 11 · p. 25, §C.5 · p. 25

Related papers

Loading 2609.34212v1…