MaskCoFT co-fine-tunes routers and experts with a learnable mask, cutting expert fetches per token by 23.7% on Mixtral-8x7B and 10.1% on DeepSeek-V2-Lite
Related research and updates
Synopsis
The work proposes MaskCoFT, a masked co-adaptive fine-tuning method that trains routers and experts together using only the cross-entropy loss: during fine-tuning a learnable binary mask restricts each layer's Top-K routing to a subset of experts so the experts adapt to the redirected tokens, and at inference the learned mask becomes a soft prior that re-ranks experts while every expert remains selectable; with a simulated GPU cache of 4 experts per layer for Mixtral-8x7B and 12 for DeepSeek-V2-Lite, expert fetches per token drop by 23.7% and 10.1% relative to the base model, time per output token in real offloading system serving falls by up to 16.4% and 5.5%, and average accuracy over nine benchmarks stays above the base model by 0.92 and 0.53 points.
Figure 1: Routed expert weights dominate memory, and loading them dominates offloaded decoding latency. MaskCoFT reduces expert fetches while preserving accuracy. (a) Routed experts hold 96.6% of Mixtral-8×7B’s parameters and 91.7% of DeepSeek-V2-Lite’s. (b) Decode latency breakdown of Mixtral-8×7B in MoE-Offloading ( Eliseev and Mazur, 2023 ) , measured with NVIDIA Nsight Systems ( NVIDIA Corporation, 2026 ) . (c) Expert fetches per token and average accuracy over nine benchmarks under GPU cache holding 4 experts per layer for Mixtral-8×7B and 12 for DeepSeek-V2-Lite, both relative to the base model. The GPU cache holds 4 experts per layer for Mixtral-8×7B and 12 for DeepSeek-V2-Lite. MaskCoFT cuts fetches by 23.7% on Mixtral-8×7B and by 10.1% on DeepSeek-V2-Lite. It keeps average accuracy above the base model.
arXivInterpretation
It introduces masked co-adaptive fine-tuning: a learnable binary mask restricts each layer's Top-K routing to a subset of experts during fine-tuning, and routers and experts are updated together under the cross-entropy loss so experts adapt to the tokens the new routing sends them. Prior router-only fine-tuning reshapes routing while keeping experts frozen, so experts cannot adapt to the tokens the new routing assigns; MaskCoFT lets routers and experts adapt jointly and needs no loss term beyond cross-entropy. The abstract states the mechanism and the training objective (cross-entropy loss alone) explicitly and reports results at two model scales, which is method-level self-reported evidence.
At inference the learned mask acts as a soft prior that re-ranks experts while keeping every expert selectable, so expert fetches are reduced without hard-excluding any expert. Unlike hard pruning of routing to a fixed expert set, the soft prior preserves the selectability of all experts, turning the fine-tuning-time routing constraint into an inference-time ranking preference. The abstract states that 'every expert remains selectable' but does not give the concrete form of the soft prior or an ablation.
Under a simulated GPU cache (4 experts per layer for Mixtral-8x7B and 12 for DeepSeek-V2-Lite), expert fetches per token fall by 23.7% and 10.1% relative to the base model. It converts a routing-level change directly into the fetch metric that dominates offloading, quantifying the gain under cache-constrained conditions. Relative reductions for two models and two cache configurations, as reported in the abstract; the cache sizes are simulated settings.
In real offloading system serving, time per output token drops by up to 16.4% (Mixtral-8x7B) and 5.5% (DeepSeek-V2-Lite), while average accuracy over nine benchmarks stays above the base model by 0.92 and 0.53 points. It provides evidence on both end-to-end serving latency and multi-benchmark accuracy, indicating that fewer fetches translate into real decoding speedups without an accuracy penalty. The abstract reports 'up to' latency reductions and nine-benchmark average accuracy differences, without per-benchmark detail, variance, or significance information.
Perspective
The work targets memory-constrained MoE inference that relies on expert offloading: most experts reside in host memory and are loaded on demand, so decoding speed depends on how many experts each token must fetch. The method applies to MoE layers with Top-K routing, requires fine-tuning on the target model, and participates in expert re-ranking as a soft prior at inference. The reported gains correspond to two specific models and two simulated cache sizes (4 experts per layer for Mixtral-8x7B and 12 for DeepSeek-V2-Lite), so its scope is best read as this class of cache-constrained offloading serving rather than all MoE deployments. For engineering readers aiming to reduce single-GPU memory footprint while holding multi-benchmark accuracy, this line offers a reproducible direction for fine-tuning and inference modification.
The abstract does not describe the mask parameterization or training details, how the soft prior participates in re-ranking, the fine-tuning data and compute cost, or per-benchmark accuracy, latency distributions, and variance; the cache sizes are simulated settings, and behavior on real hardware and at different batch sizes remains to be observed. Gains are reported as 'up to', so the lower end across configurations is not clear. In addition, the current text is abstract-level and contains no figures or experiment tables, so open questions such as training stability and item-by-item comparison against baselines like router-only fine-tuning are best resolved in the paper's methods section.