PatchMoE: Fine-Grained Mixture-of-Experts

• PatchMoE is a novel MoE architecture that uses patch-level routing and expert partitioning to enhance representation learning in time-series data.
• It replaces standard FFN layers with modules that integrate shared and routed experts using recurrent noisy gating for dynamic, context-aware token processing.
• Empirical studies demonstrate that PatchMoE achieves state-of-the-art performance in forecasting, anomaly detection, imputation, and classification by balancing temporal and channel dependencies.

PatchMoE refers to a class of Mixture-of-Experts (MoE) neural architectures in which routing, expert partitioning, and load balancing operate at the granularity of "patches" or tokens—typically contiguous segments of an input sequence, image, or multivariate timeseries. Unlike traditional MoE variants predominantly used in natural language processing or computer vision, PatchMoE is tailored for input modalities or problems where fine-grained, context-aware expert allocation can significantly improve representation learning and task specialization. Its recent implementations, such as in time-series analytics, demonstrate state-of-the-art performance by combining task-aware recurrent routing, channel-temporal tokenization, and specialized load-balancing losses to address the heterogeneity and hierarchical abstraction requirements of distinct downstream tasks (Wu et al., 26 Sep 2025). Parallel theoretical and empirical work further establishes the sample complexity and generalization advantages of patch-level MoE for structured data, especially compared to single-expert models (Chowdhury et al., 2023).

1. Motivation and Limitations of Standard MoE

Time-series analytics encompasses forecasting, anomaly detection, imputation, and classification, each exhibiting distinct layerwise representational requirements. Canonical MoE architectures, such as Switch Transformer and Multi-gate Mixture-of-Experts (MMoE), employ a task-agnostic router for expert selection, making per-layer or per-task gating decisions in isolation. This static routing neither adapts to diverse task-specific abstraction needs nor exploits joint patterns across channels. Furthermore, most transformer-based time-series MoEs employ channel-independent tokenization, treating temporal patches in each channel separately and missing cross-channel dependencies. PatchMoE directly addresses these limitations by enabling task- and context-aware routing and expert allocation along both temporal and channel axes (Wu et al., 26 Sep 2025).

2. PatchMoE Architecture

PatchMoE augments a standard Transformer stack by replacing every Feed-Forward (FFN) layer with a patchwise Mixture-of-Experts module. The input multivariate time series $X\in\mathbb{R}^{N\times T}$ (channels × timesteps) is segmented into $n$ patches per channel, each patch of length $p$, and linearly embedded into tokens, resulting in $X_{token}\in\mathbb{R}^{N\times n\times d}$. Within each Transformer layer:

• Multi-Head Self-Attention (MSA) computes intermediate representations.
• PatchMoE contains $N_s$ "shared experts" (global, task-agnostic two-layer FFNs) and $N_r$ "routed experts" (patch- and channel-specific dynamic selection).
• Each token is processed by the shared experts, while routed experts are selected via a Recurrent Noisy Gating mechanism. The forward computation per token aggregates outputs from both expert types and normalizes the sum with LayerNorm.

This design supports fine-grained, token-level routing that is both channel- and temporally aware and can learn complex dependencies across both axes (Wu et al., 26 Sep 2025).

3. Recurrent Noisy Gating: Task-Aware Routing

Conventional noisy gating (cf. Shazeer et al. 2017) performs token routing layer-wise, with no explicit memory of prior routing choices. PatchMoE's Recurrent Noisy Gating (RNG) employs a GRU cell to fuse prior hidden states with current MSA outputs, producing a per-token latent embedding. This recurrent state allows the router to track and adapt expert usage through the network and across layers, thereby encoding hierarchical, task-specific routing policies. The gating logits are then parameterized as Gaussians, and during training, injected noise encourages exploration and regularization. The sparse $k$-hot gating vector is determined by the KeepTopK and softmax operations, with randomness removed at inference. This adaptive, hierarchical mechanism enables PatchMoE to allocate capacity according to the depth and abstraction profile required by each specific downstream task—shallow for forecasting or anomaly detection, deeper for imputation and classification as established through CKA analyses (Wu et al., 26 Sep 2025).

