IntBMoE: Integrating Block-Level Conditioning into Expert Composition for Full-Participation Mixture-of-Experts

AI summary

Overview

Research area: Efficient deep learning architectures, specifically Mixture-of-Experts (MoE) layers for Transformers, evaluated on image classification, language modeling, and sequential recommendation.

Technical level: Intermediate. The paper assumes familiarity with Transformer FFN sublayers, sparse top-k routing, and the standard dense-vs-sparse MoE trade-off, but its framing along three resource axes is self-contained.

Scope: The paper proposes IntBMoE, a block-conditioned MoE layer that separates how many experts contribute knowledge to a token from how many are computed and how many expert-sized parameter sets are stored, and validates it on ImageNet-1K, MiniPile, IntTravel, plus a production deployment at AMap.

What This Paper Is About

Standard MoE designs force an unwanted coupling between three quantities for any given token: how many experts contribute knowledge to its output (participation), how many are actually computed (execution, i.e. compute cost), and how many expert-sized parameter sets must be built and stored (materialization, i.e. memory cost). Sparse routing keeps execution and materialization low but shrinks participation to a few experts; dense output-mixing restores full participation but makes execution grow with the number of experts; parameter merging keeps execution at one expert but makes materialization grow with the number of routing decisions. The goal of IntBMoE is to set all three quantities independently, so every token benefits from the full expert pool without dense execution or unbounded parameter materialization.

Key Contributions

1. Decoupling participation, execution, and materialization. IntBMoE uses a small codebook of learned embeddings to define the blocks a module can use, a shared hypernetwork to build each block by merging that layer's expert bases into a single composed expert, and a router that sends each token to only a few blocks. Participation is pool-wide, execution stays sparse, and materialization is bounded by the codebook size rather than by the input.

2. Dual-Path Residual Gating (DPRG). Each block's expert bases are merged twice — into a value path and a gate path — whose product is nonlinear in those bases. This adds expressiveness without enlarging the expert pool, at only a constant factor in parameter-synthesis and execution cost.

3. Visual evaluation and cross-domain generalization. IntBMoE is compared against representative sparse-routing and dense-participation MoE baselines on ImageNet-1K and improves on all of them; the same architecture is also tested on language modeling (MiniPile) and sequential recommendation (IntTravel).

4. Production deployment. IntBMoE is fully deployed in AMap's generative recommendation system, serving hundreds of millions of users under a 60 ms latency budget with a 2.4% relative UVCTR gain in large-scale online A/B testing.

Main Findings

• ImageNet-1K accuracy lead: IntBMoE reaches 73.76% Top-1 and 91.48% Top-5 on ImageNet-1K (1.28 million training images, 50,000 validation images, 1,000 classes, eight-layer DeiT-Tiny-style backbone, mean over three random seeds). Against the dense backbone this is +7.36 and +3.79 percentage points. Against the strongest baseline, SMEAR (71.78% Top-1, 90.33% Top-5), the margin is +1.98 Top-1 and +1.15 Top-5 percentage points. All MoE methods use comparable total parameter budgets of roughly 24M; IntBMoE activates 23.111M parameters.

• Inference cost: IntBMoE's total inference cost is 4.063 GFLOPs per image without caching and 3.457 GFLOPs per image with caching, measured at batch size 1. By comparison, μMoE (CP) reports 9.120 GFLOPs and SMEAR 1.493 GFLOPs.

• Ablations — the gate path matters most: On ImageNet-1K, removing the composed gate path drops Top-1 from 0.7376 to 0.6804 (the largest single drop), followed by the parameter-matched 1-Layer variant at 0.6878. Fixed λ gives 0.7244, softmax-normalized coefficients 0.7263, w/o Shared 0.7332, and w/o Filter 0.7372. The full model wins on Top-5 as well (0.9148 vs. 0.8831 for w/o Gate).

• Hyperparameter sensitivity: From the default configuration (K, E, k, L) = (8, 16, 2, 2), expanding the codebook from K=8 to K=32 raises Top-1 by only 0.07 percentage points, and expanding the expert pool from E=16 to E=64 raises it by only 0.15 percentage points. Increasing k consistently improves accuracy; L=2 achieves the highest Top-1 accuracy.

