Patch-level Routing in Mixture-of-Experts is Provably Sample-efficient
  for Convolutional Neural Networks
In deep learning, mixture-of-experts (MoE) activates one or few experts
(sub-networks) on a per-sample or per-token basis, resulting in significant
computation reduction. The recently proposed \underline{p}atch-level routing in
\underline{MoE} (pMoE) divides each input into $n$ patches (or tokens) and
sends $l$ patches ($l\ll n$) to each expert through prioritized routing. pMoE
has demonstrated great empirical success in reducing training and inference
costs while maintaining test accuracy. However, the theoretical explanation of
pMoE and the general MoE remains elusive. Focusing on a supervised
classification task using a mixture of two-layer convolutional neural networks
(CNNs), we show for the first time that pMoE provably reduces the required
number of training samples to achieve desirable generalization (referred to as
the sample complexity) by a factor in the polynomial order of $n/l$, and
outperforms its single-expert counterpart of the same or even larger capacity.
The advantage results from the discriminative routing property, which is
justified in both theory and practice that pMoE routers can filter
label-irrelevant patches and route similar class-discriminative patches to the
same expert. Our experimental results on MNIST, CIFAR-10, and CelebA support
our theoretical findings on pMoE's generalization and show that pMoE can avoid
learning spurious correlations.
Mohammed Nowaz Rabbani Chowdhury, Shuai Zhang, Meng Wang, Sijia Liu, Pin-Yu Chen
Patch-level Routing in Mixture-of-Experts is Provably Sample-efficient for Convolutional Neural Networks

###### Abstract

In deep learning, mixture-of-experts (MoE) activates one or few experts (sub-networks) on a per-sample or per-token basis, resulting in significant computation reduction. The recently proposed patch-level routing in MoE (pMoE) divides each input into \(n\) patches (or tokens) and sends \(l\) patches (\(l\ll n\)) to each expert through prioritized routing. pMoE has demonstrated great empirical success in reducing training and inference costs while maintaining test accuracy. However, the theoretical explanation of pMoE and the general MoE remains elusive. Focusing on a supervised classification task using a mixture of two-layer convolutional neural networks (CNNs), we show for the first time that pMoE provably reduces the required number of training samples to achieve desirable generalization (referred to as the sample complexity) by a factor in the polynomial order of \(n/l\), and outperforms its single-expert counterpart of the same or even larger capacity. The advantage results from the discriminative routing property, which is justified in both theory and practice that pMoE routers can filter label-irrelevant patches and route similar class-discriminative patches to the same expert. Our experimental results on MNIST, CIFAR-10, and CelebA support our theoretical findings on pMoE's generalization and show that pMoE can avoid learning spurious correlations.

Machine Learning, ICML

## 1 Introduction

Deep learning has demonstrated exceptional empirical success in many applications at the cost of high computational and data requirements. To address this issue, mixture-of-experts (MoE) only activates partial regions of a neural network for each data point and significantly reduces the computational complexity of deep learning _without hurting the performance_ in applications such as machine translation and natural image classification (Shazeer et al., 2017; Yang et al., 2019). A conventional MoE model contains multiple experts (subnetworks of the backbone architecture) and one learnable router that routes each input sample to a few but not all the experts (Ramachandran and Le, 2018). Position-wise MoE has been introduced in language models (Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2022), where the routing decisions are made on embeddings of different positions of the input separately rather than routing the entire text-input. Riquelme et al. (2021) extended it to vision models where the routing decisions are made on image patches. Zhou et al. (2022) further extended where the MoE layer has one router for each expert such that the router selects partial patches for the corresponding expert and discards the remaining patches. We termed this routing

Figure 1: An illustration of pMoE. The image is divided into \(20\) patches while the router selects \(4\) of them for each expert.
mode as _patch-level routing_ and the MoE layer as _patch-level MoE_ (pMoE) layer (see Figure 1 for an illustration of a pMoE). Notably, pMoE achieves the _same_ test accuracy in vision tasks with 20% less training compute, and 50% less inference compute compared to its single-expert (i.e., one expert which is receiving all the patches of an input) counterpart of the same capacity (Riquelme et al., 2021).

Despite the empirical success of MoE, it remains elusive in theory, why can MoE maintain test accuracy while significantly reducing the amount of computation? To the best of our knowledge, only one recent work by Chen et al. (2022) shows theoretically that a conventional sample-wise MoE achieves higher test accuracy than convolutional neural networks (CNN) in a special setup of a binary classification task on data from linearly separable clusters. However, the sample-wise analyses by Chen et al. (2022) do not extend to patch-level MoE, which employ different routing strategies than conventional MoE, and their data model might not characterize some practical datasets. This paper addresses the following question theoretically:

_How much computational resource does pMoE save from the single-expert counterpart while maintaining the same generalization guarantee?_

In this paper, we consider a supervised binary classification task where each input sample consists of \(n\) equal-sized patches including _class-discriminative_ patterns that determine the labels and _class-irrelevant_ patterns that do not affect the labels. The neural network contains a pMoE layer1 and multiple experts, each of which is a two-layer CNN2 of the same architecture. The router sends \(l\) (\(l\ll n\)) patches to each expert. Although we consider a simplified neural network model to facilitate the formal analysis of pMoE, the insights are applicable to more general setups. Our major results include:

