Original: S PARSE U PCYCLING : T RAINING M IXTURE - OF -E XPERTS FROM D ENSE C HECKPOINTS
Revised: S PARSE U PCYCLING : T RAINING M IXTURE - OF -E XPERTS FROM D ENSE C HECKPOINTS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Training large, deep neural networks to convergence can be prohibitively expensive.
Revised: Training large, deep neural networks to convergence can be prohibitively expensive.

Original: As a result, often only a small selection of popular, dense models are reused across different contexts and tasks.
Revised: As a result, often only a small selection of popular, dense models are reused across different contexts and tasks.

Original: Increasingly, sparsely activated models, which seek to decouple model size from computation costs, are becoming an attractive alternative to dense models.
Revised: Increasingly, sparsely activated models, which seek to decouple model size from computation costs, are becoming an attractive alternative to dense models.

Original: Although more efﬁcient in terms of quality and computation cost, sparse models remain data-hungry and costly to train from scratch in the large scale regime.
Revised: Although more efﬁcient in terms of quality and computation cost, sparse models remain data-hungry and costly to train from scratch in the large scale regime.

Original: In this work, we propose sparse upcycling – a simple way to reuse sunk training costs by initializing a sparsely activated Mixture-of-Experts model from a dense checkpoint.
Revised: In this work, we propose sparse upcycling – a simple way to reuse sunk training costs by initializing a sparsely activated Mixture-of-Experts model from a dense checkpoint.

Original: We show that sparsely upcycled T5 Base, Large, and
Revised: We show that sparsely upcycled T5 Base, Large, and

Original: XL language models and Vision Transformer Base and Large models, respectively, signiﬁcantly outperform their dense counterparts on SuperGLUE and ImageNet, using only ∼ 50% of the initial dense pretraining sunk cost.
Revised: XL language models and Vision Transformer Base and Large models, respectively, signiﬁcantly outperform their dense counterparts on SuperGLUE and ImageNet, using only ∼ 50% of the initial dense pretraining sunk cost.

Original: The upcycled models also outperform sparse models trained from scratch on 100% of the initial dense pretraining computation budget.
Revised: The upcycled models also outperform sparse models trained from scratch on 100% of the initial dense pretraining computation budget.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Increased scale is one of the main drivers of better performance in deep learning.
Revised: Increased scale is one of the main drivers of better performance in deep learning.

Original: From BERT (Devlin et al., 2019) to GPT-3 (Brown et al., 2020) to PaLM (Chowdhery et al., 2022) in natural language processing, or from AlexNet (Krizhevsky et al., 2017) to ViT-G (Zhai et al., 2022) in vision, breakthroughs in performance have been obtained from larger hardware, datasets, and architectures.
Revised: From BERT (Devlin et al., 2019) to GPT-3 (Brown et al., 2020) to PaLM (Chowdhery et al., 2022) in natural language processing, or from AlexNet (Krizhevsky et al., 2017) to ViT-G (Zhai et al., 2022) in vision, breakthroughs in performance have been obtained from larger hardware, datasets, and architectures.

Original: This trend holds true in many other domains too, including speech (Baevski et al., 2020), reinforcement learning (Schrittwieser et al., 2020), multimodal learning (Yu et al., 2022), and scientiﬁc applications of deep learning (Jumper et al., 2021).
Revised: This trend holds true in many other domains too, including speech (Baevski et al., 2020), reinforcement learning (Schrittwieser et al., 2020), multimodal learning (Yu et al., 2022), and scientiﬁc applications of deep learning (Jumper et al., 2021).

Original: However, most state-of-the-art neural networks are trained from-scratch; that is, starting from randomly initialized weights.
Revised: However, most state-of-the-art neural networks are trained from-scratch; that is, starting from randomly initialized weights.

Original: The cost for training such networks is growing rapidly.
Revised: The cost for training such networks is growing rapidly.

Original: For example, in language, BERT-Large (345M parameters, proposed in 2018) required an estimated 0 .
Revised: For example, in language, BERT-Large (345M parameters, proposed in 2018) required an estimated 0 .

Original: 5 ZFLOPS to train, while GPT-3 (175B parameters, from 2020) required 314 ZFLOPS (Brown et al., 2020), and PaLM (540B parameters, from 2022) required 2527 ZFLOPS (Chowdhery et al., 2022).
Revised: 5 ZFLOPS to train, while GPT-3 (175B parameters, from 2020) required 314 ZFLOPS (Brown et al., 2020), and PaLM (540B parameters, from 2022) required 2527 ZFLOPS (Chowdhery et al., 2022).

Original: As a result of these computation costs, research into new large language models is often limited to a small number of teams with access to lots of resources.
Revised: As a result of these computation costs, research into new large language models is often limited to a small number of teams with access to lots of resources.

Original: To enable signiﬁcant further progress, we must develop cheaper ways of training giant models.
Revised: To enable signiﬁcant further progress, we must develop cheaper ways of training giant models.

Original: In this paper, we explore model upcycling : upgrading an existing model with a relatively small additional computational budget.
Revised: In this paper, we explore model upcycling : upgrading an existing model with a relatively small additional computational budget.

Original: In particular, we focus on upcycling dense models into larger, sparsely activated Mixture-of-Experts (MoEs).
Revised: In particular, we focus on upcycling dense models into larger, sparsely activated Mixture-of-Experts (MoEs).

Original: We do not use any new unique sources of data (Wei et al., 2021; Ouyang et al., 2022).
Revised: We do not use any new unique sources of data (Wei et al., 2021; Ouyang et al., 2022).

Original: We assume the existence of a pretrained dense Transformer checkpoint (e.g. (Wolf et al., 2020)), that we then use to warm-start the training of a MoE. The idea is that, in this way we can obtain a high-performing MoE with relatively little cost; i.e. leverage the additional capacity of  the MoE to obtain a model more performant than the original Transformer, at a smaller cost than was used to train it.
Revised: We assume the existence of a pretrained dense Transformer checkpoint (e.g. (Wolf et al., 2020)), that we then use to warm-start the training of a MoE. By leveraging the additional capacity of from the MoE layers, we obtain an MoE model more performant than the original model, at a smaller cost than was used to train the original model.

Original: Across all model sizes  we study for  language and vision, with less than 40% additional budget, upcycling improves the network’s performance beyond what would be achieved by continued training the original Transformer.
Revised: Across all model sizes that we study for both language and vision, with less than 40% additional budget, upcycling improves the network’s performance beyond what would be achieved by continued training the original Transformer model.

Original: Sparse upcycling may be particularly valuable in two scenarios: (i) One has access to a pretrained
Revised: Sparse upcycling may be particularly valuable in two scenarios: (i) One has access to a pretrained

Original: Transformer (there are many publicly available) and wants to improve it with a modest  computational budget.
Revised: Transformer (there are many publicly available) and wants to improve it with a modest or constrained computational budget.

Original: (ii) One is planning to train a large model, and do not know whether a denseTransformer or MoE  would be more effective (the latter often being more performant, but more technically challenging to train): one can have both by ﬁrst training the dense model, then upcycling it into an
Revised: (ii) One is planning to train a large model, and do not know whether a dense or MoE model would be more effective (the latter often being more performant, but more technically challenging to train): one can have both by ﬁrst training the dense model, then upcycling it into an

Original: MoE  once it saturates.
Revised: MoE model once the dense model saturates.

Original: A central challenge in model upcycling is overcoming the initial performance decrease entailed by changing a trained network’s structure.
Revised: A central challenge in model upcycling is overcoming the initial performance decrease entailed by changing a trained network’s structure.

Original: We present a model surgery recipe that is effective in both vision and language, and numerous ablations for the key components that make it work well.
Revised: We present a model surgery recipe that is effective in both vision and language, and numerous ablations for the key components that make it work well.

Original: In experiments on Vision Transformers (Dosovitskiy et al., 2021) and T5 language models (Raffel et al., 2020), we show that upcycling is highly effective when one’s computation budget lies between +20% and +60% of the cost to train the original (dense) network.
Revised: In experiments on Vision Transformers (Dosovitskiy et al., 2021) and T5 language models (Raffel et al., 2020), we show that upcycling is highly effective when the computation budget lies between +10% and +60% of the cost to train the original (dense) network.

Original: For example, increasing the performance of ViT-B/16 by at least 1% on ImageNet 10-shot requires an additional 58% extra training time (relative to the original checkpoint) if we continue training the dense model; however, it only takes 13% extra training time with the upcycled version.
Revised: For example, increasing the performance of ViT-B/16 by at least 1% on ImageNet 10-shot requires an additional 58% extra training time (relative to the original checkpoint) if we continue training the dense model; however, it only takes 13% extra training time with the upcycled version.

Original: Similarly, upcycled T5-Large and T5-Base models outperform their dense counterparts by 1.5-2 absolute points on SuperGLUE using 46% and 55% extra training, respectively.
Revised: Similarly, upcycled T5-Large and T5-Base models outperform their dense counterparts by 1.5-2 absolute points on SuperGLUE using 46% and 55% extra training, respectively.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: In this section we recap of the main components used in sparse upcycling: Transformer-based language and vision models, and sparsely activated Mixture-of-Experts (MoEs).
Revised: In this section we recap of the main components used in sparse upcycling: Transformer-based language and vision models, and sparsely activated Mixture-of-Experts (MoEs).

Original: 2. S PARSELY ACTIVATED M IXTURE - OF -E XPERTS (M O E)
Revised: 2. S PARSELY ACTIVATED M IXTURE - OF -E XPERTS (M O E)

Original: Dense models apply all parameters to every input.
Revised: Dense models apply all parameters to every input.

Original: Accordingly, growing the model capacity results in increased computational cost.
Revised: Accordingly, growing the model capacity results in increased computational cost.

Original: Sparse models attempt to alleviate this fundamental issue by only activating a subset of parameters for each input.
Revised: Sparse models attempt to alleviate this fundamental issue by only activating a subset of parameters for each input.

Original: Sparsely activated Mixture-of-Experts (MoE) models are an accelerator friendly family of sparse models that allow training of models with up to trillions of parameters (Shazeer et al., 2017; Fedus et al., 2022).
Revised: Sparsely activated Mixture-of-Experts (MoE) models are an accelerator friendly family of sparse models that allow training of models with up to trillions of parameters (Shazeer et al., 2017; Fedus et al., 2022).

Original: MoE models typically alternate standard dense Transformer blocks with MoE blocks.
Revised: MoE models typically alternate standard dense Transformer blocks with MoE blocks.

Original: In particular, we usually replace the MLPs in a Transformer block with a number of “experts” (typically themselves
Revised: In particular, we usually replace the MLPs in a Transformer block with a number of “experts” (typically themselves

Original: MLPs) with different learnable parameters and a router—a small neural network—that decides which expert is applied to each individual token.
Revised: MLPs) with different learnable parameters and a router—a small neural network—that decides which expert is applied to each individual token.

Original: A number of routing algorithms have been developed, for example Top-K
Revised: A number of routing algorithms have been developed, for example Top-K

Original: (Shazeer et al., 2017), BASE and Sinkhorn-BASE layers (Lewis et al., 2021;
Revised: (Shazeer et al., 2017), BASE and Sinkhorn-BASE layers (Lewis et al., 2021;

Original: Clark et al., 2022), Hash layers (Roller et al., 2021), and Expert Choice routing (Zhou et al., 2022).
Revised: Clark et al., 2022), Hash layers (Roller et al., 2021), and Expert Choice routing (Zhou et al., 2022).

