PaperlayerView exact revision on arXiv

Source-linked AI summary

Go Wider Instead of Deeper

Fuzhao Xue, Ziji Shi, Futao Wei, Yuxuan Lou, Yong Liu, Yang You

arXiv:2107.11817v3cs.LGcs.AIcs.CV

TL;DR

WideNet addresses the limited modeling capacity of depth-compressed transformers without incurring the large models and advanced parallelism associated with straightforward width scaling. It shares parameters across depth, replaces FFNs with MoE, and keeps individual normalization layers to diversify semantic representations. Across vision and language benchmarks, WideNet outperforms parameter-matched transformer baselines, including ViT, ViT-MoE, ALBERT, and factorized-embedding BERT variants.

  • Problem

    Depth-wise parameter sharing reduces trainable parameters but limits modeling capacity, while wider MoE models can become too large and require advanced parallelism.

    Source: Introduction · p. 1

  • Method

    WideNet shares transformer parameters across depth, replaces FFN layers with MoE, and retains individual normalization layers to transform diverse semantic representations.

    Source: Introduction · p. 2, Conclusion · p. 7

  • Results

    WideNet achieves the best performance with fewer trainable parameters across computer vision and natural language processing backbones, including 80.1 Top-1 accuracy on ImageNet-1K with 63M parameters.

    Source: Introduction · p. 2, Conclusion · p. 7

  • Takeaways & Limitations

    The results support going wider instead of deeper as a parameter-efficient transformer deployment framework across vision and language tasks.

    Source: Introduction · p. 2, Conclusion · p. 7

    One natural language processing dataset used in the evaluation is not publicly available.

    Source: Natural Language Processing · p. 5

Abstract

from arXiv · show

More transformer blocks with residual connections have recently achieved impressive results on various tasks. To achieve better performance with fewer trainable parameters, recent methods are proposed to go shallower by parameter sharing or model compressing along with the depth. However, weak modeling capacity limits their performance. Contrastively, going wider by inducing more trainable matrixes and parameters would produce a huge model requiring advanced parallelism to train and inference. In this paper, we propose a parameter-efficient framework, going wider instead of deeper. Specially, following existing works, we adapt parameter sharing to compress along depth. But, such deployment would limit the performance. To maximize modeling capacity, we scale along model width by replacing feed-forward network (FFN) with mixture-of-experts (MoE). Across transformer blocks, instead of sharing normalization layers, we propose to use individual layernorms to transform various semantic representations in a more parameter-efficient way. To evaluate our plug-and-run framework, we design WideNet and conduct comprehensive experiments on popular computer vision and natural language processing benchmarks. On ImageNet-1K, our best model outperforms Vision Transformer (ViT) by $1.5\%$ with $0.72 \times$ trainable parameters. Using $0.46 \times$ and $0.13 \times$ parameters, our WideNet can still surpass ViT and ViT-MoE by $0.8\%$ and $2.1\%$, respectively. On four natural language processing datasets, WideNet outperforms ALBERT by $1.8\%$ on average and surpass BERT using factorized embedding parameterization by $0.8\%$ with fewer parameters.

Introduction

WideNet addresses the trade-off between shallow parameter sharing and wider but costly MoE scaling by combining depth sharing with shared MoE capacity and individual normalization layers. The resulting framework targets parameter-efficient modeling across vision and language tasks.

Introduction

  • Shallow parameter-sharing models retain limited modeling capacity, whereas large MoE models require advanced parallelism and do not scale performance linearly.MoE sparsity also scales poorly on relatively small datasets, while depth compression gives all blocks the same knowledge and causes capacity loss.Source: Introduction · p. 1
  • WideNet combines depth-wise parameter sharing, a shared MoE layer, and individual normalization layers to deploy transformer parameters more efficiently.The architecture replaces FFN layers with MoE, shares trainable parameters across blocks except normalization layers, and uses independent normalization to diversify representations.Source: Introduction · pp. 1–2, Figure 1 · p. 2
  • WideNet uses one shared MoE layer instead of one shared FFN layer, allowing every transformer block to access stronger capacity and more experts to model complex representations.Independent normalization layers provide semantic diversification across blocks, while experts receive token subsets corresponding to different input tokens.Source: Introduction · p. 2
  • Sharing the MoE layer across transformer blocks lets experts receive diversified token representations and become more fully trained.The shared experts process representations from different blocks, improving parameter utilization while preserving a stronger width-wise modeling component.Source: Introduction · p. 2
  • Individual normalization layers add few trainable parameters while transforming hidden vectors into different semantic representations for shared attention or MoE layers.This design allows blocks using the same core trainable matrices to model different semantics.Source: Introduction · p. 2
  • WideNet is presented as a parameter-efficient framework evaluated on both computer vision and natural language processing tasks.The authors summarize the framework as going wider instead of deeper and implement it as a transformer backbone named WideNet.Source: Introduction · p. 2

Mixture-of-Experts

