block-ffn-sparse-mixture-experts

name	block-ffn-sparse-mixture-experts
title	BlockFFN: Towards End-Side Acceleration-Friendly Mixture-of-Experts with Chunk-Level Activation Sparsity
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2507.08771
keywords	["Mixture-of-Experts","Model Compression","Sparsity","Edge Acceleration","Token-Level Computation"]
description	Deploy efficient MoE models on resource-constrained edge devices by learning chunk-level activation sparsity that achieves 3.67× speedup. Use when you need to compress LLMs for on-device inference while maintaining reasoning quality and supporting speculative decoding acceleration.

BlockFFN: Sparse Mixture-of-Experts Optimized for End-Device Deployment

Standard Mixture-of-Experts architectures suffer from two fundamental constraints that prevent efficient on-device deployment: non-differentiable routing mechanisms that limit training convergence, and low chunk-level sparsity where consecutive tokens activate overlapping expert sets, blocking parallelization opportunities. BlockFFN addresses both problems through differentiable routers and specialized loss functions that enforce structured sparsity patterns compatible with hardware acceleration.

The core insight is that chunk-level sparsity (whether the same expert is active across consecutive tokens) is what enables acceleration kernels, not just token-level sparsity. By learning when to deactivate experts across sequences of tokens, the model becomes compatible with speculative decoding and GEMM-based acceleration while remaining highly parameter-efficient.

Core Concept

BlockFFN replaces TopK routing with a differentiable linear router combined with ReLU gating, enabling end-to-end optimization. Critically, it introduces two novel loss functions: Activation Locality Loss encourages neighboring tokens to activate similar expert sets, and Chunk Sparsification Loss directly minimizes the probability that an expert activates within L-token windows.

This architecture design enables hardware accelerators to compute only the union of activated experts across n tokens in parallel, masking unnecessary activations post-computation. The result is genuine speedup on edge devices rather than theoretical parameter reduction.

Architecture Overview

• Router Module: Linear transformation with ReLU activation followed by RMSNorm, enabling differentiable expert selection unlike non-differentiable TopK
• Expert Modules: Non-gated MLPs with Swish activation, sized smaller than typical (d_e << d_h) for fine-grained sparsity control
• Activation Pattern Learning: Sigmoid-sharpened activation patterns for smooth gradients during Activation Locality Loss computation
• Chunk-Level Sparsity Mechanism: Probabilistic computation of expert activation within L-token windows using binary models
• CUTLASS Acceleration Kernel: GEMM-based implementation processing n=32 tokens simultaneously, computing only the union of activated experts

Implementation

Router Architecture Setup

The router combines simplicity with differentiability, avoiding TopK bottlenecks while maintaining effective expert selection.

import torch
import torch.nn as nn
import torch.nn.functional  F

 (nn.Module):
     ():
        ().__init__()
        .hidden_dim = hidden_dim
        .num_experts = num_experts
        .sharpness = sharpness

        
        .router_linear = nn.Linear(hidden_dim, num_experts, bias=)
        .rms_norm = nn.RMSNorm(num_experts)

     ():
        

        
        scores = .router_linear(x)  

        
        scores = F.relu(scores)

        
        scores = .rms_norm(scores)

        
        probs = torch.sigmoid(.sharpness * scores)

         probs, scores

 (nn.Module):
     ():
        ().__init__()
        
        .fc1 = nn.Linear(hidden_dim, expert_dim)
        .fc2 = nn.Linear(expert_dim, hidden_dim)
        .dropout = nn.Dropout(dropout)
        .activation = nn.SiLU()  

     ():
        x = .fc1(x)
        x = .activation(x)
        x = .dropout(x)
        x = .fc2(x)
         x