Transformer-based models are composed of multiple blocks, where each block consists of multi-head attention (MHA) and feed-forward layers (FFLs) (Vaswani et al., 2017). MoEs could thus be applied to either MHA or FFLs, or both. In this work, we only explore MoE FFLs in the design space; this is primarily due to the runtime overhead introduced by dynamic behavior, which we found to be prohibitively high for the already expensive attention layers. planer’s first phase explores the large design space composed of different configurations of MHAs, FFLs, and MoE layers. The inputs to the first phase are the design space, the backbone of the baseline network architecture, and a target latency, expressed as a ratio w.r.t. the baseline latency.

For real-world networks, the design space of alternative architectures often gets prohibitively large; for instance, the Transformer-XL Base network on the enwik8 dataset yields a search space size of over 68 billion architectures. To keep the search tractable, we deploy a differentiable NAS strategy, which has been shown to be significantly more efficient than reinforcement-learning-based approaches (Zoph and Le, 2016). We follow a NAS algorithm similar to the one proposed by Wu et al. (2019).

Phase 1 first composes a search architecture using the baseline network’s backbone as depicted in Figure 7. The backbone includes details on the number of blocks (MHA or FFLs) and their configuration (number of heads or hidden dimension). Using the input backbone, each of the MHA or FFL blocks in the baseline network are replaced with Super Blocks (SB), which includes all the search options in the design space. The goal is to find the best option for each block so that overall accuracy is maximized and the latency target is achieved. Figure 8 depicts the formulation of super blocks. Each of the search options ${Block}_i$ is accompanied by corresponding architectural weights ${\alpha }_i$, which are trained using gradient descent to represent the benefit factor of the search option Wu et al. (2019). To make the optimization graph differentiable with respect to the architecture weights, the output of the super block is formulated as:

Where the $GumbelSoftmax$ generates probability values by sampling the Gumbel distribution based on $\alpha$ weights.

This formulation yields two sets of parameters to be trained in Phase 1. The first group contains the actual network weights (${Block}_i$), and the second group the architectural weights (${\alpha }_i$). Training of each parameter group is done sequentially in each epoch, using separate optimizers. Thus each epoch of training in phase 1 consists of optimizing the network weights using $100\%$ of the training samples, and then training the architecture weights using $20\%$ of the randomly sampled training data. We use soft sampling for $GumbelSoftmax$ during architecture optimization, and hard-sampling while training the network weights to reduce the overheads associated with the super blocks. To ensure that neither of the network weight sets are starved due to the hard-sampling of $GumbelSoftmax$, the architecture optimization is initially disabled for $10\%$ of the epochs, and an annealing temperature scheduling is used for later epochs. These settings allow the blocks to be randomly sampled for the appropriate number of search epochs.

To incorporate latency optimization in the search phase, we formulate an auxiliary loss based on the latencies of the search and baseline network, as well as the target latency. We use an estimation for the end-to-end latency of the search network as well as baseline in phase 1, using lookup tables filled with individual block latencies similar to prior work Wu et al. (2019). Equation (2) presents the formulation for the estimated latency which is composed of accumulating the latencies of each super block ($Lat\_SB$).

Here, ${Lat}_i$ represents the profiled latency of $Bloc{k}_i$ in isolation, and $P_{bi}$ values correspond to the probability values for super block of $b$ as sampled in Equation (1) with respect to the architecture weights.

The latency loss $La{t}_{Loss}$ is implemented as the ratio of the estimated latency of the search network ($Lat$) over the normalized baseline latency with respect to the target latency.

During the training of the architecture weights, the latency loss will be automatically activated depending on whether the estimated latency of the search network is meeting the target latency requirement. For example, if the target latency is set to $50\%$ of the baseline, the latency loss will only get included if the estimated latency is higher than $0.5\times La{t}_{Baseline}$. Otherwise, the scalar factor of $\beta$ would be 0 in Equation 3, leading the optimizer to adjust the architecture weights solely in the direction of minimizing the $C{E}_{Loss}$. This novel dynamic functionality helps the search progress towards the user latency target without the need for additional hyper-parameter tuning.

