Source-linked AI summary

Meta-Learning Where to Allocate Experts: Task-Conditioned Layer-Wise Compression for MoEs

Rongfeng Wang, Shichao Weng, Zhiqiang Wang, Xinyu Liu, Yang Yi, Peilong Zhou, Hongwei Tang

arXiv:2608.26650v1cs.CL

TL;DR

Fixed expert counts do not reflect task- and depth-dependent MoE demand, while existing approaches provide only offline layer allocation or local token adaptation. MetaNet uses a support-set controller to predict per-layer budgets and bounded routing biases with the backbone and router frozen. It achieves tunable accuracy–activation trade-offs on DeepSeek-MoE-16B-Chat and transfers to C-Eval without retraining, while also showing broader transfer across backbones and architectures.

  • Problem

    MoE inference commonly fixes expert activation across layers and tasks despite evidence that expert redundancy and capacity demand vary with depth and task.

  • Method

    MetaNet infers per-layer expert budgets and bounded routing biases from support-set routing statistics while leaving the backbone, original router, and expert pool frozen.

  • Results

    MetaNet produces tunable accuracy–expert-activation trade-offs across MMLU and C-Eval, with the MMLU-trained controller transferring to C-Eval without retraining.

  • Takeaways & Limitations

    Across multiple backbones and architectures, modeling task- and depth-varying expert demand gives better accuracy–expert-activation trade-offs than uniform fixed-k routing.

  • Takeaways & Limitations

    MetaNet reduces routed-expert workload, not model storage or total end-to-end computation; proportional latency gains require runtime support for smaller expert sets.

Abstract

from arXiv · show

Mixture-of-Experts (MoE) models route each token to a subset of expert networks, increasing capacity while keeping per-token computation sparse. In many deployed MoEs, the number of active experts is fixed across layers and tasks, although layer roles and expert redundancy vary with depth and demand varies with difficulty. Existing approaches address only part of this setting: layer-wise allocations are usually determined offline and reused for all tasks, while token-level methods vary expert activation using local routing signals without task-level context. We propose MetaNet, a support-set controller that predicts, for each layer, an expert-retention threshold and a bounded routing bias. The backbone, experts, and router remain frozen. On DeepSeek-MoE-16B-Chat, MetaNet provides a tunable accuracy-expert-activation trade-off. Relative to fixed k=6, a conservative setting activates 3.61 experts on average (40% fewer) and achieves comparable MMLU accuracy (0.489 vs. 0.474), whereas an aggressive setting activates 2.28 experts on average (62% fewer) with accuracy approximately 3.7 percentage points lower. The MMLU-trained controller also transfers to C-Eval without retraining, activating 2.90 experts on average (52% fewer than fixed k=6) at 0.386 accuracy.

1. Introduction

MetaNet addresses task-conditioned, layer-wise expert allocation for frozen MoE inference, replacing fixed budgets with support-set policies. It offers a tunable accuracy–activation trade-off, including comparable MMLU accuracy with 40% fewer activated experts and stronger compression with a modest accuracy decline.

  • Motivation: Standard MoE inference typically fixes expert counts across layers, tokens, and tasks despite depth- and task-dependent redundancy.Existing layer-wise allocations are fixed offline, while token-level routing lacks task-level context.
  • Approach: MetaNet predicts per-layer retention thresholds and bounded routing biases from a small support set while keeping the backbone, router, and experts frozen.The controller infers how much expert capacity each layer receives without modifying pretrained parameters.
  • Results: 62% fewer activated experts reduces the mean from 6.00 to 2.28, with a 3.7 percentage-point MMLU accuracy drop under the aggressive setting.The two operating points use the same controller family under different deployment budgets.
  • Results: The controller transfers from MMLU to C-Eval without retraining while maintaining a similar reduction in activated experts.

2. Related Work

Prior MoE compression methods separately address expert representation, token-adaptive routing, layer-wise allocation, or router modification. MetaNet combines task-conditioned layer-wise budgeting with a bounded bias while retaining the original router and frozen backbone.

  • MoE background: MoE models route each token to k of E expert FFNs, keeping per-step compute proportional to k while preserving large total capacity.
  • Inference compression: Expert pruning, merging, and residual approximation determine compressed expert representations before deployment.
  • Existing approaches: Token-adaptive methods vary expert activation from local input signals, whereas layer-wise methods assign offline budgets reused across tasks.Neither category assigns task-conditioned per-layer capacity from a support set.
  • Router modification: MetaNet keeps the original MoE router frozen and adds only a bounded task-conditioned bias to expert selection.
  • Systems support: Specialized systems address dynamic expert workloads, communication, and parallelism because routing sparsity alone does not guarantee proportional end-to-end speedup.MetaNet changes routing policy rather than runtime and is complementary to these optimizations.
  • Meta-learning: Unlike parameter-generating or parameter-updating meta-learning methods, MetaNet outputs a per-layer routing policy from a support set.

3. Method