• Every expert basis contributes: Removing each of all 16 expert bases from one MoE layer at a time, across Layers 0, 2, 4 and 6, gives 64 removal settings, and every removal reduces Top-1 accuracy. Mean drops are 1.60, 0.49, 1.17 and 0.78 percentage points for Layers 0, 2, 4 and 6. Contributions are highly uneven in Layer 0, where removing E1, E12 or E15 costs 4.10, 4.99 and 9.15 percentage points; later layers are more balanced (0.26–0.74 points in Layer 2, 0.60–1.95 in Layer 4, 0.45–1.04 in Layer 6). The smallest decrease overall is 0.26 percentage points.

• Routing is class-conditioned and depth-dependent: With eight blocks, uniform allocation is 12.5%. Different classes favor different blocks in the same layer, and preferences shift with depth — at layer 4 hen favors b4 while Boston bull favors b6; at layer 6 their preferred blocks move to b0 and b3.

• Blocks learn distinct recipes, increasingly so with depth: Mean pairwise cosine similarity between block composition recipes falls monotonically across depth. For the value path it drops from 0.796 at layer 0 to 0.079, 0.019 and 0.010 at layers 2, 4 and 6; the gate path follows the same trend from 0.726 to 0.126, 0.083 and 0.043.

• Caching decouples request-time cost from pool size: Without caching, peak inference memory rises from 54.57 MB to 627.83 MB and inference cost from 3.495 to 8.305 GFLOPs as the expert pool E grows from 1 to 128. With caching both remain constant at 104.33 MB and 3.457 GFLOPs. Caching uses more memory for small expert pools but becomes more memory-efficient from 16 experts onward.

• Generalization to language modeling: On MiniPile (a 6-GB subset of the deduplicated Pile, using an 18-layer Llama-style causal Transformer), IntBMoE reaches a test loss of 2.6802 and perplexity 14.5878. Relative to μMoE (CP), the strongest baseline at 15.0306 PPL, this is a 2.9% perplexity reduction; relative to the dense backbone at 16.6621 PPL, it is 12.4%.

• Generalization to recommendation: On IntTravel (162.8 million users, 7.3 million POIs, 4.13 billion interactions; "Where" task predicting the destination POI of the next journey), IntBMoE achieves the highest HR@1 (0.6852), HR@5 (0.8692) and NDCG@5 (0.7850). Sparse MoE is reduced to causal form here by constructing each slot from the current and preceding tokens only, and SMEAR uses token-level composition.

• Ablations hold across domains: On both MiniPile and IntTravel the gate path is again among the largest contributors (w/o Gate gives 2.7660 loss / 15.8947 PPL on MiniPile), the parameter-matched 1-Layer variant underperforms the full model, and unconstrained variance-scaled coefficients beat softmax-normalized composition.

• Production result: In AMap's generative recommendation system for the initial map screen (predicting the Top-10 POIs a user is likely to visit, which determine the map viewport), cached IntBMoE delivers a 2.4% relative UVCTR gain in online A/B testing under a 60 ms latency budget. The paper content provided is truncated mid-sentence in this section, so the remaining online-evaluation details are not available here.

Methodology in Plain English

The key move is to separate building expert transformations from applying them.

Each IntBMoE module keeps a codebook of K learned embeddings — one per candidate "block" — and a pool of E expert bases at every internal layer. A small hypernetwork reads only the codebook embedding (never the input token) and outputs two sets of coefficients, a value recipe and a gate recipe. Those coefficients mix the layer's expert bases into a single composed expert, so every block draws on the entire pool. Because the hypernetwork never sees the token, all K blocks can be computed once and reused.

At run time, a router looks at each token and picks the top k blocks. Before entering a block, the token is multiplied element-wise by a soft mask computed from the token and that block's codebook embedding, so different blocks get different views of the same token. Inside a block, each of its L layers computes a value path and a gate path from independently composed weights, normalizes the gate path with RMSNorm, and multiplies them together as value × (1 + λ · SiLU(gate)), with LayerNorm between layers and a learnable residual scale λ. The outputs of the selected blocks are combined using the router's softmax weights and added to an always-active shared SwiGLU expert.