Footnote 1: In practice, pMoEs are usually placed in the last layers of deep models. Our analysis can be extended to this case as long as the input to the pMoE layer satisfies our data model (see Section 4.2).

Footnote 2: We consider CNN as expert due to its wide applications, especially in vision tasks. Moreover, the pMoE in (Riquelme et al., 2021; Zhou et al., 2022) uses two-layer Multi-Layer Perceptrons (MLPs) as experts in vision transformer (ViT), which operates on image patches. Hence, the MLPs in (Riquelme et al., 2021; Zhou et al., 2022) are effectively non-overlapping CNNs.

**1.** To the best of our knowledge, this paper provides **the first theoretical generalization analysis of pMoE**. Our analysis reveals that pMoE with two-layer CNNs as experts can achieve the same generalization performance as conventional CNN while reducing the sample complexity (the required number of training samples to learn a proper model) and model complexity. Specifically, we prove that as long as \(l\) is larger than a certain threshold, pMoE reduces the sample complexity and model complexity by a factor polynomial in \(n/l\), indicating an improved generalization with a smaller \(l\).

**2. Characterization of the desired property of the pMoE router.** We show that a desired pMoE router can dispatch the same class-discriminative patterns to the same expert and discard some class-irrelevant patterns. This discriminative property allows the experts to learn the class-discriminative patterns with reduced interference from irrelevant patterns, which in turn reduces the sample complexity and model complexity. We also prove theoretically that a separately trained pMoE router has the desired property and empirically verify this property on practical pMoE routers.

**3. Experimental demonstration of reduced sample complexity by pMoE in deep CNN models.** In addition to verifying our theoretical findings on synthetic data prepared from the MNIST dataset (LeCun et al., 2010), we demonstrate the sample efficiency of pMoE in learning some benchmark vision datasets (e.g., CIFAR-10 (Krizhevsky, 2009) and CelebA (Liu et al., 2015)) by replacing the last convolutional layer of a ten-layer wide residual network (WRN) (Zagoruyko & Komodakis, 2016) with a pMoE layer. These experiments not only verify our theoretical findings but also demonstrate the applicability of pMoE in reducing sample complexity in deep-CNN-based vision models, complementing the existing empirical success of pMoE with vision transformers.

## 2 Related Works

**Mixture-of-Experts.** MoE was first introduced in the 1990s with dense sample-wise routing, i.e. each input sample is routed to all the experts (Jacobs et al., 1991; Jordan & Jacobs, 1994; Chen et al., 1999; Tresp, 2000; Rasmussen & Ghahramani, 2001). Sparse sample-wise routing was later introduced (Bengio et al., 2013; Eigen et al., 2013), where each input sample activates few of the experts in an MoE layer both for joint training (Ramachandran & Le, 2018; Yang et al., 2019) and separate training of the router and experts (Collobert et al., 2001, 2003; Ahmed et al., 2016; Gross et al., 2017). Position/patch-wise MoE (i.e., pMoE) recently demonstrated success in large language and vision models (Shazeer et al., 2017; Lepikhin et al., 2020; Riquelme et al., 2021; Fedus et al., 2022). To solve the issue of load imbalance (Lewis et al., 2021), Zhou et al. (2022) introduces the _expert-choice routing_ in pMoE, where each expert uses one router to select a fixed number of patches from the input. This paper analyzes the sparse patch-level MoE with expert-choice routing under both joint-training and separate-training setups.

**Optimization and generalization analyses of neural networks (NN).** Due to the significant nonconvexity of deep learning problem, the existing generalization analyses are
limited to linearized or shallow neural networks. The Neural-Tangent-Kernel (NTK) approach (Jacot et al., 2018; Lee et al., 2019; Du et al., 2019; Allen-Zhu et al., 2019; Zou et al., 2020; Chizat et al., 2019; Ghorbani et al., 2021) considers strong over-parameterization and approximates the neural network by the first-order Taylor expansion. The NTK results are independent of the input data, and performance gaps in the representation power and generalization ability exist between the practical NN and the NTK results (Yehudai and Shamir, 2019; Ghorbani et al., 2019, 2020; Li et al., 2020; Malach et al., 2021). Nonlinear neural networks are analyzed recently through higher-order Taylor expansions (Allen-Zhu et al., 2019; Bai and Lee, 2019; Arora et al., 2019; Ji and Telgarsky, 2019) or employing a model estimation approach from Gaussian input data (Zhong et al., 2017; Ma et al., 2020; Ma et al., 2020; Li et al., 2022), but these results are limited to two-layer networks with few papers on three-layer networks (Allen-Zhu et al., 2019; Allen-Zhu and Li, 2019; 2020; Li et al., 2022).