MetaNet derives task-conditioned, layer-wise expert budgets and bounded routing biases from support-set routing statistics while keeping the MoE backbone and original router frozen. It separates capacity allocation from fine-grained token routing and trains the controller with task, compression, alignment, and consistency objectives.

  • Support profiling and policy: MetaNet maps a task support set to per-layer expert budgets, replacing fixed top-K_nat with bounded layer budgets k_τ,l.The controller uses k_min=1 and k_max=12 for DeepSeek-MoE, allowing fewer experts in redundant layers and more in diffuse layers.
  • Support profiling and policy: The support profile summarizes task-favored experts using per-layer routing distributions and concentration statistics.The statistics include normalized entropy, maximum expert probability, and top-R routing mass.
  • MetaNet controller: The two controller branches predict a retention threshold for expert count and a per-expert routing bias for close gate scores.The budget branch uses routing distributions and statistics, while the prior branch provides a weak task-level preference.
  • MetaNet controller: The bounded bias is limited to magnitude τ_b, with α(1 − β)=0.005 in the main configuration, so the pretrained router remains the main source of expert ordering.The bias is intended to weakly influence selection when gate scores are similar.
  • Inference: At query time, MetaNet adds the task-level bias to the original gate logits and applies layer-specific top-k selection without modifying backbone parameters.The policy is computed once from the support set and reused across query examples.
  • Training objective: Training uses episodic support-query learning with task loss, one-sided quality-preserving budget pressure, alignment, consistency, diversity, entropy, and ranking terms.The quality penalty compares dynamic-routing loss with a detached native-router reference and uses a Huber transition.

4. Experimental setup

Experiments evaluate MetaNet on MMLU and C-Eval, with additional cross-backbone, cross-architecture, and ablation studies. Evaluation treats each subject as a task, separates support from query examples, and reports accuracy alongside routed-expert workload and implementation-specific latency.

  • Main setup: The main DeepSeek-MoE-16B-Chat setup freezes 27 routed layers with 64 experts per layer and native top-k K_nat=6.MMLU uses 8 support and 16 disjoint query examples; C-Eval uses five labeled development examples for support and up to 16 validation examples for queries.
  • Experiments: The study evaluates MMLU meta-test performance, zero-shot C-Eval transfer, cross-backbone transfer to OLMoE, cross-architecture vision tasks, and component and layer-position ablations.The vision check uses a frozen GoogLeNet backbone and a 3-of-4 branch policy.
  • Task splits: MMLU uses 34 meta-train, 11 meta-validation, and 12 disjoint meta-test subjects, producing 192 held-out query predictions.Support examples construct routing profiles, while query examples determine the reported metric.
  • Task splits: C-Eval-Full evaluates the MMLU-trained controller on all 52 subjects without updates, while C-Eval-10 is a separate 10-subject subset for expanded baselines and sweeps.The two C-Eval panels are not compared directly.
  • Robustness checks: Repeated-sampling checks vary support and query sampling while keeping subject splits fixed and examples disjoint.Two seeds and two episodes per held-out MMLU subject yield 384 query predictions per seed.
  • Baselines: Baselines include fixed top-k values, a support-statistics HyperRouter-style head, and adapted AdaMoE, Dynamic MoE, and Probe Pruning rules.The adapted baselines are placed in the same frozen DeepSeek harness rather than reproduced with their original training procedures.
  • Metrics: Accuracy is averaged over task-level exact-match rates, while mean dynamic top-k and active expert ratio measure routed-expert workload rather than total model FLOPs.Latency is wall-clock time in the common three-GPU harness and is reported separately from accuracy and workload.

5. Results

MetaNet improves the accuracy–expert-activation trade-off through task-conditioned, layer-wise routing, including transfer across backbones and tasks. Results show that adaptive allocation, branch coordination, and depth-aware compression outperform or complement uniform fixed-k routing.

  • Main results: 0.438 accuracy at mean top-k=2.28 achieves a 62% reduction in activated experts, outperforming fixed k=3 at 0.417 accuracy and a larger active ratio.MetaNet’s learned layer-wise allocation better matches heterogeneous expert demand than a uniform budget.
  • Expanded baselines and operating points: 2.90 mean experts on C-Eval-Full, 52% fewer than fixed k=6, transfers the MMLU-trained controller without retraining and exceeds the nearby fixed k=2 baseline at 0.377 accuracy.The transfer result indicates that the learned controller can operate across tasks without task-specific retraining.
  • Expanded baselines and operating points: 0.573 vs. 0.521 on MMLU and 0.375 vs. 0.363 on C-Eval show cross-backbone gains on OLMoE-1B-7B at mean top-k=7.52.Because this budget is close to OLMoE’s native k=8, the improvement may reflect bounded routing bias rather than expert-count reduction alone.
  • Ablation studies: 0.943 gate agreement with both branches preserves original-gate expert selection while the budget branch reduces active experts; removing either branch weakens one of these properties.Budget-only reaches 0.375 accuracy and 0.038 gate agreement, whereas prior-only uses k=12 throughout and reaches 0.484 accuracy.
  • Ablation studies: 0.005 accuracy points lost by compressing deep layers at ratio 0.069 is the best efficiency–accuracy trade-off in the layer-range sweep, while shallow-only compression loses 0.078 points.Full-layer compression reaches ratio 0.036 and 0.438 accuracy, indicating that joint optimization can compensate partly for harder shallow-layer budgets.
  • Mechanistic analysis: 0.10 mean cross-layer Jaccard overlap and 3.4% lower latency support layer-dependent expert selection, but routed-expert sparsity does not translate proportionally into end-to-end speedup.The inference stack retains dense attention, dispatch, synchronization, and communication costs.

