PaperlayerView exact revision on arXiv

Source-linked AI summary

Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints

Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, Neil Houlsby

arXiv:2212.05055v2cs.LGcs.CLcs.CV

TL;DR

Training large neural networks is increasingly expensive, while sparse models remain costly to train from scratch. The paper introduces sparse upcycling, which initializes sparsely activated MoEs from dense checkpoints, and reports gains over dense counterparts across language and vision within roughly half the original pretraining cost.

  • Problem

    Training large neural networks requires rapidly increasing computation, while sparse models remain costly and data-hungry to train from scratch.

  • Method

    Sparse upcycling warm-starts a larger, sparsely activated Mixture-of-Experts model from a pretrained dense Transformer checkpoint.

  • Results

    Upcycled T5 and Vision Transformer models significantly outperform dense counterparts on SuperGLUE and ImageNet using approximately 50% of the initial dense pretraining sunk cost.

  • Takeaways & Limitations

    Reusing dense-model computation and weights can provide a smooth transition to sparse models with greater capacity and flexibility at inference.

Abstract

from arXiv · show

Training large, deep neural networks to convergence can be prohibitively expensive. As a result, often only a small selection of popular, dense models are reused across different contexts and tasks. Increasingly, sparsely activated models, which seek to decouple model size from computation costs, are becoming an attractive alternative to dense models. Although more efficient in terms of quality and computation cost, sparse models remain data-hungry and costly to train from scratch in the large scale regime. 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. We show that sparsely upcycled T5 Base, Large, and XL language models and Vision Transformer Base and Large models, respectively, significantly outperform their dense counterparts on SuperGLUE and ImageNet, using only ~50% of the initial dense pretraining sunk cost. The upcycled models also outperform sparse models trained from scratch on 100% of the initial dense pretraining computation budget.

1 INTRODUCTION

Training larger models improves performance but increasingly requires costly computation, motivating methods that reuse pretrained dense checkpoints to build sparse MoEs more cheaply. Sparse upcycling improves both language and vision models within modest additional budgets.

  • Motivation: Training cost has risen sharply with model scale, limiting large-model research to teams with substantial computational resources.The cited examples increase from 0.5 ZFLOPS for BERT-Large to 2527 ZFLOPS for PaLM.
  • Approach: Sparse upcycling upgrades an existing dense model into a larger, sparsely activated Mixture-of-Experts model using a relatively small additional computational budget.The approach assumes a pretrained dense Transformer checkpoint and uses it to warm-start MoE training.
  • Motivation: A pretrained dense model can first be trained and later converted into an MoE when its performance saturates, preserving flexibility between dense and sparse alternatives.This scenario avoids requiring the dense-versus-MoE choice before the initial training run.
  • Results: 13% extra training time with upcycling reaches at least a 1% ImageNet 10-shot improvement for ViT-B/16, versus 58% for continued dense training.The comparison is relative to the original checkpoint.
  • Results: Upcycled T5-Large and T5-Base outperform their dense counterparts by 1.5-2 absolute SuperGLUE points using 46% and 55% extra training, respectively.These results fall within the reported +10% to +60% budget range.

2 BACKGROUND

Sparsely activated MoEs increase model capacity while activating only selected parameters per input. The paper applies this architecture and its routing variants to upcycled T5 and Vision Transformer models.

  • Sparse Mixture-of-Experts: Dense models apply all parameters to every input, whereas sparse models activate only a subset to increase capacity with lower computation.Sparsely activated MoEs are presented as an accelerator-friendly sparse family supporting models with up to trillions of parameters.
  • Sparse Mixture-of-Experts: MoE blocks replace Transformer MLPs with multiple parameterized experts and a router that selects which expert processes each token.The background lists Top-K, BASE, Sinkhorn-BASE, Hash, and Expert Choice routing algorithms.
  • Expert Choice Routing: Expert Choice routing assigns each expert its top-T tokens by routing probability, with T = C(n/E), so token coverage can overlap or leave some tokens unprocessed.At C = 1, each expert processes exactly n/E tokens; larger C usually improves performance at higher compute cost.
  • Experimental Architectures: The experiments use T5 encoder-decoder models and Vision Transformer encoder models, applying the sparse upcycling recipe across language and vision.The studied T5 variants are Base, Large, and XL; ViT variants include B/32, B/16, L/32, and L/16.
  • Experimental Architectures: Vision models use Expert Choice routing, while language models use it in encoders and Top-K routing with K = 2 in decoders.The decoder choice addresses differences between full-batch teacher forcing during training and single-token autoregressive decoding during inference.