The optimized architecture obtained from Phase 1 is now instantiated for retraining. Since the weights of this final architecture were shared with other search points during Phase 1, a retraining step is necessary to avoid under-fitting and to obtain optimal accuracy. We construct the optimized architecture by selecting the blocks with the highest architecture weight values in each super block; from our empirical evaluation, this sampling strategy best balances additional training overheads with accuracy compared to other approaches such as the one described in  Liu et al. (2018). We retrain the sampled architecture from scratch using the same settings as the baseline.

Since MoE blocks may be part of the final architecture, we incorporate an auxiliary loss during Phase 2 to enforce a balanced load across the experts. We follow the same implementation of the auxiliary loss for load balancing ($Balanc{e}_{Loss}$) as Switch Transformer Fedus et al. (2021). Consider an MoE layer with $E$ experts:

Here, $F_e$ represents the fraction of the tokens processed by expert $e$, and $G_e$ measures the average gate score received by expert $e$ across the input tokens.

The $Balanc{e}_{Loss}$ provides an approximation for the load balancing score across experts. If the tokens are distributed uniformly across the experts by the gate function, we can expect each expert to process $\frac{1}{E}$ of the input tokens, while receiving an average score of $\frac{1}{E}$ from the gate. This would result in $Balanc{e}_{Loss}$ having an ideal value of $1$ in a fully-uniform distribution of tokens across the experts. If there is more than one MoE layer in the architecture, the $Balanc{e}_{Loss}$ is the average of the individual loss values across the MoE layers.

Figure (a)a compares the Phase 2 training progress of a Transformer-XL architecture with multiple MoE layers under two scenarios: (1) when the $Balanc{e}_{Loss}$ term is excluded from the loss function (Relaxed Load Balancing), and (2) when the loss function includes the $Balanc{e}_{Loss}$ term (Enforced Load Balancing). From the figure, we notice that trends for the ${CE}_{Loss}$ term are similar in both scenarios, highlighting the fact that overall accuracy of the network is unaffected by load balancing constraints. From our experiments, we also notice that a balanced load improves the runtime of MoE layers by reducing tail latency - we illustrate this in Figure (b)b. Here, we notice a runtime speedup of up to $1.16\times$ for MoE layers when load balancing is enforced.

We use Transformer-XL (TXL) Base on the WikiText-103 (WT103) and enwik8 datasets as our baseline networks. The backbone architecture for both datasets uses a model dimension of 512 and an interleaved pattern of multi-head attention (MHA) with 8 heads and feed-forward layer (FFLs) with an inner dimension of 2048. The total number of blocks (MHA/FFL) is 24 and 32 for enwik8 and WT103, respectively ¹¹1The number of MHA/FFL blocks is $2\times$ of the number of Transformer blocks.. The search space for phase 1 includes: (1) Skip connection, (2) MHA with 1, 2, 4, or 8 heads, (3) FFL with inner dimension of 2048, and (4) MoE FFL with inner dimension of 2048, 8 experts, where each token is processed by either 1 or 2 experts ($To{p}_K=1or2$).

To evaluate the performance of planer, we compare the latency and accuracy of the optimized models with the baseline TXL model and two prior papers: Sandwich Transformer Press et al. (2019) and PAR Transformer Mandava et al. (2020). The design space is explored using planer’s 2-phase methodology (described in more detail in Section 3) with target latencies ranging from $50\%$ to $95\%$.

All training is performed on a node with 8 NVIDIA V100 GPUs. We use the settings published by NVIDIA for hyper-parameters NVIDIA . The exact hyper-parameters used for each dataset are:

• WikiText-103 - Network Weights (Phase 1 and 2): JITLamb optimizer, learning rate of 0.01, batch size of 256, target and memory length of 192, dropout rate of 0.1 for non-MoE layers and 0.2 for MoE layers, and 40000 iterations.

• WikiText-103 - Architecture Weights (Phase 1): Adam optimizer, learning rate of 0.01, initial temperature of 5 for the Gumbel Softmax, and temperature annealing rate of 0.6.

• enwik8 - Network Weights (Phase 1 and 2): JITLamb optimizer, learning rate of 0.004, batch size of 64, target and memory length of 512, dropout rate of 0.1 for non-MoE layers and 0.3 for MoE layers, and 120000 iterations.