WideNet combines depth-wise parameter sharing with width-wise MoE capacity, while retaining individual normalization layers to diversify representations across transformer blocks.

Conditional Computation with MoE

Routing

  • The routing function applies noisy softmax scores followed by TopK selection, producing sparse expert assignments when K is much smaller than E.The routing transformation maps representations from R^D to R^E, and Gaussian noise supports exploration.Source: Routing · p. 3

Balanced Loading

  • Balanced loading combines per-expert capacity limits with a differentiable auxiliary loss so routing avoids both expert overload and underuse.The capacity ratio controls how many tokens each expert preserves, while the loss encourages uniform assignments.Source: Balanced Loading · p. 3

Sharing MoE across transformer blocks

  • WideNet shares the MoE layer, router, and attention parameters across transformer blocks while keeping token representations block-specific.The shared MoE receives different hidden representations at each block, allowing experts to process more varied tokens.Source: Sharing MoE across transformer blocks · pp. 3–4

Individual Layer Normalization

  • Individual layer-normalization parameters diversify inputs to shared attention and MoE layers with only a few additional trainable parameters.Each normalization layer uses trainable vectors γ and β, while shared matrix parameters process more varied semantic representations.Source: Individual Layer Normalization · p. 4

Optimization

  • The optimization objective adds a load-balancing term weighted by λ to the main task loss, with λ set to 0.01 in this work.For supervised image classification, the main loss is cross-entropy.Source: Optimization · p. 4

Experiments

Experiments evaluate WideNet against ViT and ViT-MoE on ImageNet-1K under matched settings, with additional evidence from downstream finetuning benchmarks.

Experimental Settings

  • ImageNet-1K experiments use ViT and ViT-MoE baselines across three WideNet scales with common augmentation and optimizer settings.The models are evaluated on ImageNet-1K, which contains 1k classes and 1.3M images.Source: Experimental Settings · p. 4
  • WideNet is reported as robust to hyper-parameters under baseline-matched pretraining and finetuning settings.Finetuning removes label smoothing and warm-up and uses SGD with momentum.Source: Experimental Settings · p. 4, Comparison with baselines · p. 5

Comparison with baselines

  • WideNet-H exceeds ViT-B by 1.5% Top-1 accuracy on ImageNet-1K with fewer trainable parameters, while WideNet-L surpasses all baselines with 0.13× ViT-L parameters.WideNet-B remains comparable to ViT-L and ViT-MoE-B with over 4× fewer trainable parameters.Source: Table 1 · p. 4, Comparison with baselines · p. 5
  • MoE does not improve ViT on ImageNet-1K, a dataset 200 times smaller than the original ViT-MoE pretraining dataset.The paper contrasts this result with MoE training on JFT-300M and C4.Source: Comparison with baselines · p. 5

Natural Language Processing

WideNet is evaluated on NLP benchmarks using ALBERT-matched pretraining and factorized embeddings, then compared directly on GLUE and SQuAD tasks. With more experts, it outperforms ALBERT and, at 16 experts, BERT with factorized embedding parameterization on all four downstream tasks.

Natural Language Processing

  • NLP experiments use English Wikipedia and BOOKCORPUS for pretraining, with factorized embedding parameterization and ALBERT-matched hyperparameters.The WordPiece embedding size E is 128, and models are pretrained on 256 TPUv3 cores using LAMB.Source: Natural Language Processing · p. 5
  • Evaluation covers GLUE and two versions of the SQuAD dataset, with GLUE results reported as medians over five development-set runs.Source: Natural Language Processing · p. 5
  • One referenced dataset is not publicly available.Source: Natural Language Processing · p. 5

Downstream Evaluation

  • WideNet with 16 experts uses slightly fewer trainable parameters than BERT with factorized embedding parameterization and outperforms BERT on all four downstream tasks.This supports the parameter-efficiency and effectiveness of going wider instead of deeper.Source: Downstream Evaluation · p. 5
  • WideNet with 4 experts surpasses ALBERT by 1.2% in average downstream performance.Source: Downstream Evaluation · p. 5

MoE Analysis

The MoE analysis examines how expert count and routing diversity affect performance. More experts can increase modeling capacity but cause overfitting on ImageNet-1K, while fewer routing groups reduce performance by decreasing token diversity.

MoE Analysis

  • The first experiment scales the number of experts in every transformer block to control the tokens fed into each expert on ImageNet-1K.The analysis approximates the number of training tokens per expert using image count, patch tokens, selected experts, and capacity ratio.Source: MoE Analysis · p. 5
  • More experts increase trainable parameters and modeling capacity but lead to overfitting on ImageNet-1K.Training accuracy is lower than testing accuracy because of data augmentation.Source: MoE Analysis · p. 5, Figure 2 · p. 6
  • WideNet groups transformer blocks by shared routing assignments, with each group containing D/G blocks when there are G groups and D total blocks.Source: MoE Analysis · p. 5
  • Fewer routing groups cause an obvious performance drop because each expert receives less diversified token representations.The analysis concludes that more diversified tokens are required to train MoE models on smaller datasets.Source: MoE Analysis · p. 6