Original: We generally focus on Expert Choice routing, which works as follows.
Revised: We generally focus on Expert Choice routing, which works as follows.

Original: Let E denote the total number of experts in a MoE layer, and n the total number of tokens.
Revised: Let E denote the total number of experts in a MoE layer, and n the total number of tokens.

Original: The router outputs a matrix
Revised: The router outputs a matrix

Original: R ∈ R n × E with the routing probabilities, where row r i ∈ R E corresponds to the i -th token and is a distribution over E experts ( r ij ≥ 0 and (cid:80) j r ij = 1 ).
Revised: R ∈ R n × E with the routing probabilities, where row r i ∈ R E corresponds to the i -th token and is a distribution over E experts ( r ij ≥ 0 and (cid:80) j r ij = 1 ).

Original: Then, every expert e independently chooses the T tokens with highest probabilities for e (i.e., we perform top-T per column) and processes them.
Revised: Then, every expert e independently chooses the T tokens with highest probabilities for e (i.e., we perform top-T per column) and processes them.

Original: We parameterize T as T = C ( n/E ) , where C is a capacity factor that we control to choose more or fewer tokens per expert.
Revised: We parameterize T as T = C ( n/E ) , where C is a capacity factor that we control to choose more or fewer tokens per expert.

Original: When C = 1 , each expert processes exactly n/E tokens; note that some tokens may be processed by several experts, while others by none.
Revised: When C = 1 , each expert processes exactly n/E tokens; note that some tokens may be processed by several experts, while others by none.

Original: This allows for a model parameter count increase with minimal FLOPs overhead.
Revised: This allows for a model parameter count increase with minimal FLOPs overhead.

Original: 2 Letting C > 1 usually leads to higher performance at a higher compute cost.
Revised: 2 Letting C > 1 usually leads to higher performance at a higher compute cost.

Original: 2. A RCHITECTURES
Revised: 2. A RCHITECTURES

Original: We apply the same sparse upcycling recipe to both language and vision tasks, focusing on the
Revised: We apply the same sparse upcycling recipe to both language and vision tasks, focusing on the

Original: T5 (encoder-decoder) (Raffel et al., 2020; Narang et al., 2021) and Vision Transformer (encoder)
Revised: T5 (encoder-decoder) (Raffel et al., 2020; Narang et al., 2021) and Vision Transformer (encoder)

Original: (Dosovitskiy et al., 2021) architectures, respectively.
Revised: (Dosovitskiy et al., 2021) architectures, respectively.

Original: We generally adopt the same gating function and MoE hyperparameters in the encoders of both models.
Revised: We generally adopt the same gating function and MoE hyperparameters in the encoders of both models.

Original: See Section 3. for speciﬁc design choices and  for differences between the vision and language upcycling setups
Revised: See Section 3.1 for speciﬁc design choices and Appendix A for differences between the vision and language upcycling setups.

Original: [Figure]
Revised: [Figure]

Original: Vision .
Revised: Vision .

Original: Vision Transformers (ViT) are encoder-only Transformer architectures (Liu et al., 2021;
Revised: Vision Transformers (ViT) are encoder-only Transformer architectures (Liu et al., 2021;

Original: Radford et al., 2021; Touvron et al., 2021;
Revised: Radford et al., 2021; Touvron et al., 2021;

Original: He et al., 2022) which tokenize and embed images.
Revised: He et al., 2022) which tokenize and embed images.

Original: We upcycle models based on the B/32, B/16, L/32 and L/16 variants.
Revised: We upcycle models based on the B/32, B/16, L/32 and L/16 variants.

Original: The resultant MoEs broadly follow
Revised: The resultant MoEs broadly follow

Original: Vision MoE Transfomers (“V-MoE”)
Revised: Vision MoE Transfomers (“V-MoE”)

Original: (Riquelme et al., 2021), with two differences; we perform global average pooling (Zhai et al., 2022) and use Expert Choice routing.
Revised: (Riquelme et al., 2021), with two differences; we perform global average pooling (Zhai et al., 2022) and use Expert Choice routing.

Original: Language .
Revised: Language .

Original: We experiment with the T5 (Raffel et al., 2020) encoder-decoder as our archetypal language model.
Revised: We experiment with the T5 (Raffel et al., 2020) encoder-decoder as our archetypal language model.

Original: We upcycle the Base, Large, and XL variants of the model.
Revised: We upcycle the Base, Large, and XL variants of the model.

Original: We sparsify both the encoder and decoder.
Revised: We sparsify both the encoder and decoder.

Original: As in our vision setup, the model encoder applies Expert Choice routing.
Revised: As in our vision setup, the model encoder applies Expert Choice routing.

Original: We use Top-K routing in the decoder with K = 2 ; see also Section 3.1.
Revised: We use Top-K routing in the decoder with K = 2 ; see also Section 3.1.

Original:  T HE U PCYCLING A LGORITHM
Revised:  T HE U PCYCLING A LGORITHM

Original: The algorithm is illustrated in Figure 1.
Revised: The algorithm is illustrated in Figure 1.

Original: To upcycle a model, we need a dense model’s parameters (i.e. a checkpoint ).
Revised: To upcycle a model, we need a dense model’s parameters (i.e. a checkpoint ).

Original: The number and shape of Transformer blocks in the new model is identical to that in the original dense model.
Revised: The number and shape of Transformer blocks in the new model is identical to that in the original dense model.

Original: A subset of the of the MLP layers are expanded into MoE layers.
Revised: A subset of the of the MLP layers are expanded into MoE layers.

Original: The remaining MLP layers, along with all of the layer-norm and attention layers, and the embedding and output layers are copied across from the original model to the new model.
Revised: The remaining MLP layers, along with all of the layer-norm and attention layers, and the embedding and output layers are copied across from the original model to the new model.

Original: Each MoE layer contains a ﬁxed number of experts.
Revised: Each MoE layer contains a ﬁxed number of experts.

Original: Each expert is initialized as a copy of the original MLP.
Revised: Each expert is initialized as a copy of the original MLP.

Original: In addition, we add a router whose weights are randomly initialized.
Revised: In addition, we add a router whose weights are randomly initialized.

Original: In Section 4.2.2, we experiment with different variations on this basic recipe.
Revised: In Section 4.2.2, we experiment with different variations on this basic recipe.

Original: After the new model is loaded and initialized, we continue training it for a number of additional steps depending on the available budget and resources.
Revised: After the new model is loaded and initialized, we continue training it for a number of additional steps depending on the available budget and resources.

Original: We use the original hyperparameters: same batch size, learning rate schedule, and weight decay leading to the original checkpoint; see also Appendix A for full training details.
Revised: We use the original hyperparameters: same batch size, learning rate schedule, and weight decay leading to the original checkpoint; see also Appendix A for full training details.

Original: We plan to open-source the upcycling code for both vision and language models in case of acceptance.

Original: 3. D ESIGN D ECISIONS
Revised: 3. D ESIGN D ECISIONS

Original: An upcycled model’s performance is heavily inﬂuenced by the conﬁguration of the MoE layers.
Revised: An upcycled model’s performance is heavily inﬂuenced by the conﬁguration of the MoE layers.

Original: Increasing the model capacity by increasing the number of upcycled layers, number of experts or expert capacity will generally lead to a higher quality model, but will also increase the computational cost and/or result in a greater initial quality drop, due to the more drastic reconﬁguration of the layers.
Revised: Increasing the model capacity by increasing the number of upcycled layers, number of experts or expert capacity will generally lead to a higher quality model, but will also increase the computational cost and/or result in a greater initial quality drop, due to the more drastic reconﬁguration of the layers.

Original: Router type .
Revised: Router type .

Original: For upcycled vision models and for the encoder of upcycled language models, we use
Revised: For upcycled vision models and for the encoder of upcycled language models, we use

Original: Expert Choice routing with capacity factor C = 2 .
Revised: Expert Choice routing with capacity factor C = 2 .

Original: To avoid train time (full batch teacher forcing) versus inference time (single token auto-regressive decoding) discrepancies, we use Top-K ( K = 2 ) routing in the language decoder.
Revised: To avoid train time (full batch teacher forcing) versus inference time (single token auto-regressive decoding) discrepancies, we use Top-K ( K = 2 ) routing in the language decoder.

Original: In Section 4.2.2, we show that Expert Choice routing outperforms standard Top-K routing for upcycling, while both beat dense continuations.
Revised: In Section 4.2.2, we show that Expert Choice routing outperforms standard Top-K routing for upcycling, while both beat dense continuations.

Original: Number layers to upcycle.
Revised: Number layers to upcycle.

Original: Adding more MoE layers increases the model capacity dramatically, at the expense of increasing the model’s cost, and also causing the quality of the upcycled model to initially drop further relative to the original dense model.
Revised: Adding more MoE layers increases the model capacity dramatically, at the expense of increasing the model’s cost, and also causing the quality of the upcycled model to initially drop further relative to the original dense model.

Original: Based on our ablation in Section 4.2.2 and prevailing conventions in the MoE literature (Lepikhin et al., 2021), unless otherwise speciﬁed, we replace half of the MLP layers in our upcycled models with MoE layers.
Revised: Based on our ablation in Section 4.2.2 and prevailing conventions in the MoE literature (Lepikhin et al., 2021), unless otherwise speciﬁed, we replace half of the MLP layers in our upcycled models with MoE layers.

Original: Number of experts to add in upcycled layers .
Revised: Number of experts to add in upcycled layers .

Original: Each new expert provides new learnable parameters that extend the model capacity.
Revised: Each new expert provides new learnable parameters that extend the model capacity.

Original: The expert capacity—the number of tokens expert processes–is inversely proportional to the number of experts, thus adding more experts does not signiﬁcantly affect the FLOPS or the run time of the model.
Revised: The expert capacity—the number of tokens expert processes–is inversely proportional to the number of experts, thus adding more experts does not signiﬁcantly affect the FLOPS or the run time of the model.

Original: However, with a very large number of experts, the upcycled model experiences a larger initial quality drop relative to the baseline dense model.
Revised: However, with a very large number of experts, the upcycled model experiences a larger initial quality drop relative to the baseline dense model.

Original: Given sufﬁcient upcycling compute, this initial drop can be overcome.
Revised: Given sufﬁcient upcycling compute, this initial drop can be overcome.

Original: In our studies, we upcycle with +20% to +100% of the initial dense baseline model’s computational cost, and in this regime we ﬁnd thatexperts provides a good compromise.
Revised: In our studies, we upcycle with +20% to +100% of the initial dense baseline model’s computational cost, and in this regime we ﬁnd thatexperts provides a good compromise.

Original: We explore varying the number of experts in Section 4.2.2.
Revised: We explore varying the number of experts in Section 4.2.2.

Original: Expert capacity .
Revised: Expert capacity .

Original: By tuning the expert capacity, C , we control the number of experts that process each token on average.
Revised: By tuning the expert capacity, C , we control the number of experts that process each token on average.

Original: 3 Larger expert capacity generally yields larger quality but also increases the FLOPS and run time.
Revised: 3 Larger expert capacity generally yields larger quality but also increases the FLOPS and run time.

Original: Although increasing the expert capacity yields quality gains on a per step basis, we ﬁnd that C = 2 generally offers good quality on a compute time basis.
Revised: Although increasing the expert capacity yields quality gains on a per step basis, we ﬁnd that C = 2 generally offers good quality on a compute time basis.

Original: We ablate through different capacity factors in Section 4.2.2.
Revised: We ablate through different capacity factors in Section 4.2.2.

Original: Resuming optimizer state (vision only) .
Revised: Resuming optimizer state (vision only) .