The composition coefficients are deliberately not softmax- or sigmoid-normalized: they can be negative and need not sum to one, so composition spans the linear hull of the expert bases rather than only their convex hull. A 1/√E factor keeps the composed parameter scale roughly stable as the pool grows. The cost analysis shows composition costs O(KED) once, routed execution costs O(TkD), and the amortized per-token cost is O(KED/T + kD); with caching, the composition term disappears at request time, so request-time computation no longer grows with E for fixed K and k.

Why This Matters

Impact on research. The paper reframes the MoE design space along three axes — participation, execution, materialization — and shows that existing strategies are not forced trade-offs but consequences of coupling composition with execution. This gives the field a vocabulary for comparing MoEs that isn't just "which routing rule," and it identifies a concrete mechanism (a finite, input-independent codebook plus a shared hypernetwork) for escaping the usual trilemma. The expert-removal study is also a useful methodological contribution, since it provides a scale-compensated way to test whether nominally participating experts actually matter.

Real-world applications:

• Map and POI recommendation: the deployed AMap setting predicts the Top-10 POIs a user is likely to visit and uses them to set the map viewport on app open.
• Sequential next-destination prediction, as measured by the IntTravel "Where" task predicting the destination POI of a user's next journey.
• Image classification on ImageNet-1K-scale data using a DeiT-Tiny-style vision backbone.
• Causal language modeling, evaluated on the 6-GB MiniPile subset with an 18-layer Llama-style Transformer.

Industry relevance. The deployment result is the practical headline: full expert participation with bounded memory, running under a 60 ms latency budget for hundreds of millions of users and returning a 2.4% relative UVCTR gain in online A/B testing. The caching analysis matters directly for serving costs, because it shows that request-time memory and FLOPs can be made independent of expert-pool size once the blocks are precomputed — the property that makes larger pools affordable at inference.

Future Directions

• How far can K and E be pushed under caching? The sensitivity study finds diminishing returns from K=8 to K=32 (+0.07 points) and E=16 to E=64 (+0.15 points) on ImageNet-1K, but the caching analysis shows request-time cost is flat in E from 16 experts onward. Whether that headroom pays off at much larger scale, or in domains with more heterogeneous inputs, is untested here.

• Adaptive codebooks or block counts. The codebook size K and the number of selected blocks k are fixed hyperparameters. The paper's own evidence that routing preference is class-conditioned and depth-dependent, and that later layers learn much more diverse recipes, suggests per-layer or input-adaptive K and k could be worth exploring — but the paper does not test this.

• Training-time composition cost. The complexity analysis notes that during training all images in a minibatch share the same composed blocks, so composition is performed once per batch and amortized. For very large E and K, or for small batches, this remains an open cost, and no training-cost measurements are reported.

• Extending the evaluation beyond the tested causal tasks. The generalization study covers two causal settings and modifies Soft MoE and SMEAR to be causal for comparison. Whether the block-conditioned design holds up on other modalities (multimodal models and large-scale language models are named in the introduction

Authors’ abstract

Mixture-of-Experts (MoE) scales capacity, but existing designs cannot set three quantities independently. For a single token, participation is how many experts contribute knowledge to its output, execution is how many are actually computed (compute cost), and materialization is how many expert-sized parameter sets must be built and stored (memory cost). Sparse routing keeps execution and materialization low, but shrinks participation: for each token, only a few experts contribute. Dense output-mixing restores full participation, but its execution grows with the number of experts. Parameter-merging keeps execution at one expert, but its materialization grows with the number of routing decisions. We propose IntBMoE, a block-conditioned MoE that decouples all three by pairing dense expert composition with sparse block execution. Its blocks come from a small learned codebook, one per entry. At each internal layer, a lightweight hypernetwork merges all expert bases in that layer's pool into one composed expert. Participation is full, because every composed expert draws on the entire pool. Execution stays sparse, because a router sends each token to only a few blocks. Materialization is bounded, because the codebook, not the input, fixes how many blocks exist. Dual-Path Residual Gating (DPRG) further couples two independently composed paths through multiplicative gating. Experiments on image classification show consistent gains over representative sparse and dense MoE baselines. Additional experiments on language modeling and sequential recommendation validate its generalization beyond vision. IntBMoE is fully deployed in AMap's generative recommendation system, serving hundreds of millions of users under a 60ms latency budget, with a 2.4% relative UVCTR gain in online A/B testing. Our code is available at https://github.com/AMAP-ML/DreamX-Rec/.