Layer Norm Analysis

The layer-norm analysis studies whether independent normalization diversifies representations when large trainable matrices are shared across transformer blocks. WideNet shows larger γ and β diversity than ViT, while the supplied ablation caption identifies independent normalization, MoE width scaling, and depth compression as the tested modifications.

Layer Norm Analysis

  • WideNet shares MHA and FFN or MoE trainable matrices across blocks, so independent normalization is intended to encourage more diversified inputs across semantic representations.Source: Layer Norm Analysis · p. 6
  • The analysis computes each γ or β element’s distance from corresponding elements across vectors from other transformer blocks.N denotes the number of transformer blocks, and M denotes the vector dimension.Source: Layer Norm Analysis · p. 6
  • WideNet has larger γ and β distance values than ViT, indicating that its MoE receives more diversified input representations.The result supports using individual normalization layers to model varied semantics with shared large trainable matrices.Source: Layer Norm Analysis · p. 6, Figure 3 · p. 6
  • The ablation study evaluates independent layer normalization, MoE-based width scaling, and parameter-sharing-based depth compression on ImageNet-1K.Source: Table 3 · p. 6

Ablation Study: Contributions of key modifications

The ablations test independent normalization, MoE width scaling, and depth compression through parameter sharing, then compare WideNet with ViT under similar computation. These studies show that the modifications improve performance or efficiency, although routing and capacity settings increase training cost.

Ablation Study: Contributions of key modifications

  • Replacing individual layer normalizations with shared ones causes a performance drop with almost the same trainable parameters.The comparison supports independent normalization as an effective design choice.Source: Ablation Study: Contributions of key… · p. 6
  • Removing MoE causes severe optimization or performance problems, including gradient explosion for WideNet-B with substantially fewer trainable parameters.The ablation reports a significant performance drop without MoE.Source: Ablation Study: Contributions of key… · p. 6
  • Removing parameter sharing slightly reduces performance while substantially increasing parameters, and WideNet-H runs out of memory on 256 TPUv3 cores.Parameter sharing therefore contributes to the framework's parameter efficiency and trainability.Source: Ablation Study: Contributions of key… · p. 6

Ablation Study: Comparison with comparable speed or computation cost

  • WideNet uses K=2 selected experts and capacity ratio C=1.2, introducing extra computation relative to vanilla ViT.The second ablation set evaluates the framework against baselines with comparable speed or computation cost.Source: Ablation Study: Comparison with… · p. 7
  • WideNet-L outperforms ViT-L by 0.7% with slightly less training time and 13.1% parameters when using 12 transformer blocks.With comparable per-token computation, widening ViT-L's FFN does not improve its performance, whereas WideNet remains stronger but requires more training time.Source: Ablation Study: Comparison with… · p. 7

Conclusion

The paper proposes going wider instead of deeper through WideNet, combining depth-wise parameter sharing with MoE layers and independent normalization. Across vision and language benchmarks, WideNet achieves stronger performance with fewer trainable parameters.

Conclusion

  • WideNet combines parameter sharing across transformer blocks, MoE layers replacing FFNs, and individual layer normalization to deploy parameters efficiently.The framework is intended to maximize each block's modeling ability while diversifying semantic representations.Source: Conclusion · p. 7
  • WideNet achieves the best performance with fewer trainable parameters across computer vision and natural language processing backbones.On ImageNet-1K, the best model reaches 80.1 Top-1 accuracy with 63M parameters; on four NLP datasets, it outperforms ALBERT and surpasses BERT with fewer parameters.Source: Conclusion · p. 7
  • The investigation reports that MoE scaling does not work well on smaller datasets because each expert requires enough tokens for training.The paper also finds that individual normalization can transform hidden representations into more diversified semantics.Source: Conclusion · p. 7

Appendix Finetuning results on computer vision

WideNet transfers pretrained performance effectively to computer-vision finetuning, with WideNet-L remaining the strongest model on CIFAR10.

Appendix Finetuning results on computer vision

Hyper-parameters of experiments on computer vision

The computer-vision experiments use shared hyper-parameters across models, with cosine learning-rate decay after 30 warmup epochs on ImageNet-1K. The reported setup is described as robust.

Hyper-parameters of experiments on computer vision

Hyper-parameters of experiments on natural language processing

The NLP experiments use task-specific hyper-parameters, including a larger capacity ratio and λ setting for SQuAD1.1 because limited training data may create an unbalanced token distribution.

  • For SQuAD1.1, the experiments set λ to 0 and the capacity ratio C to 2.0.The larger C is motivated by limited training data that may induce an unbalanced token distribution.Source: Hyper-parameters of experiments on… · p. 10
  • Table 7 lists the hyper-parameters used across downstream natural language processing tasks.Source: Table 7 · p. 10
Loading 2107.11817v3…