3 THE UPCYCLING ALGORITHM

The upcycling algorithm expands selected dense MLP layers into MoE layers while copying the remaining model components from a dense checkpoint. Its design balances added capacity and performance against computational cost and the initial quality drop.

  • Initialization: Upcycling preserves the Transformer block structure, copies non-MoE components, and initializes every new expert as a copy of the replaced dense MLP.The new MoE layers also require a router, which is absent from the original dense architecture.
  • Design Trade-offs: Increasing MoE layers, experts, or expert capacity generally raises quality but also increases computation and can deepen the initial quality drop.The trade-off results from more drastic reconfiguration and larger computational requirements.
  • Routing: Expert Choice routing with capacity factor C = 2 is used for vision and language encoders, while language decoders use Top-K routing with K = 2.The paper reports that Expert Choice outperforms standard Top-K for upcycling, while both outperform dense continuations.
  • Layer Expansion: Replacing half of the MLP layers with MoE layers is the default because more replacements increase capacity and cost while causing a larger initial quality decline.This choice is based on the paper’s ablation and prevailing MoE conventions.
  • Expert Count: 32 experts provide a reported compromise in the studied regime of +20% to +100% of the initial dense baseline model’s computational cost.The paper also explores other expert counts in ablations.
  • Expert Capacity: Capacity factor C controls how many experts process each token on average; larger values improve quality but increase FLOPs and runtime.The authors find C = 2 generally offers good quality relative to compute time.
  • Training State: Reusing optimizer state boosts vision performance but provides no improvement in the reported language experiments.Consequently, optimizer state is reused only for vision models.
  • Weight Normalization: Router-weight normalization helps upcycled vision models but hurts upcycled language models.The procedure normalizes each token’s router combine weights to 1 and is used only for vision in the reported setup.

4 EXPERIMENTS

Experiments compare sparse upcycling with dense continuation, scratch-trained MoEs, and dense upcycling across vision and language tasks, using quality-versus-compute evaluations and ablations.

  • Experimental setup: Quality-versus-cost plots evaluate upstream or downstream performance against TPU-core-days or training steps.The experiments use training steps when compared models have the same cost per step.
  • Core results: Upcycled models show clear gains over dense continuation after a non-trivial amount of extra compute across vision and language pretraining.At very limited additional training, the methods remain close to the original checkpoint quality.
  • Core results: Upstream vision gains generally transfer to ImageNet, while language finetuning is more variable but still favors upcycled models on SuperGLUE.Figure 3 reports full-finetuning performance after the pretraining comparisons.
  • Core results: About 120% of the original dense checkpoint’s computation budget is required for a scratch-trained language MoE to catch up with the upcycled model.Scratch-trained MoEs can gain quality faster per step, attributed to their larger learning rate and independently diversifying experts.
  • Core results: For computation budgets below 100% of the initial dense model’s budget, sparse upcycling is more resource-efficient; at larger budgets, scratch-trained MoEs may become preferable.The scratch-trained model is expected to eventually catch up when the available computation exceeds the initial budget.
  • Core results: Sparse upcycling outperforms dense upcycling from a T5 Base checkpoint, while dense upcycling uses depth tiling rather than width tiling.The dense-upcycling models quickly improve over the original checkpoint but remain below the sparse model.
  • Ablations: Upcycling gains remain fairly consistent across different amounts of initial dense pretraining, while Expert Choice routing is preferred over slower Top-K routing on a train-time basis.The ablation uses C = 1 for comparable per-step FLOPS and runtime; separately, C = 2 generally gives the best quality per compute time.

5 RELATED WORK