The above works consider arbitrary input data or Gaussian input. To better characterize the practical generalization performance, some recent works analyze structured data models using approaches such as feature mapping (Li and Liang, 2018), where some of the initial model weights are close to data features, and feature learning (Daniely and Malach, 2020; Shalev-Shwartz et al., 2020; Shi et al., 2021; Allen-Zhu and Li, 2022; Li et al., 2023), where some weights gradually learn features during training. Among them, Allen-Zhu and Li (2020); Brutzkus and Globerson (2021); Karp et al. (2021) analyze CNN on learning structured data composed of class-discriminative patterns that determine the labels and other label-irrelevant patterns. This paper extends the data models in Allen-Zhu and Li (2020); Brutzkus and Globerson (2021); Karp et al. (2021) to a more general setup, and our analytical approach is a combination of feature learning in routers and feature mapping in experts for pMoE.

## 3 Problem Formulation

This paper considers the supervised binary classification3 problem where given \(N\) i.i.d. training samples \(\{(x_{i},y_{i})\}_{i=1}^{N}\) generated by an unknown distribution \(\mathcal{D}\), the objective is to learn a neural network model that maps \(x\) to \(y\) for any \((x,y)\) sampled from \(\mathcal{D}\). Here, the input \(x\in\mathbb{R}^{nd}\) has \(n\) disjoint patches, i.e., \(x^{\intercal}=[x^{(1)\intercal},x^{(2)\intercal},...,x^{(n)\intercal}]\), where \(x^{(j)}\in\mathbb{R}^{d}\) denotes the \(j\)-th patch of \(x\). \(y\in\{+1,-1\}\) denotes the corresponding label.

Footnote 3: Our results can be extended to multiclass classification problems. See Section M in the Appendix for details.

### Neural Network Models

We consider a pMoE architecture that includes \(k\) experts and the corresponding \(k\) routers. Each router selects \(l\) out of \(n\) (\(l<n\)) patches for each expert separately. Specifically, the router for each expert \(s\) (\(s\in[k]\)) contains a trainable gating kernel \(w_{s}\in\mathbb{R}^{d}\). Given a sample \(x\), the router computes a routing value \(g_{j,s}(x)=\langle w_{s},x^{(j)}\rangle\) for each patch \(j\). Let \(J_{s}(x)\) denote the index set of top-\(l\) values of \(g_{j,s}\) among all the patches \(j\in[n]\). Only patches with indices in \(J_{s}(x)\) are routed to the expert \(s\), multiplied by a gating value \(G_{j,s}(x)\), which are selected differently in different pMoE models.

Each expert is a two-layer CNN with the same architecture. Let \(m\) denote the total number of neurons in all the experts. Then each expert contains \((m/k)\) neurons. Let \(w_{r,s}\in\mathbb{R}^{d}\) and \(a_{r,s}\in\mathbb{R}\) denote the hidden layer and output layer weights for neuron \(r\) (\(r\in[m/k]\)) in expert \(s\) (\(s\in[k]\)), respectively. The activation function is the rectified linear unit (ReLU), where \(\textbf{ReLU}(z)=\text{max}(0,z)\).

Let \(\theta=\{a_{r,s},w_{r,s},w_{s},\forall s\in[k],\forall r\in[m/k]\}\) include all the trainable weights. The pMoE model denoted as \(f_{M}\), is defined as follows:

\[f_{M}(\theta,x)=\sum_{s=1}^{k}\sum_{r=1}^{\overline{k}}\frac{a_{r,s}}{l}\sum_ {j\in J_{s}(w_{s},x)}\textbf{ReLU}(\langle w_{r,s},x^{(j)}\rangle)G_{j,s}(w_{ s},x) \tag{1}\]

An illustration of (1) is given in Figure 2.

The learning problem solves the following empirical risk minimization problem with the logistic loss function,

\[\text{min}_{\theta}:\quad L(\theta)=\frac{1}{N}\!\!\sum_{i=1}^{N}\log\left(1+e ^{-y_{i}f_{M}(\theta,x_{i})}\right) \tag{2}\]

We consider two different training modes of pMoE, _Separate-training_ and _Joint-training_ of the routers and the experts. We also consider the conventional CNN architecture for comparison.

Figure 2: An illustration of the pMoE model in (1) with \(k=3,m=6,n=6\), and \(l=2\).
(I) **Separate-training pMoE**: Under the setup of the so-called _hard mixtures of experts_(Collobert et al., 2003; Ahmed et al., 2016; Gross et al., 2017), the router weights \(w_{s}\) are trained first and then fixed when training the weights of the experts. In this case, the gating values are set as

\[G_{j,s}(w_{s},x)\equiv 1,\ \forall j,s,x \tag{3}\]

We select \(k=2\) in this case to simplify the analysis.

(II) **Joint-training pMoE**: The routers and the experts are learned jointly, see, e.g., (Lepikhin et al., 2020; Riquelme et al., 2021; Fedus et al., 2022). Here, the gating values are softmax functions with

\[G_{j,s}(w_{s},x)=e^{g_{j,s}(x)}/(\sum_{i\in J_{s}(x)}g_{i,s}(x)) \tag{4}\]

(III) **CNN single-expert counterpart**: The conventional two-layer CNN with \(m\) neurons, denoted as \(f_{C}\), satisfies,