Original: When upcycling a model, we can resume the optimizer state from the original dense checkpoint together with the model parameters.
Revised: When upcycling a model, we can resume the optimizer state from the original dense checkpoint together with the model parameters.

Original: In Appendix B.6, we ﬁnd that reusing the optimizer state gives a performance boost for vision models.
Revised: In Appendix B.6, we ﬁnd that reusing the optimizer state gives a performance boost for vision models.

Original: We did not, however, see any improvement from reusing the dense model optimizer state in our language experiments, so we only reuse the optimizer state for vision models.
Revised: We did not, however, see any improvement from reusing the dense model optimizer state in our language experiments, so we only reuse the optimizer state for vision models.

Original: Normalize weights after routing (vision only) .
Revised: Normalize weights after routing (vision only) .

Original: In an effort to reduce the performance drop when applying the upcycling model surgery, we attempted to normalize the router combine weights of each token to 1 .
Revised: In an effort to reduce the performance drop when applying the upcycling model surgery, we attempted to normalize the router combine weights of each token to 1 .

Original: This follows that intuition that each token was previously only processed by a single “expert” MLP in the dense model.
Revised: This follows that intuition that each token was previously only processed by a single “expert” MLP in the dense model.

Original: Appendix B.7 shows that router weight normalization helps upcycled vision models, but hurts the performance of upcycled language models.
Revised: Appendix B.7 shows that router weight normalization helps upcycled vision models, but hurts the performance of upcycled language models.

Original: One hypothesis for this different behavior is that the vision models use Expert Choice routing everywhere, but the language models use Expert Choice in the encoder and Top-K routing in the decoder.
Revised: One hypothesis for this different behavior is that the vision models use Expert Choice routing everywhere, but the language models use Expert Choice in the encoder and Top-K routing in the decoder.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section, we present the main experimental results of the paper.
Revised: In this section, we present the main experimental results of the paper.

Original: We also share the takeaways of a number of ablations aimed at identifying the key aspects of our algorithm; full results are included in Appendix B. Most of the results are presented as quality vs. cost plots, where we use the upstream or downstream performance to measure quality, and training time in terms of TPU-core-days (as prominent cost metrics (Dehghani et al., 2021)) or training steps (when the cost per step is the same for all the compared models) to measure computation cost.
Revised: We also share the takeaways of a number of ablations aimed at identifying the key aspects of our algorithm; full results are included in Appendix B. Most of the results are presented as quality vs. cost plots, where we use the upstream or downstream performance to measure quality, and training time in terms of TPU-core-days (as prominent cost metrics (Dehghani et al., 2021)) or training steps (when the cost per step is the same for all the compared models) to measure computation cost.

Original: 4. E XPERIMENTAL S ETUP
Revised: 4. E XPERIMENTAL S ETUP

Original: All upcycled experiments begin from a pretrained dense model checkpoint.
Revised: All upcycled experiments begin from a pretrained dense model checkpoint.

Original: Because all of our starting dense checkpoints are trained with an inverse square root learning rate schedule, training can be continued without discontinuities in the learning rate schedule.
Revised: Because all of our starting dense checkpoints are trained with an inverse square root learning rate schedule, training can be continued without discontinuities in the learning rate schedule.

Original: We upcycle models and continue training, showing performance for varying amounts of continued training.
Revised: We upcycle models and continue training, showing performance for varying amounts of continued training.

Original: As a baseline, we also continue the training of the original dense model (“dense continuation”).
Revised: As a baseline, we also continue the training of the original dense model (“dense continuation”).

Original: Vision experiments .
Revised: Vision experiments .

Original: MoE Vision Transfomers (“V-MoE”) models are trained broadly following the protocol of Riquelme et al.
Revised: MoE Vision Transfomers (“V-MoE”) models are trained broadly following the protocol of Riquelme et al.

Original: Upstream pretraining is done on JFT300M (Sun et al., 
Revised: Upstream pretraining is done on JFT300M (Sun et al., 2017), with validation metrics computed on a held-out set of 894,574 examples.

Original:  3 For Expert Choice routing, more capacity 
Revised: Few-shot transfer 3 For Expert Choice routing, more capacity means that each

Original: Extra Pretraining Time (TPU-core-days) 38% 43% 49% 53% 58%
Revised: Extra Pretraining Time (TPU-core-days) 38% 43% 49% 53% 58%

Original: J F T
Revised: J F T

Original: V a li d a t i o n
Revised: V a li d a t i o n

Original: P r e c
Revised: P r e c

Original: i s i o n a t
Revised: i s i o n a t

Original: B/
Revised: B/

Original: B/
Revised: B/

Original: L/ method Dense
Revised: L/ method Dense

Original: Upcycling variant
Revised: Upcycling variant

Original: B/32 B/
Revised: B/32 B/

Original: L/ 

Revised: L/

Original: Extra Pretraining Time (TPU-core-days) 68% 70% 72% 73% 75%C
Revised: Extra Pretraining Time (TPU-core-days) 68% 70% 72% 73% 75%

Original: V a li d a t i o n

Original: T o

Original: [Figure]

Original: e n

Original: A cc u r a

Original: [Figure]
Revised: [Figure]

Original: y

Original: Base
Revised: Base

Original: Large
Revised: Large

Original: XL method Dense
Revised: XL [Figure]

Original: Upcycling variant
Revised: variant

Original: Base Large XL 2017), with validation metrics computed on a held-out set of 894,574 examples.
Revised: [Figure]

Original: Few-shot transfer follows Dosovitskiy et al.
Revised: XL follows Dosovitskiy et al.

Original: (2021), whereby a least-squares regressor predicts one-hot classes given frozen image representations.
Revised: (2021), whereby a least-squares regressor predicts one-hot classes given frozen image representations.

Original: We further validate our results on ImageNet using 10-shot – i.e.training examples per class.
Revised: We further validate our results on ImageNet using 10-shot – i.e.training examples per class.

Original: We do this for 5 different training sets, and report average accuracy across them.
Revised: We do this for 5 different training sets, and report average accuracy across them.

Original: For full ﬁnetuning, we replace the pretraining head with a randomly initialized head, and ﬁnetune the entire network.
Revised: For full ﬁnetuning, we replace the pretraining head with a randomly initialized head, and ﬁnetune the entire network.

Original: See Appendix A.2.2 for further details.
Revised: See Appendix A.2.2 for further details.

Original: Language experiments .
Revised: Language experiments .

Original: Our language experiments follow the setup of Raffel et al.
Revised: Our language experiments follow the setup of Raffel et al.

Original: (2020): we pretrain using the span corruption task on the English C4 dataset (Raffel et al., 2020) and ﬁnetune on a proportional mix of all SuperGLUE (Wang et al., 2019) tasks simultaneously.
Revised: (2020): we pretrain using the span corruption task on the English C4 dataset (Raffel et al., 2020) and ﬁnetune on a proportional mix of all SuperGLUE (Wang et al., 2019) tasks simultaneously.

Original: We include speciﬁc training details in Appendix A.2, but highlight one important aspect here: For Base model sizes, for which we perform the majority of our ablations, we pretrain the dense baseline starting checkpoint ourselves.
Revised: We include speciﬁc training details in Appendix A.2, but highlight one important aspect here: For Base model sizes, for which we perform the majority of our ablations, we pretrain the dense baseline starting checkpoint ourselves.

Original: To highlight the versatility of our upcycling algorithm, for Large and XL models, we instead begin all experiments from ofﬁcial T5 1.1 checkpoints (Narang et al., 2021; Roberts et al., 2022).
Revised: To highlight the versatility of our upcycling algorithm, for Large and XL models, we instead begin all experiments from ofﬁcial T5 1.1 checkpoints (Narang et al., 2021; Roberts et al., 2022).

Original: 4. E XPERIMENTAL R ESULTS 4.2. C ORE RESULTS
Revised: 4. E XPERIMENTAL R ESULTS 4.2. C ORE RESULTS

Original: Figure 2 shows a detailed comparison of upstream metrics of upcycled models and dense continuation models at various model sizes both for vision (left panel) and language (right panel).
Revised: Figure 2 shows a detailed comparison of upstream metrics of upcycled models and dense continuation models at various model sizes both for vision (left panel) and language (right panel).

Original: For any given model size and task, we observe that the dense and upcycled models perform close to each other when we apply a very limited extra training budget – indeed, close to their discontinuous horizontal line representing the original checkpoint’s performance.
Revised: For any given model size and task, we observe that the dense and upcycled models perform close to each other when we apply a very limited extra training budget – indeed, close to their discontinuous horizontal line representing the original checkpoint’s performance.

Original: Once we apply a non-trivial amount of extra compute, a clear pattern emerges showing the strong gains delivered by the upcycled architecture.
Revised: Once we apply a non-trivial amount of extra compute, a clear pattern emerges showing the strong gains delivered by the upcycled architecture.

Original: Figure 3 shows the performance after ﬁnetuning the models trained in Figure 2.
Revised: Figure 3 shows the performance after ﬁnetuning the models trained in Figure 2.

Original: For vision (left panel), the upstream performance gains generally transfer fairly cleanly downstream.
Revised: For vision (left panel), the upstream performance gains generally transfer fairly cleanly downstream.

Original: For language (right panel), there is substantial variance in the performance after ﬁnetuning.
Revised: For language (right panel), there is substantial variance in the performance after ﬁnetuning.

Original: There is a lack of performance monotonicity with respect to additional (upstream) pretraining even for the dense Large models.
Revised: There is a lack of performance monotonicity with respect to additional (upstream) pretraining even for the dense Large models.

Original: Nevertheless, the trend favors the upcycled language models.
Revised: Nevertheless, the trend favors the upcycled language models.

Original: Figure 4 compares sparse upcycling with sparse models trained from scratch.
Revised: Figure 4 compares sparse upcycling with sparse models trained from scratch.

Original: As training from scratch does not reuse the computation cost already sunk into the dense checkpoint, it takes longer, on a extra train time basis, to catch up with the upcycled models.
Revised: As training from scratch does not reuse the computation cost already sunk into the dense checkpoint, it takes longer, on a extra train time basis, to catch up with the upcycled models.

Original: The language MoE model trained from scratch requires about 120% of the original dense checkpoint’s computation budget to catch up to the 
Revised: The language MoE model trained from scratch requires about 120% of the original dense checkpoint’s computation budget to catch up to the upcycled model.

Revised: The relatively faster quality gains, on a per step basis, of the MoE models

Original: [Figure]
Revised: [Figure]