Related work reduces training cost through parameter reuse, progressive model growth, pruning, and evolving sparse architectures. Sparse upcycling differs by growing trained dense models into sparse models within a fraction of the original training budget.

  • Parameter reuse: Prior warm-start methods reuse trained parameters through function-preserving initialization, shared-layer training, or gradual parameter unsharing.These approaches target deeper or wider models or reduce training iterations by reusing existing parameters.
  • Progressive growth: Progressive-growth methods decompose training into stages that enlarge models by copying weights or stacking new layers.Some stages update only newly added layers to reduce cost.
  • Sparse upcycling: Sparse upcycling targets compute budgets that are a fraction of original training and preserves sparse inference cost while adding model capacity.This combines dense-checkpoint reuse with growth into a sparsely activated architecture.
  • Pruning: Unlike pruning, which typically constructs smaller and faster models, sparse upcycling grows dense models into larger sparse models.The distinction is architectural direction: expansion rather than post-training reduction.
  • Sparsely activated MoE: MoE models increase parameter count with sublinear computation growth by routing each input through only a subset of parameters.The paper applies this sparsification to existing dense models.

6 CONCLUSIONS

Sparse upcycling reuses pretrained dense checkpoints to initialize more powerful sparse models, extending model capacity and flexibility without repeating full training costs. The paper’s experiments support strong gains across vision and language, while its ablations identify choices that shape performance–cost trade-offs.

  • Sparse upcycling reuses pretrained dense checkpoints to initialize more powerful sparsely activated Mixture-of-Experts models.
  • The recipe provides a smooth transition to sparse models with more capacity and flexibility at inference.
  • Experiments across vision and language models at various scales show large performance gains relative to continuing dense models.
  • Ablations identify algorithmic choices that matter for performance–cost trade-offs under specific compute budgets.
  • Upcycling techniques can support a dynamic model ecosystem by growing existing models, replicating components, and resuming training after model surgery.

A TRAINING AND EVALUATION DETAILS

The experiments combine training curves before and after upcycling to characterize how the method changes model performance over continued training.

  • Figure 7 illustrates combined training curves obtained before and after upcycling.

A.1 UPSTREAM TRAINING

The experiments pretrain dense and upcycled language and vision models under specified optimization, hardware, and evaluation protocols. Upcycled models continue from dense checkpoints with copied model settings and added sparse capacity.

  • Upstream training: Dense Base language models are pretrained from scratch for 1 million steps before upcycled models continue from the resulting checkpoint.
  • Upstream training: Upcycled models use the corresponding dense model’s hyperparameters and continue its inverse square root learning-rate schedule.
  • Upstream training: Every other language-model layer is upcycled with 32 experts, beginning at the second layer, without dropout.
  • Hardware: Upcycling uses TPU v4 accelerators with 64 chips for Base and Large models and 256 chips for XL models.
  • Vision training: Vision models are pretrained on JFT300M with Adafactor, decoupled weight decay, batch size 4096, and a scheduled learning-rate decay.
  • Evaluation: SuperGLUE finetuning averages accuracy across three runs over all tasks trained in a proportional mix for 200K steps.
  • Evaluation: Few-shot linear evaluation averages test accuracy over five random seeds with fixed L2 regularization of 1024.
  • Evaluation: Full ImageNet2012 finetuning sweeps two training schedules and six learning rates, selecting the best validation accuracy.

B ABLATIONS AND ADDITIONAL EXPERIMENTS

The ablations examine routing, capacity, expert count, MoE-layer placement, expert initialization, and optimizer-state reuse across language and vision settings. They reveal configuration-dependent performance–compute trade-offs and favor careful initialization and selective sparsification.

  • Setup: Vision ablations use B/16 models with 32 experts, capacity factor C = 1, and six MoE layers in the final blocks after 14 dense-pretraining epochs.
  • Router type: Top-K routing matches Expert Choice for vision on a per-step basis when combined with Batch Priority Routing.
  • Router type: Expert Choice routing outperforms Top-2 and Switch Top-1 routing for language on a train-time basis.
  • Capacity factor: For vision, intermediate capacity factors C = 2 and C = 3 offer better trade-offs than C = 1 and C = 5; for language, C = 2 is best per compute.
  • Number of experts: With six vision MoE layers and 2–128 experts, more experts generally improve fixed-compute performance, with downstream diminishing returns.
  • Number of MoE layers: Sparsifying the first Transformer block is problematic, so the fully sparsified 12-layer configuration is omitted.
  • Number of MoE layers: For a 12-block B/16 vision model, sparsifying 40–50% of layers, between Last-5 and Last-6, offers the most attractive fixed-compute trade-off.
  • Expert initialization: Randomly initialized experts take a long time to catch up with experts copied from dense MLPs, while partial expert upcycling also underperforms.