6. Conclusion

MetaNet models task- and depth-dependent expert demand during frozen MoE inference, improving the accuracy–expert-activation trade-off over uniform fixed-k routing.

  • MetaNet predicts per-layer expert budgets and a bounded routing bias from support-set information without updating backbone parameters.

Limitations

MetaNet reduces routed-expert workload but does not directly reduce model storage or total end-to-end computation, and its latency benefits depend on runtime support.

  • MetaNet reduces activated-expert workload rather than model storage or total end-to-end computation.The full expert pool remains resident in memory, so proportional parameter-storage or peak-memory savings do not follow in the current implementation.
  • Wall-clock speedup need not scale with mean top-k reduction without specialized dispatch and communication support.Integrating MetaNet with optimized MoE runtimes is left for future work.
  • The main MMLU result uses 192 held-out queries and only two seeds in the repeated-sampling check.These two seeds do not support a reliable confidence interval or significance test.
  • Evaluation focuses mainly on academic multiple-choice benchmarks, while open-ended, long-context, and safety-sensitive applications may exhibit different routing patterns.

Ethical Considerations

The work targets inference-time efficiency for released MoE models and reports indirect misuse risks without adding generative capabilities or collecting personal data.

  • The study evaluates inference-time efficiency for released MoE models rather than introducing a new model, dataset, or deployment system.
  • Lower activated-expert workload can reduce hardware cost and energy use when paired with suitable sparse dispatch kernels.
  • Cheaper inference may make existing LLM misuse scenarios easier to scale, representing an indirect potential risk.
  • The method does not add new generative capabilities, collect personal data, or rely on human-subject experiments.

B. Additional experimental details

Additional experiments use fixed subject splits, separate C-Eval scopes, no C-Eval parameter updates, and inference-compatible baseline adaptations rather than exact method reproductions.

  • The fixed MMLU split contains 34 meta-train, 11 meta-validation, and 12 meta-test subjects.
  • The subject split is fixed across all sampling seeds.
  • C-Eval-Full uses 52 subjects and 828 labeled queries, while C-Eval-10 reports ten specified subjects separately.
  • MetaNet receives no C-Eval parameter updates in either evaluation scope.
  • Adapted baselines preserve the DeepSeek backbone and original router but are not exact reproductions of methods requiring different training or pruning procedures.

C. Additional experimental results

The additional comparison uses a common inference harness and reports realized operating points rather than assuming nominal training targets are achieved. Table 8 presents an operating-point sweep with active ratio and latency as key measurements.

  • C. Additional experimental results: The comparison includes adapted inference-compatible decision rules, while some original methods require training or replacing routing components.Thus, the expanded comparison does not necessarily reproduce every method exactly.
  • C.1. Operating-point sweep: Four independently trained controller profiles are compared in the additional operating-point sweep.The profiles are labeled bt0p3, bt0p4, bt0p5, and bt0p8.
  • C.1. Operating-point sweep: Nominal base budget targets u_0 = 0.3, 0.4, 0.5, and 0.8 configure training objectives rather than realized active ratios.The reported metrics therefore use realized values, and the rows are not assumed to be universally Pareto-optimal.
  • C.1. Operating-point sweep: Active ratio is the realized mean top-k divided by 64, while latency is measured in the same three-GPU harness.These measurements characterize routed-expert workload and wall-clock inference behavior under the common harness.

C.2. Sensitivity to episode sampling

The repeated-sampling check varies support and query sampling while holding the subject split fixed. Across two seeds, MetaNet profiles show different accuracy and mean top-k values, but the small sample limits statistical conclusions.

  • C.2. Sensitivity to episode sampling: The fixed subject split is evaluated under two seeds that change support/query sampling only.Each seed contains two episodes per subject and 384 query predictions, with support and query examples disjoint.
  • C.2. Sensitivity to episode sampling: Accuracy is 0.505 ± 0.030 for fixed k=6, 0.517 ± 0.020 for fixed k=12, 0.506 ± 0.022 for conservative MetaNet, and 0.477 ± 0.022 for aggressive MetaNet.Values are reported as mean ± sample standard deviation across the two seeds.
  • C.2. Sensitivity to episode sampling: Mean top-k is 6.00 ± 0.00 for fixed k=6, 12.00 ± 0.00 for fixed k=12, 4.94 ± 0.27 for conservative MetaNet, and 4.15 ± 0.46 for aggressive MetaNet.The repeated evaluation reports the realized mean top-k for each operating profile.
  • C.2. Sensitivity to episode sampling: With only two seeds, the repeated-sampling results measure sampling sensitivity rather than establish statistical significance.The reported variation should therefore be interpreted as sensitivity evidence, not a significance test.
Loading 2608.26650v1…