• enwik8 - Architecture Weights (Phase 1): Adam optimizer, learning rate of 0.01, initial temperature of 5 for the Gumbel Softmax, and temperature annealing rate of 0.7.

Table 1 lists the accuracy numbers obtained by planer and compares them with the baseline architectures. We notice that all the TXL variants, including ones produced by planer, maintain baseline accuracy levels. We provide a detailed comparison of the different architectures in Appendix A.

Figure 14 shows the speedups obtained by planer and the various baselines (described in Section 4.1) across both datasets and varying batch sizes. From the Figure, we notice that planer provides speedups of over $2\times$ at larger batch sizes. On smaller batch sizes, PAR Transformer outperforms planer; this is primarily due to the unoptimized MoE layers used in our current implementation. Specifically, our current implementation computes the outputs of each expert sequentially, where a batch of sequences with $N$ tokens are sequentially processed in mini-batches of size $\frac{To{p}_K\times N}{Experts}$. This consequently leads to under-utilization of the compute units.

Figure 15 provides a more detailed overview of the current deficiencies in the sequential implementation of the MoE layers; here, we provide a runtime comparison of the FFL, MHA, and MoE layers across different batch sizes normalized with respect to FFL runtime. At lower batch sizes, MoE layers have an overhead of $7\times$ over FFL, which is also higher than the MHA layers. However, as batch size increases, GPU resource utilization goes up, consequently decreasing the overhead of MoE layers to less than $3\times$. The oracle implementation (dashed orange line in Figure) shows the theoretically optimal runtime of the MoE layer. Since we use a $To{p}_k$ value of $2$ (viz., each input token is processed by 2 experts), we notice a corresponding $2\times$ runtime overhead over the baseline FFL. Note that the oracle runtime does not take overheads related to gate function evaluation and the gathering/scattering of tokens across experts into account - the real-world runtime is thus likely to be higher. We are currently working on a more optimized parallel implementation of MoE layers, which will help plug this performance gap across various batch sizes.

We also compare planer to an iso-parameter setup, which replaces the MoE with a scaled FFL in the search space. The scaled FFL has an inner dimension of 16384, which results in the same number of parameters as the MoE with 8 experts. The goal of the iso-parameter experiment is to analyze the effectiveness of different model scaling solutions in compensating for accuracy drops caused by aggressive attention pruning.

Figure 16 presents the comparison of the Pareto frontiers of the architectures obtained by planer with different latency targets on the WikiText-103 dataset. From the Figure, we clearly notice that the use of MoE layers results in higher performance architectures across the board at different accuracy levels. Further performance benchmarking reveals that scaled FFL layers are at least $2\times$ slower than our (relatively unoptimized) MoE layers and actually approach the runtime of the much slower MHA layers with 8 heads. Naively scaling up the size of FFLs is thus not an ideal option for either improving accuracy or performance.

In this section, we analyze the performance of the dynamic latency loss used in Phase 1. Figure (a)a shows the correlation between input target latency and the estimated latency of the architectures sampled at the end of Phase 1, while Figure (b)b shows the correlation between estimated latency and profiled end-to-end latency. We make two important observations from the figures: (1) our dynamic latency loss formulation successfully steers the NAS towards architectures that match the input target latency, and (2) the latency estimated in Equation (2) is highly correlated with real-world latency, making it an appropriate option for planer’s Phase 1 search.

To evaluate and validate the reproducibility of our experiments, and observe any potential variations in the final architectures, we also repeat the planer optimization of the architectures evaluated in Section 4.2. For this experiment, we keep all hyper-parameters fixed, but repeat planer’s search process four times. Figure 22 presents the achieved accuracy and speedup numbers from our experiment. We notice that all the accuracy values are within $0.5\%$ of the baseline, with speedups consistently over $2\times$. The variations in the final architectures across the two datasets are presented in Appendix B. Although the architectures do not match exactly, we notice a strong similarity in the number of heads in the attention layers. We also noticed that MoE layers tend to be concentrated towards the end of the networks across both datasets.