\[f_{C}(\theta,x)=\sum_{r=1}^{m}\!a_{r}\left(\frac{1}{n}\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!
**(II). Both the sample complexity and the required number of hidden nodes of pMoE reduce by a polynomial factor of \(n/l\) over CNN.** We prove that as long as \(l\), the number of patches per expert, is greater than a threshold (that decreases as the separation between class-discriminative and class-irrelevant patterns increases), the sample complexity and the required number of neurons of learning pMoE are \(\Omega(l^{8})\) and \(\Omega(l^{10})\) respectively. In contrast, the sample and model complexities of the CNN are \(\Omega(n^{8})\) and \(\Omega(n^{10})\) respectively, indicating improved generalization by pMoE.

**(III). Larger separation among class-discriminative and class-irrelevant patterns reduces the sample complexity and model complexity of pMoE.** Both the sample complexity and the required number of neurons of pMoE is polynomial in \(\delta\), which decreases when the separation among patterns increases.

### Data Model Assumptions and Rationale

The input \(x\) is comprised of one class-discriminative pattern and \(n-1\) class-irrelevant patterns, and the label \(y\) is determined by the class-discriminative pattern only.

**Distributions of class-discriminative patterns**: The unit vectors \(o_{1}\) and \(o_{2}\in\mathbb{R}^{d}\) denote the _class-discriminative_ patterns that determine the labels. The separation between \(o_{1}\) and \(o_{2}\) is measured as \(\delta_{d}:=\langle o_{1},o_{2}\rangle\in(-1,1)\). \(o_{1}\) and \(o_{2}\) are equally distributed in the samples, and each sample has exactly one of them. If \(x\) contains \(o_{1}\) (or \(o_{2}\)), then \(y\) is \(+1\) (or \(-1\)).

**Distributions of class-irrelevant patterns.**_Class-irrelevant_ patterns are unit vectors in \(\mathbb{R}^{d}\) belonging to \(p\) disjoint pattern sets \(S_{1},S_{2},....,S_{p}\), and these patterns distribute equally for both classes. \(\delta_{r}\) measures the separation between class-discriminative patterns and class-irrelevant patterns, where \(|\langle o_{i},q\rangle|\leq\delta_{r}\), \(\forall i\in[2]\), \(\forall q\in S_{j}\), \(j=1,...,p\). Each \(S_{j}\) belongs to a ball with a diameter of \(\Theta(\sqrt{(1-\delta_{r}^{2})/dp^{2}})\). Note that NO separation among class-irrelevant patterns themselves is required.

**The rationale of our data model.** The data distribution \(\mathcal{D}\) captures the locality of the label-defining features in image data. It is motivated by and extended from the data distributions in recent theoretical frameworks (Yu et al., 2019; Brutzkus & Globerson, 2021; Karp et al., 2021; Chen et al., 2022). Specifically, Yu et al. (2019) and Brutzkus & Globerson (2021) require orthogonal patterns, i.e., \(\delta_{r}\) and \(\delta_{d}\) are both \(0\), and there are only a fixed number of non-discriminative patterns. Karp et al. (2021) and Chen et al. (2022) assume that \(\delta_{d}=-1\) and a possibly infinite number of patterns drawn from zero-mean Gaussian distribution. In our model, \(\delta_{d}\) takes any value in \((-1,1)\), and the class-irrelevant patterns can be drawn from \(p\) pattern sets that contain an infinite number of patterns that are not necessarily Gaussian or orthogonal.

Define

\[\delta=1/(1-\max(\delta_{d}^{2},\delta_{r}^{2})) \tag{9}\]

\(\delta\) decreases if (1) \(o_{1}\) and \(o_{2}\) are more separated from each other, and (2) Both \(o_{1}\) and \(o_{2}\) are more separated from any set \(S_{i}\), \(i\in[p]\). We also define an integer \(l^{*}\) (\(l^{*}\leq n\)) that measures _the maximum number of class-irrelevant patterns per sample that are sufficiently closer to \(o_{1}\) than \(o_{2}\), and vice versa_. Specifically, a class-irrelevant pattern \(q\) is called \(\delta^{\prime}\)-closer (\(\delta^{\prime}>0\)) to \(o_{1}\) than \(o_{2}\), if \(\langle o_{1}-o_{2},q\rangle>\delta^{\prime}\) holds. Similarly, \(q\) is \(\delta^{\prime}\)-closer to \(o_{2}\) than \(o_{1}\) if \(\langle o_{2}-o_{1},q\rangle>\delta^{\prime}\). Then, let \(l^{*}-1\) be the maximum number of class-irrelevant patches that are either \(\delta^{\prime}\)-closer to \(o_{1}\) than \(o_{2}\) or vice versa with \(\delta^{\prime}=\Theta(1-\delta_{d})\) in any \(x\) sampled from \(\mathcal{D}\). \(l^{*}\) depends on \(\mathcal{D}\) and \(\delta_{d}\). When \(\mathcal{D}\) is fixed, a smaller \(\delta_{d}\) corresponds to a larger separation between \(o_{1}\) and \(o_{2}\) and leads to a small \(l^{*}\). In contrast to linearly separable data in (Yu et al., 2019; Brutzkus et al., 2018; Chen et al., 2022), our data model is **NOT** linearly separable as long as \(l^{*}=\Omega(1)\) (see section K in Appendix for the proof).

### Main Theoretical Results

#### 4.3.1 Generalization Guarantee of Separate-training pMoE

Lemma 4.1 shows that as long as the number of patches per expert, \(l\), is greater than \(l^{*}\), then the separately learned routers by solving (6) always send \(o_{1}\) to expert 1 and \(o_{2}\) to expert 2. Based on this discriminative property of the learned routers, Theorem 4.2 then quantifies the sample complexity and network size of separate-training pMoE to achieve a desired generalization error \(\epsilon\). Theorem 4.3 quantifies the sample and model complexities of CNN for comparison.

**Lemma 4.1** (Discriminative Property of Separately Trained Routers).: _For every \(l\geq l^{*}\), w.h.p. over the random initialization defined in (7), after doing mini-batch SGD with batch-size \(B_{r}=\Omega\left(n^{2}/(1-\delta_{d})^{2}\right)\) and learning rate \(\eta_{r}=\Theta(1/n)\), for \(T_{r}=\Omega\left(1/(1-\delta_{d})\right)\) iterations, the returned \(w_{1}\) and \(w_{2}\) satisfy_

\[\underset{j\in[n]}{\text{arg}}\left(x^{(j)}=o_{1}\right)\in J_{1}(w_{1},x), \quad\forall(x,y=+1)\sim\mathcal{D}\]

\[\underset{j\in[n]}{\text{arg}}\left(x^{(j)}=o_{2}\right)\in J_{2}(w_{2},x), \quad\forall(x,y=-1)\sim\mathcal{D}\]

_i.e., the learned routers always send \(o_{1}\) to expert 1 and \(o_{2}\) to expert 2._

The main idea in proving Lemma 4.1 is to show that the gradient in each iteration has a large component along the directions of \(o_{1}\) and \(o_{2}\). Then after enough iterations, the inner product of \(w_{1}\) and \(o_{1}\) (similarly, \(w_{2}\) and \(o_{2}\)) is sufficiently
large. The intuition of requiring \(l\geq l^{*}\) is that because there are at most \(l^{*}-1\) class-irrelevant patches sufficiently closer to \(o_{1}\) than \(o_{2}\) (or vice versa), then sending \(l\geq l^{*}\) patches to one expert will ensure that one of them is \(o_{1}\) (or \(o_{2}\)). Note that the batch size \(B_{r}\) and the number of iterations \(T_{r}\) depend on \(\delta_{d}\), the separation between \(o_{1}\) and \(o_{2}\), but are independent of the separation between class-discriminative and class-irrelevant patterns.

We then show that the separate-training pMoE reduces both the sample complexity and the required model size (Theorem 4.2) compared to the CNN (Theorem 4.3).

**Theorem 4.2** (Generalization guarantee of separate-training pMoE).: _For every \(\epsilon>0\) and \(l\geq l^{*}\), for every \(m\geq M_{S}=\Omega\left(l^{10}p^{12}\delta^{6}\big{/}\epsilon^{16}\right)\) with at least \(N_{S}=\Omega(l^{8}p^{12}\delta^{6}/\epsilon^{16})\) training samples, after performing minibatch SGD with the batch size \(B=\Omega\left(l^{4}p^{6}\delta^{3}\big{/}\epsilon^{8}\right)\) and the learning rate \(\eta=O\big{(}1/(mpoly(l,p,\delta,1/\epsilon,\log m))\big{)}\) for \(T=O\left(l^{4}p^{6}\delta^{3}\big{/}\epsilon^{8}\right)\) iterations, it holds w.h.p. that_

\[\mathop{\mathbb{P}}_{(x,y)\sim\mathcal{D}}\big{[}yf_{M}(\theta^{(T)},x)>0 \big{]}\geq 1-\epsilon\]

Theorem 4.2 implies that to achieve generalization error \(\epsilon\) by a separate-training pMoE, we need \(N_{S}=\Omega(l^{8}p^{12}\delta^{6}/\epsilon^{16})\) training samples and \(M_{S}=\Omega\left(l^{10}p^{12}\delta^{6}\big{/}\epsilon^{16}\right)\) hidden nodes. Therefore, both \(N_{S}\) and \(M_{S}\) increase polynomially with the number of patches \(l\) sent to each expert. Moreover, both \(N_{S}\) and \(M_{S}\) are polynomial in \(\delta\) defined in (9), indicating an improved generalization performance with stronger separation among patterns.

The proof of Theorem 4.2 is inspired by Li & Liang (2018), which analyzes the generalization performance of fully-connected neural networks (FCN) on structured data, but we have new technical contributions in analyzing pMoE models. In addition to analyzing the pMoE routers (Lemma 4.1), which do not appear in the FCN analysis, our analyses also significantly relax the separation requirement on the data, compared with that by Li & Liang (2018). For example, Li & Liang (2018) requires the separation between the two classes, measured by the smallest \(\ell_{2}\)-norm distance of two points in different classes, being \(\Omega(n)\) to obtain a sample complexity bound of poly(\(n\)) for the binary classification task. In contrast, the separation between the two classes in our data model is \(\min\{\sqrt{2(1-\delta_{d})},2\sqrt{1-\delta_{r}}\}\), much less than \(\Omega(n)\) required by Li & Liang (2018).

**Theorem 4.3** (Generalization guarantee of CNN).: _For every \(\epsilon>0\), for every \(m\geq M_{C}=\Omega\left(n^{10}p^{12}\delta^{6}\big{/}\epsilon^{16}\right)\) with at least \(N_{C}=\Omega(n^{8}p^{12}\delta^{6}\big{/}\epsilon^{16})\) training samples, after performing minibatch SGD with the batch size \(B=\Omega\left(n^{4}p^{6}\delta^{3}\big{/}\epsilon^{8}\right)\) and the learning rate \(\eta=O\big{(}1/(mpoly(n,p,\delta,1/\epsilon,\log m))\big{)}\) for \(T=O\left(n^{4}p^{6}\delta^{3}\big{/}\epsilon^{8}\right)\) iterations, it holds w.h.p. that_

\[\mathop{\mathbb{P}}_{(x,y)\sim\mathcal{D}}\big{[}yf_{C}(\theta^{(T)},x)>0 \big{]}\geq 1-\epsilon\]

Theorem 4.3 implies that to achieve a generalization error \(\epsilon\) using CNN in (5), we need \(N_{C}=\Omega(n^{8}p^{12}\delta^{6}/\epsilon^{16})\) training samples and \(M_{C}=\Omega\left(n^{10}p^{12}\delta^{6}\big{/}\epsilon^{16}\right)\) neurons.

**Sample-complexity gap between single CNN and mixture of CNNs.** From Theorem 4.2 and Theorem 4.3, the sample-complexity ratio of the CNN to the separate-training pMoE is \(N_{C}/N_{S}=\Theta\big{(}(n/l)^{8}\big{)}\). Similarly, the required number of neurons is reduced by a factor of \(M_{C}/M_{S}=\Theta\big{(}(n/l)^{10}\big{)}\) in separate-training pMoE4.

Footnote 4: The bounds for the sample complexity and model size in Theorem 4.2 and Theorem 4.3 are sufficient but not necessary. Thus, rigorously speaking, one can not compare sufficient conditions only. In our analysis, however, the bounds for MoE and CNN are derived with exactly the same technique with the only difference to handle the routers. Therefore, it is fair to compare these two bounds to show the advantage of pMoE.

#### 4.3.2 Generalization Guarantee of

Joint-training pMoE with Proper Routers

Theorem 4.5 characterizes the generalization performance of joint-training pMoE assuming the routers are properly trained in the sense that after some SGD iterations, for each class at least one of the \(k\) experts receives all class-discriminative patches of that class with the largest gating-value (see Assumption 4.4).

**Assumption 4.4**.: There exists an integer \(T^{\prime}<T\) such that for all \(t\geq T^{\prime}\), it holds that:

There exists an expert \(s\in[k]\) s.t. \(\forall(x,y=+1)\sim\mathcal{D},\)

\[j_{o_{1}}\in J_{s}(w_{s}^{(t)},x),\text{ and }G_{j_{o_{1}},s}^{(t)}(x) \geq G_{j,s}^{(t)}(x)\]

and an expert \(s\in[k]\) s.t. \(\forall(x,y=-1)\sim\mathcal{D},\)

\[j_{o_{2}}\in J_{s}(w_{s}^{(t)},x),\text{ and }G_{j_{o_{2}},s}^{(t)}(x) \geq G_{j,s}^{(t)}(x)\]

where \(j_{o_{1}}\) (\(j_{o_{2}}\)) denotes the index of the class-discriminative pattern \(o_{1}\) (\(o_{2}\)), \(G_{j,s}^{(t)}(x)\) is the gating output of patch \(j\in J_{s}(w_{s}^{(t)},x)\) of sample \(x\) for expert \(s\) at the iteration \(t\), and \(w_{s}^{(t)}\) is the gating kernel for expert \(s\) at iteration \(t\).

Assumption 4.4 is required in proving Theorem 4.5 because of the difficulty of tracking the dynamics of the routers in joint-training pMoE. Assumption 4.4 is verified on empirical experiments in Section 5.1, while its theoretical proof is left for future work.

**Theorem 4.5** (Generalization guarantee of joint-training pMoE).: _Suppose Assumption 4.4 hold. Then for every \(\epsilon>0\), for every \(m\geq M_{J}=\Omega\left(k^{3}n^{2}l^{6}p^{12}\delta^{6}\big{/}\epsilon^{16}\right)\) with at least \(N_{J}=\Omega(k^{4}l^{6}p^{12}\delta^{6}/\epsilon^{16})\) training samples, after performing minibatch SGD with the batch size \(B=\Omega\left(k^{2}l^{4}p^{6}\delta^{3}\big{/}\epsilon^{8}\right)\) and the learning rate
\(\eta=O\big{(}1/(mppoly(l,p,\delta,1/\epsilon,\log m))\big{)}\) for \(T=O\left(k^{2}l^{2}p^{6}\delta^{3}/\epsilon^{8}\right)\) iterations, it holds w.h.p. that

\[\mathop{\mathbb{P}}\limits_{(x,y)\sim\mathcal{D}}\big{[}yf_{M}(\theta^{(T)},x)> 0\big{]}\geq 1-\epsilon\]

Theorem 4.5 indicates that, with proper routers, joint-training pMoE needs \(N_{J}=\Omega(k^{4}l^{6}p^{12}\delta^{6}/\epsilon^{16})\) training samples and \(M_{J}=\Omega\left(k^{3}n^{2}l^{6}p^{12}\delta^{6}/\epsilon^{16}\right)\) neurons to achieve \(\epsilon\) generalization error. Compared with CNN in Theorem 4.3, joint-training pMoE reduces the sample complexity and model size by a factor of \(\Theta(n^{8}/k^{4}l^{6})\) and \(\Theta(n^{10}/k^{3}l^{6})\), respectively. With more experts (a larger \(k\)), it is easier to satisfy Assumption 4.4 to learn proper routers but requires larger sample and model complexities. When the number of samples is fixed, the expression of \(N_{J}\) also indicates that \(\epsilon\) sales as \(k^{1/4}l^{3/8}\), corresponding to an improved generalization when \(k\) and \(l\) decrease.

We provide the end-to-end computational complexity comparison between the analyzed pMoE models and general CNN model in Table 1 (see section N in Appendix for details). The results in Table 1 indicates that the computational complexity in joint-training pMoE is reduced by a factor of \(O(n^{5}/k^{2}l^{3})\) compared with CNN. Similarly, the reduction of computational complexity of separate-training pMoE is \(O(n^{5}/l^{5})\).

## 5 Experimental Results

### pMoE of Two-layer CNN

**Dataset**: We verify our theoretical findings about the model in (1) on synthetic data prepared from MNIST (LeCun et al., 2010) data set. Each sample contains \(n=16\) patches with patch size \(d=28\times 28\). Each patch is drawn from the MNIST dataset. See Figure 3 as an example. We treat the digits "**I**" and "**0**" as the class-discriminative patterns \(o_{1}\) and \(o_{2}\), respectively. Each of the digits from "**2**" to "**9**" represents a class-irrelevant pattern set.

**Setup**: We compare separate-training pMoE, joint-training pMoE, and CNN with similar model sizes. The separate-training pMoE contains _two_ experts with \(20\) hidden nodes in each expert. The joint-training pMoE has eight experts with five hidden nodes per expert. The CNN has \(40\) hidden nodes. All are trained using SGD with \(\eta=0.2\) until zero training error. pMoE converges much faster than CNN, which takes \(150\) epochs. Before training the experts in the separate-training pMoE, we train the router for \(100\) epochs. The models are evaluated on \(1000\) test samples.

**Generalization performance**: Figure 4 compares the test accuracy of the three models, where \(l=2\) and \(l=6\) for separate-training and joint-training pMoE, respectively. The error bars show the mean plus/minus one standard deviation

\begin{table}
\begin{tabular}{l|c|c|c} \hline \hline \multirow{2}{*}{**Complexity to**} & \multicolumn{3}{c|}{**pMoE**} & \multirow{2}{*}{**CNN**} \\ \cline{2-3} \cline{5-5}  & **Separate-training** & & **Joint-training** \\ \cline{2-3} \cline{5-5}  & \(O(Bml^{5}d/\epsilon^{8})\) & \(O(Bmk^{2}l^{3}d/\epsilon^{8})\) & \(O(Bmn^{5}d/\epsilon^{8})\) \\ \hline \multirow{2}{*}{
\begin{tabular}{l} Complexity per \\ Iteration \\ (**CompLYS/Iter**) \\ \end{tabular} } & \multirow{2}{*}{\(O(Bmld)\)} & **Router** & **Expert** \\ \cline{3-3} \cline{5-5}  & & \(O(Bknd)\) (Forward pass) & \multirow{2}{*}{\(O(Bmld)\)} \\ \cline{3-3} \cline{5-5}  & & \(O(Bkl^{2}d)\) (Backward pass) & \\ \hline Iteration required & & & \\ to converge with \(\epsilon\) & & & \\ error (**T**) & & & \\ \hline \hline \end{tabular}
\end{table}
Table 1: Computational complexity of pMoE and CNN.

Figure 4: Generalization performance of pMoE and CNN with a similar model size

Figure 5: Phase transition of sample complexity with \(l\) in separate-training pMoE

Figure 3: Sample image of the synthetic data from MNIST. Class label is "1".
of five independent experiments. pMoE outperforms CNN with the same number of training samples. pMoE only requires 60% of the training samples needed by CNN to achieve \(95\%\) test accuracy.

Figure 5 shows the sample complexity of separate-training pMoE with respect to \(l\). Each block represents 20 independent trials. A white block indicates all success, and a black block indicates all failure. The sample complexity is polynomial in \(l\), verifying Theorem 4.2. Figure 7 and 6 show the test accuracy of joint-training pMoE with a fixed sample size when \(l\) and \(k\) change, respectively. When \(l\) is greater than \(l^{*}\), which is \(6\) in Figure 7, the test accuracy matches our predicted order. Similarly, the dependence on \(k\) also matches our prediction, when \(k\) is large enough to make Assumption 4.4 hold.

**Router performance**: Figure 8 verifies the discriminative property of separately trained routers (Lemma 4.1) by showing the percentage of testing data that have class-discriminative patterns (\(o_{1}\) and \(o_{2}\)) in top \(l\) patches of the separately trained router. With very few training samples (such as \(300\)), one can already learn a proper router that has discriminative patterns in top-\(4\) patches for 95% of data. Figure 9 verifies the discriminative property of jointly trained routers (Assumption 4.4). With only \(300\) training samples, the jointly trained router dispatches \(o_{1}\) with the largest gating value to a particular expert for 95% of class-1 data and similarly for \(o_{2}\) in 92% of class-2 data.

### pMoE of Wide Residual Networks (WRNs)

**Neural network model**: We employ the 10-layer WRN (Zagoruyko & Komodakis, 2016) with a widening factor of 10 as the expert. We construct a patch-level MoE counterpart of WRN, referred to as WRN-pMoE, by replacing the last convolutional layer of WRN with an pMoE layer of an equal number of trainable parameters (see Figure 18 in Appendix for an illustration). WRN-pMoE is trained with the joint-training method5. All the results are averaged over five independent experiments.

Footnote 5: Code is available at [https://github.com/nowazrabbani/pMoE_CNN](https://github.com/nowazrabbani/pMoE_CNN)

**Datasets**: We consider both CelebA (Liu et al., 2015) and CIFAR-10 datasets. The experiments on CIFAR-10 are deferred to the Appendix (see section A). We down-sample the images of CelebA to \(64\times 64\). The last convolutional layer of WRN receives a (\(16\times 16\times 640\)) dimensional feature map. The feature map is divided into \(16\) patches with size \(4\times 4\times 640\) in WRN-pMoE. \(k=8\) and \(l=2\) for the pMoE layer.

Figure 8: Percentage of properly routed discriminative patterns by a separately trained router.

Figure 6: Change of test accuracy in joint-training pMoE with \(k\) for fixed sample sizes

Figure 7: Change of test accuracy in joint-training pMoE with \(l\) for fixed sample sizes

Figure 9: Percentage of properly routed discriminative patterns by a jointly trained router. \(l=6\).
**Performance Comparison**: Figure 11 shows the test accuracy of the binary classification problem on the attribute "smiling." WRN-pMoE requires less than _one-fifth_ of the training samples needed by WRN to achieve 86% accuracy. Figure 11 shows the performance when the training data contain spurious correlations with the hair color as a spurious attribute. Specifically, 95% of the training images with the attribute "smiling" also have the attribute "black hair," while 95% of the training images with the attribute "not-smiling" have the attribute "blond hair." The models may learn the hair-color attribute rather than "smiling" due to spurious correlation and, thus, the test accuracies are lower in Figure 11 than those in Figure 11. Nevertheless, WRN-pMoE outperforms WRN and reduces the sample complexity to achieve the same accuracy.

Figure 12 shows the test accuracy of multiclass classification (four classes with class attributes: "Not smiling, Eyeglass," "Smiling, Eyeglass," "Smiling, No eyeglass," and "Not smiling, No eyeglass") in CelebA. The results are consistent with the binary classification results. Furthermore, Table 2 empirically verifies the computational efficiency of WRN-pMoE over WRN on multiclass classification in CelebA6. Even with same number of training samples, WRN-pMoE is still more computationally efficient than WRN, because WRN-pMoE requires fewer iterations to converge and has a lower per-iteration cost.

Footnote 6: An NVIDIA RTX 4500 GPU was used to run the experiments, training FLOPs are calculated as Training FLOPs \(=\) Training time (second) \(\times\) Number of GPUs \(\times\) peak FLOP/second \(\times\) GPU utilization rate

## 6 Conclusion

MoE reduces computational costs significantly without hurting the generalization performance in various empirical studies, but the theoretical explanation is mostly elusive. This paper provides the first theoretical analysis of patch-level MoE and proves its savings in sample complexity and model size quantitatively compared with the single-expert counterpart. Although centered on a classification task using a mixture of two-layer CNNs, our theoretical insights are verified empirically on deep architectures and multiple datasets. Future works include analyzing other MoE architectures such as MoE in Vision Transformer (ViT) and connecting MoE with other sparsification methods to further reduce the computation.

## Acknowledgements

This work was supported by AFOSR FA9550-20-1-0122, NSF 1932196 and the Rensselaer-IBM AI Research Collaboration ([http://airc.rpi.edu](http://airc.rpi.edu)), part of the IBM AI Horizons Network ([http://ibm.biz/AIHorizons](http://ibm.biz/AIHorizons)). We thank Yihua Zhang at Michigan State University for the help in experiments with CelebA dataset. We thank all anonymous reviewers.