Patch-level Routing in Mixture-of-Experts is Provably Sample-efficient for Convolutional Neural Networks

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:

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

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

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.

Let $\tilde{\theta}$ denote the parameters of the learned model by solving (1). The predicted label for a test sample $x$ by the learned model is $\textrm{sign}(f_{M}(\tilde{\theta},x))$. The generalization accuracy, i.e., the fraction of correct predictions of all test samples equals $\underset{(x,y)\sim\mathcal{D}}{\mathbb{P}}\left[yf_{M}(\theta,x)>0\right]$. This paper studies both separate and joint training of pMoE and compares their performance with CNN, from the perspective of sample complexity to achieve a desirable generalization accuracy.

In the following algorithms, we fix the output layer weights $a_{r,s}$ and $a_{r}$ at their initial values randomly sampled from the standard Gaussian distribution $\mathcal{N}(0,1)$ and do not update them during the training. This is a typical simplification when analyzing NN, as used in (Li & Liang, 2018; Brutzkus et al., 2018; Allen-Zhu et al., 2019a; Arora et al., 2019).

(I) Separate-training pMoE: The routers are separately trained using $N_{r}$ training samples ($N_{r}<N$), denoted by $\{(x_{i},y_{i})\}_{i=1}^{N_{r}}$ without loss of generality. The gating kernels $w_{1}$ and $w_{2}$ are obtained by solving the following minimization problem:

To solve (6), we implement the mini-batch SGD with batch size $B_{r}$ for $T_{r}=N_{r}/B_{r}$ iterations, starting from the random initialization as follows:

where, $\sigma_{r}=\Theta\big{(}1\big{/}(n^{2}\log{(\textrm{poly}(n))}\sqrt{d})\big{)}$.

After learning the routers, we train the hidden-layer weights $w_{r,s}$ by solving (2) while fixing $w_{1}$ and $w_{2}$. We implement mini-batch SGD of batch size $B$ for $T=N/B$ iterations starting from the initialization

(II) Joint-training pMoE: $w_{s}$ and $w_{r,s}$ in (1) are updated simultaneously by mini-batch SGD of batch size $B$ for $T=N/B$ iterations starting from the initialization in (7) and (8).

(III) CNN: $w_{r}$ in (5) are updated by mini-batch SGD of batch size $B$ for $T=N/B$ iterations starting from the initialization in (8).

Before defining the data model assumptions and rationale in Section 4.2 and presenting the formal results in 4.3, we first summarize our key findings. We assume that the data patches are sampled from either class-discriminative patterns that determine the labels or a possibly infinite number of class-irrelevant patterns that have no impact on the label. The parameter $\delta$ (defined in (9)) is inversely related to the separation among patterns, i.e., $\delta$ decreases when (i) the separation among class-discriminative patterns increases, and/or (ii) the separation between class-discriminative and class-irrelevant patterns increases. The key findings are as follows.

(I). A properly trained patch-level router sends class-discriminative patches of one class to the same expert while dropping some class-irrelevant patches. We prove that separate-training pMoE routes class-discriminative patches of the class with label $y=+1$ (or the class with label $y=-1$) to the expert 1 (or the expert 2) respectively, and the class-irrelevant patterns that are sufficiently away from class-discriminative patterns are not routed to any expert (Lemma 4.1). This discriminative routing property is also verified empirically for joint-training pMoE (see section 5.1). Therefore, pMoE effectively reduces the interference by irrelevant patches when each expert learns the class-discriminative patterns. Moreover, we show empirically that pMoE can remove class-irrelevant patches that are spuriously correlated with class labels and thus can avoid learning from spuriously correlated features of the data.

(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.

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$ 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).

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 4 as an example. We treat the digits “1” 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 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.

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 method⁵⁵5Code is available at https://github.com/nowazrabbani/pMoE_CNN. All the results are averaged over five independent experiments.

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.

Performance Comparison: Figure 12 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 12 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 12 than those in Figure 12. 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 CelebA⁶⁶6An NVIDIA RTX 4500 GPU was used to run the experiments, training FLOPs are calculated as $\text{Training FLOPs}=\text{Training time (second)}\times\text{Number of GPUs}\times\text{peak FLOP/second}\times\text{GPU utilization rate}$. 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.