• How MoE is enabled
• Recommended presets
  • Stable baseline
  • Higher-capacity experimental setup
• Advanced features in the current implementation
  • Router types
  • Adaptive top-k
  • Hash routers
  • Grouped expert kernel path
  • MTP heads inside MoE
• Integration with ForecastingModel

MoE Guide

ForeBlocks integrates Mixture-of-Experts into the transformer feedforward path through MoEFeedForwardDMoE.

You typically do not instantiate this block directly. Instead, you enable MoE through transformer constructor arguments or the FeedForwardBlock module.

Related docs:

• Documentation Overview
• Getting Started
• Transformer
• Custom Blocks
• Advanced MoE — routers, load-balancing, expert types, production tuning

How MoE is enabled

from foreblocks import TransformerEncoder, TransformerDecoder

encoder = TransformerEncoder(
    input_size=8,
    d_model=256,
    nhead=8,
    num_layers=4,
    use_moe=True,
    num_experts=8,
    top_k=2,
)

decoder = TransformerDecoder(
    input_size=1,
    output_size=1,
    d_model=256,
    nhead=8,
    num_layers=4,
    use_moe=True,
    num_experts=8,
    top_k=2,
)

Recommended presets

Stable baseline

encoder = TransformerEncoder(
    input_size=8,
    d_model=256,
    nhead=8,
    num_layers=4,
    use_moe=True,
    num_experts=8,
    num_shared=1,
    top_k=2,
    router_type="noisy_topk",
    routing_mode="token_choice",
    z_loss_weight=1e-3,
    moe_aux_lambda=1.0,
)

Higher-capacity experimental setup

encoder = TransformerEncoder(
    input_size=8,
    d_model=384,
    nhead=8,
    num_layers=6,
    use_moe=True,
    num_experts=16,
    num_shared=2,
    top_k=2,
    routing_mode="expert_choice",
    moe_capacity_factor=1.5,
    z_loss_weight=1e-3,
    use_gradient_checkpointing=True,
)

Advanced features in the current implementation

Router types

ForeBlocks provides several router implementations in foreblocks.modules.moe.experts.routers:

• NoisyTopKRouter (default, recommended)
• StraightThroughTopKRouter
• ContinuousTopKRouter
• HashTopKRouter
• AdaptiveNoisyTopKRouter
• SoftDenseRouter
• AuxiliaryTokenRouter
• LinearRouter

See the Advanced MoE Guide for detailed router configuration.

Adaptive top-k

adaptive_noisy_topk can vary the effective number of experts selected per token.

This path also tracks per-token k statistics and supports a REINFORCE-style adaptive-k loss internally.

Hash routers

hash_topk and multi_hash_topk are available when you want routing diversity without a standard learned dense router over all experts.

Grouped expert kernel path

The implementation can use grouped expert kernels and fused top-k routing in favorable runtime conditions.

You usually do not need to tune these first. They are lower-level performance details rather than primary modeling controls.

MTP heads inside MoE

The MoE block supports optional multi-token-prediction heads:

• mtp_num_heads
• mtp_loss_weight

This is an advanced decoder-side path and should be treated as research functionality, not a default production setting.

Integration with ForecastingModel

from foreblocks import ForecastingModel

model = ForecastingModel(
    encoder=encoder,
    decoder=decoder,
    forecasting_strategy="transformer_seq2seq",
    model_type="transformer",
    target_len=24,
    output_size=1,
)