B.7 COMBINE WEIGHT NORMALIZATION AFTER ROUTING

Combine-weight normalization rescales routed expert outputs so each token’s weights sum to 1, preserving the dense-model behavior more closely for vision upcycling but not language upcycling.

  • Normalization mechanism: For vision models, combine weights are normalized so each token’s routed weights sum to 1, while unrouted tokens retain weight 0.The trick reflects that dense models previously processed each token through one expert MLP.
  • Normalization mechanism: When three experts select a token with weights 0.3, 0.2, and 0.1, the normalized output combines their expert outputs using weights proportional to those values.The original weights sum to 0.6 before normalization.
  • Normalization mechanism: When only one expert selects a token, normalization sets its output routing weight to 1.0 regardless of the original weight.This can leave such tokens with vanishing routing gradients, but did not hurt training-from-scratch vision performance.
  • Empirical effects: For language models, normalization produced comparable upstream accuracy but lower downstream SuperGLUE scores: 70.8% vs 70.7% upstream and 79.3% vs 78.8% downstream.The comparison is no normalization versus normalization, respectively.
  • Empirical effects: Figure 15 examines how capacity size ratio affects initial B/16 and L/16 performance after upcycling when routing weights are normalized.With large capacity, the upcycled model retains the dense model’s function.

B.8 CLOSER LOOK AT DESIGN CHOICES IMMEDIATELY FOLLOWING UPCYCLING

Immediately after upcycling, performance drops because routing changes the dense model’s learned function; routing and layer-placement choices strongly affect this initial drop, while later training can recover differences.

  • Initial performance: Upcycling causes an initial performance drop because the modified sparse structure deviates from the dense model’s learned function.This occurs even when MoE experts reuse the dense MLP weights.
  • Routing choices: Increasing capacity factor C reduces dropped tokens, and with normalized weights the upcycled model matches the dense model for tokens selected by at least one expert.The set of unselected tokens shrinks as C increases.
  • Routing choices: Top-K routing starts below Expert Choice routing with weight normalization, but subsequent training allows the Top-K upcycled model to catch up.The passage attributes this contrast to starting-point differences and later smoothing during training.
  • Layer placement: Upcycling bottom layers causes a larger initial drop, whereas consecutive final layers or interleaved layers yield the smallest initial drop.This comparison uses Expert Choice routing with normalized combined weights.
  • Number of experts: Routing to more experts produces a heavier initial drop, although upcycled models can eventually recover and sometimes achieve higher performance.The initial effect is reported for ImageNet 10-shot performance, while later recovery is shown separately.
  • Unsuccessful modifications: Adding noise, changing learning rates, and related router modifications did not produce significant performance improvements; larger learning rates sometimes increased instability.These unsuccessful attempts included Gaussian noise, router or expert learning-rate changes, temperature adjustments, and router initialization changes.
  • Routing choices: Increasing group size has little effect with Expert Choice routing but improves initial performance with Top-K routing.Routing operations, including top-k operations, are performed in groups.

C SELECTED RESULTS

The selected-results tables report vision and text benchmark metrics alongside additional computation, including absolute and relative costs compared with corresponding dense checkpoints.

  • Vision results: Table 4 reports selected vision results, including JFT-300M validation precision, ImageNet 10-shot accuracy, fine-tuned ImageNet accuracy, and extra TPUv3-core-days and ExaFLOPs.Costs are reported both absolutely and relative to the corresponding dense checkpoint.
  • Text results: Table 5 reports selected text results, including C4 validation token accuracy, individual SuperGLUE task metrics, overall SuperGLUE score, and extra TPUv4-core-days and ExaFLOPs.The table also reports costs relative to the corresponding dense checkpoint.
  • Comparison convention: The reported relative computation terms are defined with respect to the corresponding dense checkpoint, with zero indicating the initial dense checkpoint.This convention applies to the selected vision and text results.
Loading 2212.05055v2…