4. Routing Strategy and Temporal–Channel Load Balancing

To prevent expert collapse—where a few experts monopolize most tokens—the routing strategy incorporates dual axis load balancing: separate penalties are introduced for overuse along the temporal (patch index) and channel (variable) dimensions. Specifically, the Temporal–Channel Load Balancing Loss is formulated to encourage diverse tokens from both axes to be routed to different experts, enforcing a balanced, sparse utilization. For each patch and each channel, the loss computes the frequency and probability of each routed expert being selected, then applies a weighted penalty if expert allocations become concentrated. The final loss is a linear combination of the temporal and channel terms, weighted by hyperparameters $\alpha$ and $\beta$. This regularization promotes robust cross-channel and cross-time modeling, which is critical for high-dimensional time-series datasets (Wu et al., 26 Sep 2025).

5. Training Regimen and Robust Hyperparameters

PatchMoE is trained end-to-end with Adam optimizer and, depending on the target, uses MSE or MAE for forecasting/imputation, plus the load balancing term. The authors report robust hyperparameters suitable across multiple tasks: patch length $p=24$, window size $T=96$, $L=3$ MoE layers, $N_r=10$ routed experts, $k=3$ top routed experts per token, and $\alpha = \beta = 0.01$. Notably, no gradient clipping or "drop-last" data tricks are needed, and the shared experts stabilize learning of global features (Wu et al., 26 Sep 2025).

6. Empirical Validation and Ablation

PatchMoE demonstrates state-of-the-art (SOTA) results across primary time-series challenges:

Task Type	SOTA Metric	Previous SOTA	PatchMoE Value	Reference
Univariate Forecast	msMAPE	∼20.8%, MASE 2.02	19.95%, 1.97	(Wu et al., 26 Sep 2025)
Multivariate Forecast	MSE (Electricity)	0.163 (iTransformer)	0.158	(Wu et al., 26 Sep 2025)
Anomaly Detection	Avg F1/AUC	0.835/0.838 (CATCH)	0.842/0.861	(Wu et al., 26 Sep 2025)
Imputation	MSE (Electricity)	0.109 (TimeMixer++)	0.052	(Wu et al., 26 Sep 2025)
Classification	Avg Accuracy	73.6% (TimesNet; 10 UEA sets)	74.11%	(Wu et al., 26 Sep 2025)

Ablation studies demonstrate that removing the recurrent gating increases MSE by ~4% on ETTh1, omitting shared experts degrades Traffic accuracy by ~6.9%, and dropping the balancing loss impairs performance by 1–3%, substantiating the necessity of each PatchMoE component (Wu et al., 26 Sep 2025).

7. Theoretical Foundations and Generalization Properties

Patch-level MoE is further justified by theoretical advances. For supervised classification, partitioning an input into $n$ patches and sending top-$l$ per expert ($l\ll n$) enables polynomial reductions in both model and sample complexity compared to single-expert CNNs. Specifically, sample complexity ratios scale as $\Theta((n/l)^8)$ and model size as $\Theta((n/l)^{10})$, given proper discriminative routing (Chowdhury et al., 2023). Empirical results validate these advantages on MNIST, CelebA, and CIFAR-10, where patch-level MoEs retain SOTA accuracy while using significantly less supervision and compute. The discriminative routing property—where routers learn to direct class-relevant patches to consistent experts—prevents learning of spurious correlations and enhances generalization (Chowdhury et al., 2023).

8. Limitations and Future Directions

PatchMoE's primary limitations involve increased computational cost from noise injection and load-balancing overhead, and the need to tune expert and loss hyperparameters for very small datasets. Potential improvements include adaptive expert counts, integration with diffusion or state-space modeling, and possible hybridization with large pretrained sequence models (Wu et al., 26 Sep 2025). The fundamental principle—fine-grained, context- and task-aware patch routing combined with explicit balance enforcements—suggests broad applicability to other domains, contingent on efficient optimization and further theoretical analysis.