Original: Pretraining Time (TPU
Revised: Pretraining Time (TPU

Original: [Figure]
Revised: [Figure]

Original: core
Revised: core

Original: [Figure]
Revised: [Figure]

Original: V R
Revised: V R

Original: C
Revised: C

Original: [Figure]
Revised: [Figure]

Original: i n e t un e
Revised: i n e t un e

Original: A cc u r a c y
Revised: A cc u r a c y

Original: [Figure]
Revised: [Figure]

Original: Pretraining 
Revised: Pretraining Time

Original: [Figure]
Revised: [Figure]

Original: TPU
Revised: TPU

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: checkpoints trained 
Revised: trained from scratch can be attributed to the relatively larger learning rate and that the experts are able to independently develop and diversify from the beginning.

Original: up
Revised: Figure 4 suggests that, given a very large computation budget ( > 100% of the initial dense model’s computation budget), the MoE-fromscratch model will eventually catch the upcycled model. For such large computation regimes, it may be preferable to train MoE models from scratch. For constrained or limited compute budgets ( < 100% of the initial computation budget), sparse upcycling is a more efﬁcient use of resources.

Original: upcycled model. The relatively faster quality gains, on a per step basis, of the MoE modelstrained fromscratch can be attributed to the relatively larger learning rate and the fact that the experts are able to independently develop and diversify from the beginning.
Revised: Finally, Figure 5 compares sparse upcycling with warm starting (“dense upcycling”). We warm start larger models from the dense Base checkpoint by replicating new layers (”depth tiling”) in the same tiling patterns as in (Rae et al., 2021).

Original: 4.2. A BLATIONS
Revised: The densely upcycled models quickly see gains over the original dense checkpoint, but underperform the sparse model. We did not attempt to increase the model hidden dimensions (“width tiling”), which (Rae et al., 2021) found to be less effective. [Figure] [Figure]

Revised: are fairly consistent independent of the amount of initial pretraining.

Revised: Note: for this particular ablation, we use a capacity factor of C = 1 , to ensure that the FLOPS and run times of the dense model and sparsely upcycled model are roughly comparable on a per step basis.

Revised: 4.2. A BLATIONS

Original: In this section we summarize important architecture and training ablations relative to the baseline model described in Section 3.
Revised: In this section we summarize important architecture and training ablations relative to the baseline model described in Section 3.

Original: Full results are provided in Appendix B.
Revised: Full results are provided in Appendix B.

Original: Unless stated otherwise, vision ablations use a B/16 sparse model with 32 experts, C = 1 and 6 MoE layers placed in the last few blocks of the model.
Revised: Unless stated otherwise, vision ablations use a B/16 sparse model with 32 experts, C = 1 and 6 MoE layers placed in the last few blocks of the model.

Original: The dense checkpoint was trained for 14 epochs, and we train for an additional 7 epochs (up to a total of 21 epochs).For our language ablations, our default conﬁguration [Figure] a
Revised: The dense checkpoint was trained for 14 epochs, and we train for an additional 7 epochs (up to a total of 21 epochs).

Original: [Figure]

Original: i o n

Original: [Figure]

Original: r

Original: [Figure]

Original: c i s i o n [Figure] [Figure] is unchanged: we use a Base model with 32 experts, C = 2 and 6 MoE layers interspersed throughout the model.
Revised: For our language ablations, our default conﬁguration is unchanged: we use a Base model with 32 experts, C = 2 and 6 MoE layers interspersed throughout the model.

Original: We train for between 0.5 million and 1 million extra steps.
Revised: We train for between 0.5 million and 1 million extra steps.

Original: Amount of dense pretraining .
Revised: Amount of dense pretraining .

Original: The upcycling efﬁciency may, in principle, depend on how converged the initial dense model is.
Revised: The upcycling efﬁciency may, in principle, depend on how converged the initial dense model is.

Original: To explore this, in Figure 5, we upcycle a B/16 vision model starting from different dense checkpoints with varying amounts of pretraining.
Revised: To explore this, in Figure 6, we upcycle a B/16 vision model starting from different dense checkpoints with varying amounts of pretraining.

Original: From a given dense checkpoint, we compare upcycling and dense continuation for 200k steps.
Revised: From a given dense checkpoint, we compare upcycling and dense continuation for 200k steps.

Original: Independent of when we start upcycling, the performance improvement from doing so is fairly consistent.
Revised: Independent of when we start upcycling, the performance improvement from doing so is fairly consistent.

Original: Router type .
Revised: Router type .

Original: While our default upcycling recipe uses Expert Choice routing (in the encoder), the same recipe can be applied to other routing mechanisms.
Revised: While our default upcycling recipe uses Expert Choice routing (in the encoder), the same recipe can be applied to other routing mechanisms.

Original: For vision, Appendix B.1 shows that although Top-K routing, with Batch Prioritized Routing (BPR) (Riquelme et al., 2021), matches
Revised: For vision, Appendix B.1 shows that although Top-K routing, with Batch Prioritized Routing (BPR) (Riquelme et al., 2021), matches

Original: Expert Choice routing performance on a per step basis, as it is slightly slower, Top-K routing underperforms Expert Choice routing on a per train time basis.
Revised: Expert Choice routing performance on a per step basis, as it is slightly slower, Top-K routing underperforms Expert Choice routing on a per train time basis.

Original: Note both approaches beat dense.
Revised: Note both approaches beat dense.

Original: Expert capacity factor .
Revised: Expert capacity factor .

Original: The more tokens processed per expert, the greater the amount of compute per input example and (generally) the higher the model quality.
Revised: The more tokens processed per expert, the greater the amount of compute per input example and (generally) the higher the model quality.

Original: Our sparse models smoothly control this via the expert capacity factor C .
Revised: Our sparse models smoothly control this via the expert capacity factor C .

Original: Appendix B.2 explores how the performance-speed trade-off varies as a function of C .
Revised: Appendix B.2 explores how the performance-speed trade-off varies as a function of C .

Original: Although increasing the expert capacity yields quality gains on a per step basis, we ﬁnd that C = 2 generally offers the best quality on a compute time basis, for both language and vision models.
Revised: Although increasing the expert capacity yields quality gains on a per step basis, we ﬁnd that C = 2 generally offers the best quality on a compute time basis, for both language and vision models.

Original: Number of MoE layers .
Revised: Number of MoE layers .

Original: A key decision when upcycling a model is how many sparse layers to add.
Revised: A key decision when upcycling a model is how many sparse layers to add.

Original: Model capacity and parameter-count increases with the number of MoE layers, but at the expense of slowing down model run time.
Revised: Model capacity and parameter-count increases with the number of MoE layers, but at the expense of slowing down model run time.

Original: Appendix B.4 provides an ablation over the number of MoE layers in an upcycled B/16 vision model.
Revised: Appendix B.4 provides an ablation over the number of MoE layers in an upcycled B/16 vision model.

Original: In this case, around 6 MoE layers (out of a total of 12 layers) offers the best precision quality, although fewer sparse layers lower the compute costs.
Revised: In this case, around 6 MoE layers (out of a total of 12 layers) offers the best precision quality, although fewer sparse layers lower the compute costs.

Original: Initialization of experts .
Revised: Initialization of experts .

Original: The standard upcycling recipe copies and replicates the dense MLP to each expert.
Revised: The standard upcycling recipe copies and replicates the dense MLP to each expert.

Original: As the router directs different tokens to each expert, the experts will start to diverge from one another, and from their initial MLP weights.
Revised: As the router directs different tokens to each expert, the experts will start to diverge from one another, and from their initial MLP weights.

Original: Appendix B.5 compares the standard recipe with randomly initializing the experts (i.e. train them from scratch).
Revised: Appendix B.5 compares the standard recipe with randomly initializing the experts (i.e. train them from scratch).

Original: For limited computation budgets, randomly initializing experts underperforms the standard recipe; for larger compute budgets it eventually matches the standard upcycling recipe performance.
Revised: For limited computation budgets, randomly initializing experts underperforms the standard recipe; for larger compute budgets it eventually matches the standard upcycling recipe performance.

Original: We also tried copying the original MLP weights and adding independent (Gaussian) noise to each expert, in an effort to promote more expert diversity.
Revised: We also tried copying the original MLP weights and adding independent (Gaussian) noise to each expert, in an effort to promote more expert diversity.

Original: Adding too much noise when copying MLP weights into experts hurts performance, while adding small amounts of noise has little to no effect; see also Appendix B.9.
Revised: Adding too much noise when copying MLP weights into experts hurts performance, while adding small amounts of noise has little to no effect; see also Appendix B.9.

Original: Number of experts .
Revised: Number of experts .

Original: Adding more experts increases the number of model parameters and, up to a point, the quality of the model.
Revised: Adding more experts increases the number of model parameters and, up to a point, the quality of the model.

Original: Given that the number of tokens each expert processes is inversely proportional to the number of experts (see Section 2.1), adding more experts does not signiﬁcantly affect the model FLOPS nor its running time.
Revised: Given that the number of tokens each expert processes is inversely proportional to the number of experts (see Section 2.1), adding more experts does not signiﬁcantly affect the model FLOPS nor its running time.

Original: However, for a very large number of experts, the upcycled model may experience a larger initial quality drop relative to the baseline dense model.
Revised: However, for a very large number of experts, the upcycled model may experience a larger initial quality drop relative to the baseline dense model.

Original: Appendix B.3 explores this trade-off and shows that, at least for Base sized models, more experts usually yield better performance.
Revised: Appendix B.3 explores this trade-off and shows that, at least for Base sized models, more experts usually yield better performance.

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Reuse of trained parameters .
Revised: Reuse of trained parameters .

Original: There is a body of work with a focus on speeding upthe trainingprocess through a warm start by reusing parameters of an existing model or gradually growing parameters in multiple stages.
Revised: Prior work has focused on speeding up training through a warm start by reusing parameters of an existing model.

Original: Berner et al.
Revised: Berner et al.

Original: (2019) explore ideas of reusing the previous edition of a trained model during a long training process using an under-development environment.
Revised: (2019) explore ideas of reusing the previous edition of a trained model during a long training process using an under-development environment.

Original: They propose a “surgery” that may involve adding new parameters at locations where they initialize the new parameters in a way that has no impact on the activation at the beginning of training compared to the original model.

Original:  Net2Net (Chen et al., 2015), given a trained model, propose a function-preserving initialization to warm start training a deeper or wider model.
Revised: Given a trained model, Net2Net (Chen et al., 2015) propose a function-preserving initialization to warm start training a deeper or wider model.

Original: Recently, Gopher (Rae et al., 2021) also explored warm starting larger models from smaller models in a large compute training regime and show that the larger, warm started model eventually reaches to a quality close to that of the equivalent model trained from scratch. 
Revised: Recently, Gopher (Rae et al., 2021) also explored warm starting larger models from smaller models in a large compute training regime and show that the larger, warm started model can converge to a quality comparable to that of the equivalent model trained from scratch. In Section 4.2.1, we show that warm starting signiﬁcantly underperforms sparse upcycling.

Original: Yang et al.
Revised: Yang et al.

Original: (2021); Lin et al.
Revised: (2021); Lin et al.

Original: (2021) offer training recipes wherein they begin with models that  share parameters across layers Lan et al.
Revised: (2021) show that they can reduce the number of training iterations with models that inititally share parameters across layers (Lan et al., 2019;

Original: Dehghani et al. (2018) and then unshare (or “delink”)
Revised: Dehghani et al., 2018) but gradually unshare (or “delink”)

Original: the parameters before continuing training.
Revised: the parameters while training.

Original: They show that this technique can reduce the number of training iterations. Progressively growing a model during training (Gong et al., 2019; Dong et al., 2020; Li et al., 2020) is also explored as a way to reduce the total training cost.
Revised: In an effort to reduce total training cost, several works explore progressively growing models during training (Gong et al., 2019; Dong et al., 2020; Li et al., 2020; Shen et al., 2022).

Original: The main idea is to decompose the training process into stages where we start with training a shallow model and then grow it by copyingthe weights and stacking new layers on top in the following stages.
Revised: The core idea is to decompose the training process into stages, each of which apply growth operators to increase the model size from the previous stage by copying weights or stacking new layers on top.

Original: In some cases, each training stage will only update parameters of the layers which are newly, which saves the cost of  full backward computation (Yang et al., 2020).
Revised: In some cases, each training stage will only update parameters of the new layers, which saves the cost of a full backward computation (Yang et al., 2020).

Original: Gu et al.
Revised: Gu et al.

Original: (2020) show thatin multi-stage training of transformers, compound scaling (scaling depth, width, input length, etc. together) is favorable and propose a strategy with various growing operators on each dimension, leading to better performance while improving the training time.
Revised: (2020) show that compound scaling (scaling depth, width and input length together) is favorable and propose a strategy with various growing operators on each dimension.

Original: Sparse upcycling, which we introduce in this paper, follows a similar motivation.
Revised: Sparse upcycling, which we introduce in this paper, follows a similar motivation.

Original: However, unlike the above works, we focus on compute regimes that are a fraction of the original model’s training.
Revised: However, unlike the above works, we focus on compute regimes that are a fraction of the original model’s training.

Original: We also present a recipe for growing a trained dense model to a sparse model, instead of a larger dense model.
Revised: We also present a recipe for growing a trained dense model to a sparse model, instead of a larger dense model.

Original: This enables us to enjoy the extra capacity due to increased parameters, while maintaining the inference cost due to the sparsity of computation.
Revised: This enables us to enjoy the extra capacity due to increased parameters, while maintaining the inference cost due to the sparsity of computation.

Original: Pruning .
Revised: Pruning .

Original: Pruning is typically employed as a post-training architecture search to construct smaller and faster models from larger models (LeCun et al., 1989; Gale et al., 2019; Blalock et al., 2020).
Revised: Pruning is typically employed as a post-training architecture search to construct smaller and faster models from larger models (LeCun et al., 1989; Gale et al., 2019; Blalock et al., 2020).

Original: However, “dynamic pruning” (Evci et al., 2020) has also been used during training to ﬁnd sparser architectures from dense models.
Revised: However, “dynamic pruning” (Evci et al., 2020) has also been used during training to ﬁnd sparser architectures from dense models.

Original: Similar to pruning, sparse upcycling also introduces sparsity to a dense model, however, unlike pruning, we grow the existing dense models into a larger sparse model.
Revised: Similar to pruning, sparse upcycling also introduces sparsity to a dense model, however, unlike pruning, we grow the existing dense models into a larger sparse model.

Original: Sparsely-activated Mixture-of-Experts (MoE) .
Revised: Sparsely-activated Mixture-of-Experts (MoE) .

Original: In this work, we sparsify existing dense models into MoE models.
Revised: In this work, we sparsify existing dense models into MoE models.

Original: MoE models (Shazeer et al., 2017) offer the promise of increasing model scale (parameter count) with sublinear increases in computation cost (FLOPS).
Revised: MoE models (Shazeer et al., 2017) offer the promise of increasing model scale (parameter count) with sublinear increases in computation cost (FLOPS).

Original: Recently, there has been a growing number of MoE works achieving state-of-the-art quality and remarkable efﬁciency gains  (Lepikhin et al., 2021; Fedus et al., 2022;
Revised: Recently, there has been a growing number of MoE works achieving state-of-the-art quality and remarkable efﬁciency gains on both language and vision tasks (Lepikhin et al., 2021; Fedus et al., 2022;

Original: Riquelme et al., 2021; Artetxe et al., 2021; Du et al., 2021; Zoph et al., 2022; Mustafa et al., 2022) on both language and vision tasks.
Revised: Riquelme et al., 2021; Artetxe et al., 2021; Du et al., 2021; Zoph et al., 2022; Mustafa et al., 2022).

Original: All of these models are large and trained from scratch with randomly initialized weights and ﬁxed architectures.
Revised: All of these models are large and trained from scratch with randomly initialized weights and ﬁxed architectures.

Original: Several MoE works have also attempted to improve upon typical training algorithms by adapting or “evolving” the model architecture during training.
Revised: Several MoE works have also attempted to improve upon typical training algorithms by adapting or “evolving” the model architecture during training.

Original: Nie et al.
Revised: Nie et al.

Original: (2021) progressively sparsify MoE layers during training by slowly adjusting the gating function from a “dense setup”, where all tokens are routed to all experts, to a fully “sparse setup”, where tokens are only routed to a subset of experts.
Revised: (2021) progressively sparsify MoE layers during training by slowly adjusting the gating function from a “dense setup”, where all tokens are routed to all experts, to a fully “sparse setup”, where tokens are only routed to a subset of experts.

Original: C ONCLUSIONS
Revised: Zhang et al. (2022) observed that, for most inputs, only a small fraction of Transformer MLP activations are nonzero. Based on this, they propose sparsiﬁcation procedure that splits the parameters of MLP blocks into multiple experts and add a routing mechanism.

Revised: Similarly, Zuo et al.

Revised: (2022) split up the MLPs in a pre-trained dense model into multiple experts to form a sparse model for ﬁne-tuning.

Revised: Closely related to our work, Wu et al.

Revised: (2022) present a novel algorithm to sparsify dense models in the context of ﬁnetuning on detection and segmentation tasks.

Revised: Similar to Nie et al.

Revised: (2021), the initial performance drop when training on the original dataset is avoided by applying a dense mixture of experts in the forward pass.

Revised: However, at our target large scales, simultaneously activating all experts for each token is not feasible.

Revised: Finally, Gururangan et al.

Revised: (2022) adapt sparse, domain expert language models to new domains by initializing a new domain expert from the most probable existing expert under the domain posterior distribution.

Revised:  C ONCLUSIONS

Original: Training large neural networks on huge datasets has proven to be a remarkably successful trend in deep learning research, especially in recent years.
Revised: Training large neural networks on huge datasets has proven to be a remarkably successful trend in deep learning research, especially in recent years.

Original: It has also proven to be very computationally expensive.
Revised: It has also proven to be very computationally expensive.

Original: Pretrained models are now widely available, thus making it possible for many practitioners to further ﬁnetune and adapt ﬁxed model architectures on their data of interest.
Revised: Pretrained models are now widely available, thus making it possible for many practitioners to further ﬁnetune and adapt ﬁxed model architectures on their data of interest.

Original: However, signiﬁcant progress requires providing more ﬂexibility in adapting and improving the model architecture itself.
Revised: However, signiﬁcant progress requires providing more ﬂexibility in adapting and improving the model architecture itself.

Original: We proposed a simple recipe to reuse pretrained dense checkpoints to initialize more powerful sparse models.
Revised: We proposed a simple recipe to reuse pretrained dense checkpoints to initialize more powerful sparse models.

Original: Our algorithm leverages the pretrained model compute and weights, and provides a smooth transition to sparsely activated Mixture-of-Experts models that offer more capacity and ﬂexibility at inference.
Revised: Our algorithm leverages the pretrained model compute and weights, and provides a smooth transition to sparsely activated Mixture-of-Experts models that offer more capacity and ﬂexibility at inference.

Original: We presented experimental results both for vision and language models at various scales; these evidence large performance gains relative to continuing to the dense model.
Revised: We presented experimental results both for vision and language models at various scales; these evidence large performance gains relative to continuing to the dense model.

Original: Our ablations highlight the importance of careful algorithmic choices, and suggest key aspects to consider when trying to ﬁnd good performance-cost trade-offs for speciﬁc compute budgets.
Revised: Our ablations highlight the importance of careful algorithmic choices, and suggest key aspects to consider when trying to ﬁnd good performance-cost trade-offs for speciﬁc compute budgets.

Original: Transfer learning and prompt tuning is becoming increasingly popular, and for good reason.
Revised: Transfer learning and prompt tuning is becoming increasingly popular, and for good reason.

Original: It allows the reuse and tuning of models by a larger body of researchers and practitioners that may only have access to limited computational and data resources.
Revised: It allows the reuse and tuning of models by a larger body of researchers and practitioners that may only have access to limited computational and data resources.

Original: Accordingly, we believe that techniques aimed at growing existing models, compartmentalizing or freezing submodules, replicating and then decoupling model components, and ﬁnally smoothly resuming training after model surgery, will prove essential for a dynamic ecosystem of models.
Revised: Accordingly, we believe that techniques aimed at growing existing models, compartmentalizing or freezing submodules, replicating and then decoupling model components, and ﬁnally smoothly resuming training after model surgery, will prove essential for a dynamic ecosystem of models.

Original: We summarize such process as upcycling , and offer a ﬁrst instance in the context of sparse models.
Revised: We summarize such process as upcycling , and offer a ﬁrst instance in the context of sparse models.

Original: We look forward to new extensions and improvements on this simple idea.
Revised: We look forward to new extensions and improvements on this simple idea.

Original: A T RAINING AND EVALUATION DETAILS
Revised: A T RAINING AND EVALUATION DETAILS

Original: In this section, we describe the precise setup for our language and vision experiments.
Revised: In this section, we describe the precise setup for our language and vision experiments.

Original: Figureillustrates the type of combined training curves we obtained before and after upcycling.
Revised: Figureillustrates the type of combined training curves we obtained before and after upcycling.

Original: A. U PSTREAM TRAINING A.1. L ANGUAGE
Revised: A. U PSTREAM TRAINING A.1. L ANGUAGE

Original: We ﬁrst pretrain a dense Base model from scratch for 1 million steps with a batch size ofsequences.
Revised: We ﬁrst pretrain a dense Base model from scratch for 1 million steps with a batch size ofsequences.

Original: We use the Adafactor optimizer with an inverse square root decay and a peak learning rate of 0 .
Revised: We use the Adafactor optimizer with an inverse square root decay and a peak learning rate of 0 .

Original: This results in plateauing performance for the dense Base model.
Revised: This results in plateauing performance for the dense Base model.

Original: Upcycled models are then initialized from the 1 million step dense checkpoint, and compared, on a compute time basis, with further training of the dense model (“dense continuation”); see Figure 2 in the main text.
Revised: Upcycled models are then initialized from the 1 million step dense checkpoint, and compared, on a compute time basis, with further training of the dense model (“dense continuation”); see Figure 2 in the main text.

Original: To highlight the versatility of our upcycling algorithm, for Large and XL models, we instead begin all experiments from ofﬁcial T5 1.1 checkpoints (Narang et al., 2021).
Revised: To highlight the versatility of our upcycling algorithm, for Large and XL models, we instead begin all experiments from ofﬁcial T5 1.1 checkpoints (Narang et al., 2021).

Original: We use the same hyperparameters for the upcycled model as for the corresponding dense model that we initialized from, continuing the inverse square root learning rate schedule where the dense checkpoint left off.
Revised: We use the same hyperparameters for the upcycled model as for the corresponding dense model that we initialized from, continuing the inverse square root learning rate schedule where the dense checkpoint left off.

Original: For all sizes, every other layer was upcycled, using 32 experts, starting with the second layer.
Revised: For all sizes, every other layer was upcycled, using 32 experts, starting with the second layer.

Original: Similar to dense pretraining, we do not include any dropout (or expert dropout; see Section A.2.1 below).
Revised: Similar to dense pretraining, we do not include any dropout (or expert dropout; see Section A.2.1 below).

Original: Router parameters are initialized randomly with a zero-mean normal distribution with standard deviation 0.02.
Revised: Router parameters are initialized randomly with a zero-mean normal distribution with standard deviation 0.02.

Original: We use a maximum routing group size of 4096 tokens.
Revised: We use a maximum routing group size of 4096 tokens.

Original: For
Revised: For

Original: Top-2 routing (in the decoder) we include an auxiliary MoE loss, with scaling factor 0.01, to ensure tokens are distributed more evenly across all experts in the decoder (Shazeer et al., 2017; Fedus et al., 2022).
Revised: Top-2 routing (in the decoder) we include an auxiliary MoE loss, with scaling factor 0.01, to ensure tokens are distributed more evenly across all experts in the decoder (Shazeer et al., 2017; Fedus et al., 2022).

Original: Upcycling was performed on TPU v4 accelerators using 64 chips for Base and Large and 256 chips for XL.
Revised: Upcycling was performed on TPU v4 accelerators using 64 chips for Base and Large and 256 chips for XL.

Original: All sizes used expert partitioning, but only XL used model partitioning, with 4 partitions.
Revised: All sizes used expert partitioning, but only XL used model partitioning, with 4 partitions.

Original: A.1. V ISION
Revised: A.1. V ISION

Original: Dense ViT models are pretrained on JFT300M Sun et al.
Revised: Dense ViT models are pretrained on JFT300M Sun et al.

Original: We train with Adafactor Shazeer &
Revised: We train with Adafactor (Shazeer &

Original: Stern (2018), and decoupled weight decay (magnitude 3 on head and 0.03 on body) following Zhai et al. (2022).
Revised: Stern, 2018), and decoupled weight decay (magnitude 3 on head and 0.03 on body) following Zhai et al. (2022).

Original: We use a batch size of 4096.
Revised: We use a batch size of 4096.

Original: The learning rate schedule consists of a learning warmup of 10 000 steps, followed by reverse square root decay with timescale 100 000 steps and ending with a linear cooldown to 0 over 50 000 steps.
Revised: The learning rate schedule consists of a learning warmup of 10 000 steps, followed by reverse square root decay with timescale 100 000 steps and ending with a linear cooldown to 0 over 50 000 steps.

Original: We use a ﬁxed peak learning rate of 4 · 10 − 4 .
Revised: We use a ﬁxed peak learning rate of 4 · 10 − 4 .

Original: 5 Models employing a patch size of 32 (i.e. B/32, L/32) were trained for a total of 14 epochs, while those employing a patch size of 16 (i.e. B/16, L/16) were trained for 28 epochs.
Revised: 5 Models employing a patch size of 32 (i.e. B/32, L/32) were trained for a total of 14 epochs, while those employing a patch size of 16 (i.e. B/16, L/16) were trained for 28 epochs.

Original: By default, we begin upcycling half-way through the total number of epochs in each case.
Revised: By default, we begin upcycling half-way through the total number of epochs in each case.

Original: This roughly corresponds to the number of epochs used to train the corresponding variant in ViT (Dosovitskiy et al., 2021).
Revised: This roughly corresponds to the number of epochs used to train the corresponding variant in ViT (Dosovitskiy et al., 2021).

Original: For
Revised: For

Original: B/16, for example, we assume that a dense checkpoint trained for 14 epochs is given, and we either continue training or apply upcycling for another 14 epochs.
Revised: B/16, for example, we assume that a dense checkpoint trained for 14 epochs is given, and we either continue training or apply upcycling for another 14 epochs.

Original: A. M ODEL TRANSFER A.2. L ANGUAGE
Revised: A. M ODEL TRANSFER A.2. L ANGUAGE

Original: For ﬁnetuning on SuperGLUE, we generally adopt the conventional setup (Raffel et al., 2020; Narang et al., 2021) where we ﬁnetune on all SuperGLUE tasks, in a proportional mix, for 200K steps with a batch size of 128.
Revised: For ﬁnetuning on SuperGLUE, we generally adopt the conventional setup (Raffel et al., 2020; Narang et al., 2021) where we ﬁnetune on all SuperGLUE tasks, in a proportional mix, for 200K steps with a batch size of 128.

Original: Each example has input length 512 and target decoding length of 62.
Revised: Each example has input length 512 and target decoding length of 62.

Original: For ﬁnetuning Dense models on SuperGLUE, we use Adafactor with the default, constant learning rate of 10 − 3 and a dropout rate of 0 .
Revised: For ﬁnetuning Dense models on SuperGLUE, we use Adafactor with the default, constant learning rate of 10 − 3 and a dropout rate of 0 .

Original: 1 (Raffel et al., 2020; Narang et al., 2021).
Revised: 1 (Raffel et al., 2020; Narang et al., 2021).

Original: For ﬁnetuning upcycled models, because there are many more parameters, it can be helpful to increase the dropout rate for the experts (Fedus et al., 2022), while using the default dropout rate of 0 .
Revised: For ﬁnetuning upcycled models, because there are many more parameters, it can be helpful to increase the dropout rate for the experts (Fedus et al., 2022), while using the default dropout rate of 0 .

Original: 1 for all “dense” parameters.
Revised: 1 for all “dense” parameters.

Original: For upcycled Base models, we obtained the strongest results for a constant learning rate of 10 − 4 with an expert dropout rate of 0 . 1 .
Revised: For upcycled Base models, we obtained the strongest results for a constant learning rate of 10 − 4 with an expert dropout rate of 0 . 1 .

Original: For upcycled Large models, we found slightly stronger results with a learning rate of 10 − 3 and an expert dropout rate of 0 . 3 .
Revised: For upcycled Large models, we found slightly stronger results with a learning rate of 10 − 3 and an expert dropout rate of 0 . 3 .

Original: Decreasing (or increasing) the learning rate was not helpful for the Dense Base or Large models.
Revised: Decreasing (or increasing) the learning rate was not helpful for the Dense Base or Large models.

Original: A.2. V ISION
Revised: A.2. V ISION

Original: Few-shot linear evaluation .
Revised: Few-shot linear evaluation .

Original: The fewshot evaluation poses classiﬁcation as a linear regression task, where inputs are frozen representations computed by a pretrained model, and outputs are one-hot vectors representing the ground truth (Dosovitskiy et al., 2021).
Revised: The fewshot evaluation poses classiﬁcation as a linear regression task, where inputs are frozen representations computed by a pretrained model, and outputs are one-hot vectors representing the ground truth (Dosovitskiy et al., 2021).

Original: There are two key changes compared to prior works which used this method (Dosovitskiy et al., 2021;
Revised: There are two key changes compared to prior works which used this method (Dosovitskiy et al., 2021;

Original: Riquelme et al., 2021):
Revised: Riquelme et al., 2021):

Original: • Multiple seeds.
Revised: • Multiple seeds.

Original: The evaluation involves randomly selecting N examples per class.
Revised: The evaluation involves randomly selecting N examples per class.

Original: To reduce dependency on that choice, we run 5 random seeds, and report the average test accuracy across them.
Revised: To reduce dependency on that choice, we run 5 random seeds, and report the average test accuracy across them.

Original: • Fixed L2 regularization.
Revised: • Fixed L2 regularization.

Original: Prior works considered a range of L2 regularizations.
Revised: Prior works considered a range of L2 regularizations.

Original: The optimal value was picked based on average test accuracy across all datasets considered for few-shot evaluation.
Revised: The optimal value was picked based on average test accuracy across all datasets considered for few-shot evaluation.

Original: We ﬁx the L2 regularisation at 1024.
Revised: We ﬁx the L2 regularisation at 1024.

Original: Full ﬁnetuning .
Revised: Full ﬁnetuning .

Original: We ﬁnetune models on ImageNet2012 using SGD and a batch size of 512.
Revised: We ﬁnetune models on ImageNet2012 using SGD and a batch size of 512.

Original: We use a cosine decay learning rate schedule with a linear warmup.
Revised: We use a cosine decay learning rate schedule with a linear warmup.

Original: We sweep over two training schedules: (i) 5k steps, with a warmup of 200 steps, and (ii) 10k steps, with a warmup of 400 steps.
Revised: We sweep over two training schedules: (i) 5k steps, with a warmup of 200 steps, and (ii) 10k steps, with a warmup of 400 steps.

Original: Alongside this we sweep over learning rates [0.1, 0.03, 0.01, 0.003, 0.001, 0.0003].
Revised: Alongside this we sweep over learning rates [0.1, 0.03, 0.01, 0.003, 0.001, 0.0003].

Original: For each pretrained model, there are therefore 12 ﬁnetuning sweeps; we select based on optimal validation accuracy, and report the corresponding test accuracy.
Revised: For each pretrained model, there are therefore 12 ﬁnetuning sweeps; we select based on optimal validation accuracy, and report the corresponding test accuracy.

Original: A. M ODEL PARAMETERS
Revised: A. M ODEL PARAMETERS

Original: Table 1 gives the number of parameters for models used in the main text.
Revised: Table 1 gives the number of parameters for models used in the main text.

Original: [Table]
Revised: [Table]

Original: A. P ARALLELIZATION STRATEGIES
Revised: A. P ARALLELIZATION STRATEGIES

Original: Sparsely activated Mixture-of-Experts (MoE) models combine three types of parallelization strategies to train large models across multiple accelerator chips: data, model and expert parallelism.
Revised: Sparsely activated Mixture-of-Experts (MoE) models combine three types of parallelization strategies to train large models across multiple accelerator chips: data, model and expert parallelism.

Original: We use data parallelism to shard the training batch across devices.
Revised: We use data parallelism to shard the training batch across devices.

Original: We use expert parallelism to partition experts across devices; for example, placing experts 1 and 2 on device 1, experts 3 and 4 on device 2, and so on.
Revised: We use expert parallelism to partition experts across devices; for example, placing experts 1 and 2 on device 1, experts 3 and 4 on device 2, and so on.

Original: Model parallelism is a third axis along which model weights (matrices) can be sharded across devices; for example, expert 1 is split across devices 1 and 2, expert 2 is split across devicesand
Revised: Model parallelism is a third axis along which model weights (matrices) can be sharded across devices; for example, expert 1 is split across devices 1 and 2, expert 2 is split across devicesand

Original: [Table]
Revised: [Table]

Original: and so on.
Revised: and so on.

Original: Model parallelism is beneﬁcial for scaling to larger model sizes.
Revised: Model parallelism is beneﬁcial for scaling to larger model sizes.

Original: See also (Fedus et al., 2022) for a more detailed discussion of these three parallelization strategies.
Revised: See also (Fedus et al., 2022) for a more detailed discussion of these three parallelization strategies.

Original: B A BLATIONS AND A DDITIONAL E XPERIMENTS
Revised: B A BLATIONS AND A DDITIONAL E XPERIMENTS

Original: In this section, we present results for a number of model ablations that try to identify good choices for the main upcycling algorithm decisions.
Revised: In this section, we present results for a number of model ablations that try to identify good choices for the main upcycling algorithm decisions.

Original: As mentioned in the main text, unless stated otherwise, vision ablations use a B/16 sparse model with 32 experts, C = 1 and 6 MoE layers placed in the last few block of the model.
Revised: As mentioned in the main text, unless stated otherwise, vision ablations use a B/16 sparse model with 32 experts, C = 1 and 6 MoE layers placed in the last few block of the model.

Original: The dense checkpoint was trained for 14 epochs, and we train for an additional 7 epochs (up to a total of 21 epochs).
Revised: The dense checkpoint was trained for 14 epochs, and we train for an additional 7 epochs (up to a total of 21 epochs).

Original: Note that, for C = 1 , comparing performance on a per step basis is a reasonably close approximation of a comparison on a per train time basis.
Revised: Note that, for C = 1 , comparing performance on a per step basis is a reasonably close approximation of a comparison on a per train time basis.

Original: For our language ablations, our default conﬁguration is unchanged: we use a Base model withexperts, C = 2 and 6 MoE layers interspersed throughout the model.
Revised: For our language ablations, our default conﬁguration is unchanged: we use a Base model withexperts, C = 2 and 6 MoE layers interspersed throughout the model.

Original: We train for between 0.
Revised: We train for between 0.

Original: million and 1 million extra steps.
Revised: million and 1 million extra steps.

Original: B. R OUTER TYPE
Revised: B. R OUTER TYPE

Original: While our default upcycling recipe uses Expert Choice routing Zhou et al.
Revised: While our default upcycling recipe uses Expert Choice routing Zhou et al.

Original: (2022) (in the encoder), the same recipe can be applied to other routing mechanisms.
Revised: (2022) (in the encoder), the same recipe can be applied to other routing mechanisms.

Original: Here, we compare with Top-K routing Shazeer et al.
Revised: Here, we compare with Top-K routing Shazeer et al.

Original: (2017), which is a very popular alternative.
Revised: (2017), which is a very popular alternative.

Original: Table 2 shows that, for vision, sparse upcycling with Top-K routing works comparably well to Expert Choice, on a per step basis, provided we also use Batch Priority Routing (BPR)
Revised: Table 2 shows that, for vision, sparse upcycling with Top-K routing works comparably well to Expert Choice, on a per step basis, provided we also use Batch Priority Routing (BPR)

Original: (Riquelme et al., 2021).
Revised: (Riquelme et al., 2021).

Original: BPR sorts tokens according to a model conﬁdence proxy so that –when experts are full– high conﬁdence tokens are given priority.
Revised: BPR sorts tokens according to a model conﬁdence proxy so that –when experts are full– high conﬁdence tokens are given priority.

Original: We suspect this may be helpful right at the beginning, when applying the upcycling, to avoid discarding important tokens.
Revised: We suspect this may be helpful right at the beginning, when applying the upcycling, to avoid discarding important tokens.

Original: Expert Choice avoids this problem by design, as experts are always balanced and select the most ‘relevant’ tokens.
Revised: Expert Choice avoids this problem by design, as experts are always balanced and select the most ‘relevant’ tokens.

Original: Notice that with Expert Choice routing, each token chooses C experts on average.
Revised: Notice that with Expert Choice routing, each token chooses C experts on average.

Original: The initial dense checkpoint was trained for 7 epochs.
Revised: The initial dense checkpoint was trained for 7 epochs.

Original: Note that these comparison are on a per-step basis, and that
Revised: Note that these comparison are on a per-step basis, and that

Original: Expert Choice upcycled models are actually slightly faster than Top-K models; see Figure
Revised: Expert Choice upcycled models are actually slightly faster than Top-K models; see Figure

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: For language, similar ablations (Figure 7) shows that Expert Choice routing outperforms both Toprouting (with BPR) and switch (Top-1) routing, on a train time basis
Revised: For language, similar ablations (Figure 8) shows that Expert Choice routing outperforms both Toprouting (with BPR) and switch (Top-1) routing, on a train time basis

Original: . B. R OUTER CAPACITY FACTOR
Revised: . B. R OUTER CAPACITY FACTOR

Original: Sparsifying the dense model increases the model capacity (number of parameters).
Revised: Sparsifying the dense model increases the model capacity (number of parameters).

Original: However, if the capacity factor C = 1 , then the FLOPS is very similar to the original, dense model (modulo the small routing costs).
Revised: However, if the capacity factor C = 1 , then the FLOPS is very similar to the original, dense model (modulo the small routing costs).

Original: We can increase the per-token compute by increasing C .
Revised: We can increase the per-token compute by increasing C .

Original: Figure 8 investigates this, and shows our results for vision (left and center panels) and language (right panel).
Revised: Figure 9 investigates this, and shows our results for vision (left and center panels) and language (right panel).

Original: For vision, we see that extreme values ( C = 1 and C = 5 ) underperform intermediate values ( C =and C = 3 ) that offer better trade-offs.
Revised: For vision, we see that extreme values ( C = 1 and C = 5 ) underperform intermediate values ( C =and C = 3 ) that offer better trade-offs.

Original: For language, the trend is even stronger: A capacity factor of
Revised: For language, the trend is even stronger: A capacity factor of

Original: C = 2 stands out as the best option on a per compute basis.
Revised: C = 2 stands out as the best option on a per compute basis.

Original: B. N UMBER OF EXPERTS
Revised: B. N UMBER OF EXPERTS

Original: Adding more experts increases the number of model parameters and, up to a point, the quality of the model.
Revised: Adding more experts increases the number of model parameters and, up to a point, the quality of the model.

Original: Given that the number of tokens each expert processes is inversely proportional to the number of experts (see Section 2.1), adding more experts usually only leads to very modest computational (and wall time) overheads.
Revised: Given that the number of tokens each expert processes is inversely proportional to the number of experts (see Section 2.1), adding more experts usually only leads to very modest computational (and wall time) overheads.

Original: However, for a very large number of experts, the upcycled model may experience a larger initial quality drop relative to the baseline dense model.
Revised: However, for a very large number of experts, the upcycled model may experience a larger initial quality drop relative to the baseline dense model.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: a t i o n
Revised: a t i o n

Original: [Figure]
Revised: [Figure]

Original: r
Revised: r

Original: [Figure]
Revised: [Figure]

Original: i
Revised: i

Original: [Figure]
Revised: [Figure]

Original: i
Revised: i

Original: [Figure]
Revised: [Figure]

Original: I L S V R C 10 s h o t
Revised: I L S V R C 10 s h o t

Original: A cc u r a c y
Revised: A cc u r a c y

Original: Figure 9 (two left panels) shows the results of a vision experiment with 6 MoE layers with a number of experts ranging from 2 to 128.
Revised: Figure 10 (two left panels) shows the results of a vision experiment with 6 MoE layers with a number of experts ranging from 2 to 128.

Original: For a ﬁxed amount of compute (value in the x-axis), we see that more experts is generally better for this B/16 model.
Revised: For a ﬁxed amount of compute (value in the x-axis), we see that more experts is generally better for this B/16 model.

Original: Figure 10 shows the ﬁnal metric values both for upstream (JFT precision at 1) and downstream (ImageNet 10-shot) with respect to the number of experts.
Revised: Figure 11 shows the ﬁnal metric values both for upstream (JFT precision at 1) and downstream (ImageNet 10-shot) with respect to the number of experts.

Original: We see steady improvements upstream, and –at some point– diminishing returns downstream.
Revised: We see steady improvements upstream, and –at some point– diminishing returns downstream.

Original: B. N UMBER OF M O E LAYERS
Revised: B. N UMBER OF M O E LAYERS

Original: Another key decision is how many layers to sparsify.
Revised: Another key decision is how many layers to sparsify.

Original: More layers leads to higher model capacity, while –especially for higher C – it introduces signiﬁcant extra wall time overhead.
Revised: More layers leads to higher model capacity, while –especially for higher C – it introduces signiﬁcant extra wall time overhead.

Original: We ablate this for vision models, as shown in Figure 9 (two right panels).
Revised: We ablate this for vision models, as shown in Figure 10 (two right panels).

Original: For a B/16 model with 12 blocks, we train upcycled versions with an increasing number of MoE layers; MoE layers are consecutive and start from the last layer.
Revised: For a B/16 model with 12 blocks, we train upcycled versions with an increasing number of MoE layers; MoE layers are consecutive and start from the last layer.

Original: For example, the model labeled as ‘5’ corresponds to a model where the last
Revised: For example, the model labeled as ‘5’ corresponds to a model where the last

Original: MLP layers are sparsiﬁed, and so on.
Revised: MLP layers are sparsiﬁed, and so on.

Original: Thus, model ‘1’ only has one MoE layer (the last one) and it is the computationally cheapest in terms of wall time.
Revised: Thus, model ‘1’ only has one MoE layer (the last one) and it is the computationally cheapest in terms of wall time.

Original: We do not include a model where all layers are sparsiﬁed (would correspond to ‘12’) as we found that sparsifying the very ﬁrst block tends to be problematic.
Revised: We do not include a model where all layers are sparsiﬁed (would correspond to ‘12’) as we found that sparsifying the very ﬁrst block tends to be problematic.

Original: We see in Figure 9 (two right panels) that more MoE layers is not always better even on a per step basis; see Figure 11 for both upstream and downstream metrics.
Revised: We see in Figure 10 (two right panels) that more MoE layers is not always better even on a per step basis; see Figure 12 for both upstream and downstream metrics.

Original: Looking at a ﬁxed value of the x -axis in Figure 9 (right panels), we conclude that something between Last-5 and Last-6 (40-50% of layers sparsiﬁed) offers the most attractive trade-off in this case.
Revised: Looking at a ﬁxed value of the x -axis in Figure 10 (right panels), we conclude that something between Last-5 and Last-6 (40-50% of layers sparsiﬁed) offers the most attractive trade-off in this case.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: B. E XPERT INITIALIZATION
Revised: B. E XPERT INITIALIZATION

Original: The standard upcycling recipe copies and replicates the dense MLP to each expert.
Revised: The standard upcycling recipe copies and replicates the dense MLP to each expert.

Original: As the router directs different tokens to each expert, the experts will start to diverge from one another, and their initial MLP weights.
Revised: As the router directs different tokens to each expert, the experts will start to diverge from one another, and their initial MLP weights.

Original: Figure 12 explores whether loading the MLPs is indeed a good idea, or whether the model would be better off learning the experts from scratch (random initialization).
Revised: Figure 13 explores whether loading the MLPs is indeed a good idea, or whether the model would be better off learning the experts from scratch (random initialization).

Original: We train forextra epochs (dense was trained for 14 epochs, and we keep training up to a total of 21).
Revised: We train forextra epochs (dense was trained for 14 epochs, and we keep training up to a total of 21).

Original: Note that the computational cost of both approaches is identical.
Revised: Note that the computational cost of both approaches is identical.

Original: It takes a long time for the model with randomly initialized experts to recover and catch up with the algorithm that upcycles the expert weights from the dense MLP layers, regardless of the number of experts.
Revised: It takes a long time for the model with randomly initialized experts to recover and catch up with the algorithm that upcycles the expert weights from the dense MLP layers, regardless of the number of experts.

Original: We also tried an intermediate approach (not shown), where we only upcycle a subset of experts and initialize the rest of scratch, but that also underperformed upcycling all of the experts.
Revised: We also tried an intermediate approach (not shown), where we only upcycle a subset of experts and initialize the rest of scratch, but that also underperformed upcycling all of the experts.

Original: [Figure]
Revised: [Figure]

Original: B. R ESUMING THE OPTIMIZER STATE
Revised: B. R ESUMING THE OPTIMIZER STATE

Original: When upcycling a model, we can resume the optimizer state from the original dense checkpoint together with the model parameters.
Revised: When upcycling a model, we can resume the optimizer state from the original dense checkpoint together with the model parameters.

Original: Figure 13 shows that reusing the optimizer state gives a performance boost for vision models, independent of the number of experts.
Revised: Figure 14 shows that reusing the optimizer state gives a performance boost for vision models, independent of the number of experts.

Original: 6 We did not, however, see any improvement from reusing the dense model optimizer state in our language experiments, so we only reuse the optimizer state for vision models.
Revised: 6 We did not, however, see any improvement from reusing the dense model optimizer state in our language experiments, so we only reuse the optimizer state for vision models.

Original: B. C OMBINE WEIGHT NORMALIZATION AFTER ROUTING
Revised: B. C OMBINE WEIGHT NORMALIZATION AFTER ROUTING

Original: A simple trick that we found useful for the upcycling of vision models was to normalize the combine weights after routing.
Revised: A simple trick that we found useful for the upcycling of vision models was to normalize the combine weights after routing.

Original: The weights of each token are normalized so that the sum is 1.
Revised: The weights of each token are normalized so that the sum is 1.

Original: This follows that intuition that each token was previously only processed by a single “expert” MLP in the dense model.
Revised: This follows that intuition that each token was previously only processed by a single “expert” MLP in the dense model.

Original: In the event that a token is not routed at all, the combine weights remain 0.
Revised: In the event that a token is not routed at all, the combine weights remain 0.

Original: We illustrate this normalization trick with two simple examples.
Revised: We illustrate this normalization trick with two simple examples.

Original: Several experts selected .
Revised: Several experts selected .

Original: Suppose a token x is selected by three different experts e 1 , e 2 and e 3 with routing weights w 1 = 0 .
Revised: Suppose a token x is selected by three different experts e 1 , e 2 and e 3 with routing weights w 1 = 0 .

Original: 3 , w 2 = 0 .
Revised: 3 , w 2 = 0 .

Original: 2 , and w 3 = 0 .
Revised: 2 , and w 3 = 0 .

Original: 1 respectively (adding up to 0.6).
Revised: 1 respectively (adding up to 0.6).

Original: The normalized weights are:
Revised: The normalized weights are:

Original: [Equation]
Revised: [Equation]

Original: The ﬁnal output x (cid:48) is:
Revised: The ﬁnal output x (cid:48) is:

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: Only one expert selected .
Revised: Only one expert selected .

Original: In this case, regardless of the selected weight w 1 , the output routing weight will be ¯ w 1 = 1 . 

Revised: In this case, regardless of the selected weight w 1 , the output routing weight will be ¯ w 1 = 1 .

Original: after normalizing it:
Revised: after normalizing it:

Original: [Equation]
Revised: [Equation]

Original: While this approach can be in principle a bit problematic
Revised: While this approach can be in principle a bit problematic

Original: (those tokens only selected by one expert have vanishing routing gradients), Table 3 shows that, even if we are training vision models from scratch, applying weight normalization does not hurt performance (while it indeed helps for upcycling).
Revised: (those tokens only selected by one expert have vanishing routing gradients), Table 3 shows that, even if we are training vision models from scratch, applying weight normalization does not hurt performance (while it indeed helps for upcycling).

Original: However, router weight normalization was not helpful for language models.
Revised: However, router weight normalization was not helpful for language models.

Original: Upstream accuracy after 1M steps was comparable: 70.8% (no normalization) vs 70.7% (normalization), but downstream average scores on SuperGLUE lagged: 79.3% (no normalization) vs 78.8% (normalization).
Revised: Upstream accuracy after 1M steps was comparable: 70.8% (no normalization) vs 70.7% (normalization), but downstream average scores on SuperGLUE lagged: 79.3% (no normalization) vs 78.8% (normalization).

Original: A similar quality degradation were observed in MoE language models trained from scratch.
Revised: A similar quality degradation were observed in MoE language models trained from scratch.

Original: One hypothesis for this different behavior is that the vision models use Expert Choice routing everywhere, but the language models use Expert Choice in the encoder and Top-K routing in the decoder.
Revised: One hypothesis for this different behavior is that the vision models use Expert Choice routing everywhere, but the language models use Expert Choice in the encoder and Top-K routing in the decoder.

Original: 5 Capacity ratio C 0.
Revised: 5 Capacity ratio C 0.

Original: I L S V
Revised: I L S V

Original: R C 10 s h o t
Revised: R C 10 s h o t

Original: A cc u r a c y
Revised: A cc u r a c y

Original: B/16 models
Revised: B/16 models

Original: ExpChoose (Norm)
Revised: ExpChoose (Norm)

Original: ExpChoose
Revised: ExpChoose

Original: Top-K= Top-K=
Revised: Top-K= Top-K=

Original: Dense     5 Capacity ratio C 0.
Revised: Dense     5 Capacity ratio C 0.

Original: I L S V
Revised: I L S V

Original: R C 10 s h o t
Revised: R C 10 s h o t

Original: A cc u r a c y
Revised: A cc u r a c y

Original: L/16 models
Revised: L/16 models

Original: B. C LOSER LOOK AT DESIGN CHOICES IMMEDIATELY FOLLOWING UPCYCLING
Revised: B. C LOSER LOOK AT DESIGN CHOICES IMMEDIATELY FOLLOWING UPCYCLING

Original: Once a model is upcycled, there is a drop in performance due to the sudden deviation from the dense model’s learned function, even with MoE experts reusing the dense MLP weights.
Revised: Once a model is upcycled, there is a drop in performance due to the sudden deviation from the dense model’s learned function, even with MoE experts reusing the dense MLP weights.

Original: Routing design decisions can make a signiﬁcant difference in ameliorating this drop.
Revised: Routing design decisions can make a signiﬁcant difference in ameliorating this drop.

Original: Section B.2 ablates the long-term upcycled model performance as a function of the capacity factor
Revised: Section B.2 ablates the long-term upcycled model performance as a function of the capacity factor

Original: C .
Revised: C .

Original: Figure 14 shows the immediate effect of modifying C on the upcycled model.
Revised: Figure 15 shows the immediate effect of modifying C on the upcycled model.

Original: Increasing C reduces the likelihood that tokens are dropped; when routing weights are normalized to sum to 1 (see
Revised: Increasing C reduces the likelihood that tokens are dropped; when routing weights are normalized to sum to 1 (see

Original: Section B.7)
Revised: Section B.7)

Original: , the upcycled model is exactly equivalent to the dense model for those tokens selected by at least one expert.
Revised: , the upcycled model is exactly equivalent to the dense model for those tokens selected by at least one expert.

Original: Note the set of tokens not selected by any expert decreases in size as we increase
Revised: Note the set of tokens not selected by any expert decreases in size as we increase

Original: C .
Revised: C .

Original: Note that although different routing mechanisms signiﬁcantly impact the starting point for upcycling, the subsequent training can smooth over many differences.
Revised: Note that although different routing mechanisms signiﬁcantly impact the starting point for upcycling, the subsequent training can smooth over many differences.

Original: For example, at the start Top-K routing is clearly worse than Expert Choice with weight normalization, but Table 2 shows that a model upcycled with Top-K routing eventually catches up.
Revised: For example, at the start Top-K routing is clearly worse than Expert Choice with weight normalization, but Table 2 shows that a model upcycled with Top-K routing eventually catches up.

Original: Figure 15 shows the effect of the group size parameter.
Revised: Figure 16 shows the effect of the group size parameter.

Original: Routing, and in particular the top-k operations, are performed in groups.
Revised: Routing, and in particular the top-k operations, are performed in groups.

Original: Using smaller group sizes will speed up routing, but will lead to higher variance in expert assignment across groups.
Revised: Using smaller group sizes will speed up routing, but will lead to higher variance in expert assignment across groups.

Original: For smaller groups, we may expect more tokens to be dropped (not selected by any expert).
Revised: For smaller groups, we may expect more tokens to be dropped (not selected by any expert).

Original: How many MoE layers we upcycle, and where we place those layers, also affects the initial performance drop.
Revised: How many MoE layers we upcycle, and where we place those layers, also affects the initial performance drop.

Original: Figure 16 shows the effect of this in the initial drop for Expert Choice routing, using normalized combined weights (Section B.7) Upcycling the bottom layers causes a larger initial performance drop.
Revised: Figure 17 shows the effect of this in the initial drop for Expert Choice routing, using normalized combined weights (Section B.7) Upcycling the bottom layers causes a larger initial performance drop.

Original: Upcycling the last layers consecutive layers or interleaving them –as in every other layer– yields the smallest initial performance drop.
Revised: Upcycling the last layers consecutive layers or interleaving them –as in every other layer– yields the smallest initial performance drop.

Original: Finally, we analyze how the number of experts per MoE layer affects the initial upcycling.
Revised: Finally, we analyze how the number of experts per MoE layer affects the initial upcycling.

Original: Figuresuggests that routing to more experts leads to a heavier drop.
Revised: Figuresuggests that routing to more experts leads to a heavier drop.

Original: Figure 11 shows that upcycled models can recover from this eventually though, and sometimes even achieve higher performance.
Revised: Figure 12 shows that upcycled models can recover from this eventually though, and sometimes even achieve higher performance.

Original: B. T HINGS THAT DID NOT WORK
Revised: B. T HINGS THAT DID NOT WORK

Original: We list several unsuccessful attempts, beyond the preceding ablations, to improve the performance of the upcycling.
Revised: We list several unsuccessful attempts, beyond the preceding ablations, to improve the performance of the upcycling.

Original: Adding (truncated) Gaussian noise to router weights, expert weights or both, in an effort to diversify the router or expert weights, did not help.
Revised: Adding (truncated) Gaussian noise to router weights, expert weights or both, in an effort to diversify the router or expert weights, did not help.

Original: Modiﬁying the learning rate of experts, routers or both, to account for the fact that the addition of sparsity may demand different learning rate from the rest of the model, generaly hurt performance; increasing the learning rates sometimes rendered the models more unstable.
Revised: Modiﬁying the learning rate of experts, routers or both, to account for the fact that the addition of sparsity may demand different learning rate from the rest of the model, generaly hurt performance; increasing the learning rates sometimes rendered the models more unstable.

Original: We attempted to add a ﬁxed temperature term to the softmax of
Revised: We attempted to add a ﬁxed temperature term to the softmax of

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: router to encourage more random or less random routing.
Revised: router to encourage more random or less random routing.

Original: We also varied the initialization scheme for the router, included the router z-loss (Zoph et al., 2022).
Revised: We also varied the initialization scheme for the router, included the router z-loss (Zoph et al., 2022).

Original: Unfortunately, none of these attempts led to any signiﬁcant performance improvement.[Figure] [Figure]
Revised: Unfortunately, none of these attempts led to any signiﬁcant performance improvement.

Original: U nd e r r e
Revised: U nd e r r e

Original: v
Revised: v

Original: i e w a s a c on f e r e n ce p a p e r a t
Revised: i e w a s a c on f e r e n ce p a p e r a t

Original: I C L
Revised: I C L

Original: R
Revised: R

Original: C S ELECTED RESULTS
Revised: C S ELECTED RESULTS

Original: [Table]
Revised: [Table]

Original: MoE
Revised: MoE

Original: B/ 38.
Revised: B/ 38.

Original: Upcycling
Revised: Upcycling

Original: B/ 47.
Revised: B/ 47.

Original: Upcycling
Revised: Upcycling

Original: B/ 46.
Revised: B/ 46.

Original: Upcycling
Revised: Upcycling

Original: L/ 53.
Revised: L/ 53.

Original: Upcycling
Revised: Upcycling

Original: L/ 51.
Revised: L/ 51.

Original: Dense
Revised: Dense

Original: B/ 48.
Revised: B/ 48.

Original: MoE
Revised: MoE

Original: B/ 50.
Revised: B/ 50.

Original: Upcycling
Revised: Upcycling

Original: B/ 52.
Revised: B/ 52.

Original: Upcycling
Revised: Upcycling

Original: B/ 54.
Revised: B/ 54.

Original: Dense
Revised: Dense

Original: L/ 54.
Revised: L/ 54.

Original: Upcycling
Revised: Upcycling

Original: L/ 56.
Revised: L/ 56.

Original: Dense
Revised: Dense

Original: L/ 55.
Revised: L/ 55.

Original: Upcycling
Revised: Upcycling

Original: L/ 57.
Revised: L/ 57.

Original: Dense
Revised: Dense

Original: L/ 56.
Revised: L/ 56.

Original: U nd e r r e
Revised: U nd e r r e

Original: v
Revised: v

Original: i e w a s a c on f e r e n ce p a p e r a t
Revised: i e w a s a c on f e r e n ce p a p e r a t

Original: I C L
Revised: I C L

Original: R
Revised: R

Original: [Table]
Revised: [Table]