Scaling MLPs: A Tale of Inductive Bias — part 2/4
Cleaned-text code points [15347, 28082). Bibliography omitted here; source.md includes it.
As a first starting point, we investigate simple MLPs with ReLU activations and isotropic design, i.e. except for the first, every layer has the same width m ∈ ℕ m\in\mathbb{N} . In order to avoid training instabilities we further enhance the standard MLP with layer normalizations ( Ba et al.,, 2016 ) placed after the activations. We thus compose several blocks of the form Block ( 𝒛 ) = σ ( 𝑾 LN ( 𝒛 ) ) \operatorname{Block}(\bm{z})=\sigma\left(\bm{W}\operatorname{LN}(\bm{z})\right) with 𝑾 ∈ ℝ m × m \bm{W}\in\mathbb{R}^{m\times m} . To embed the image 𝒙 ∈ ℝ d × d × 3 \bm{x}\in\mathbb{R}^{d\times d\times 3} we use a linear layer emb ( 𝒙 ) = 𝑾 e m b vec ( 𝒙 ) \operatorname{emb}(\bm{x})=\bm{W}^{emb}\operatorname{vec}(\bm{x}) with 𝑾 e m b ∈ ℝ m × 3 d 2 \bm{W}^{emb}\in\mathbb{R}^{m\times 3d^{2}} . Such an embedding layer is crucial since for high resolution images, 3 d 2 3d^{2} can be quite large and thus m m needs to be chosen smaller. We empirically find that such a network design is the minimal choice in order to guarantee successful training across all scales of parameter count and sample size. We will use the short cut S-MLP to denote such an architecture.
#### Inverted Bottleneck MLP.
Inspired by Lin et al., (2015) ; Tolstikhin et al., (2021) we add a bottleneck structure to an MLP block as well as skip connections as follows: Block ( 𝒛 ) = 𝒛 + 𝑾 c σ ( 𝑾 e LN ( 𝒛 ) ) \operatorname{Block}(\bm{z})=\bm{z}+\bm{W}^{c}\sigma\left(\bm{W}^{e}\operatorname{LN}\left(\bm{z}\right)\right) where 𝑾 e ∈ ℝ k m × m \bm{W}^{e}\in\mathbb{R}^{km\times m} expands the dimension to k m km for k ∈ ℕ k\in\mathbb{N} and 𝑾 ( c ) ∈ ℝ m × k m \bm{W}^{(c)}\in\mathbb{R}^{m\times km} collapses it back to width m m . For most experiments we set k = 4 k=4 . While the additions of skip connections and bottleneck layers to the architecture arguably add some amount of inductive bias, we believe that in comparison to modern architectures such enhancements remain negligible. We will denote this variant by B-MLP .
## 4 Experiments
### 4.1 Setup
In this work, we solely focus on vision tasks as inductive bias is more readily understood in this setting. Moreover, most theoretical works focus on image classification tasks, making it thus a natural test bed to assess the performance of MLPs. We study the popular tasks CIFAR10, CIFAR100 ( Krizhevsky,, 2009 ) , STL10 ( Coates et al.,, 2011 ) , TinyImageNet ( Le and Yang,, 2015 ) , ImageNet1k for evaluation, as well as ImageNet21k ( Deng et al.,, 2009 ) for pre-training. In order to limit the size of the embedding layer and the computational needs, we downscale all images to resolution 64 × 64 × 3 64\times 64\times 3 (if needed) as done in Chrabaszcz et al., (2017) . We center and normalize all the images as a pre-processing step. For data augmentations, we consider random flips and crops as well as MixUp ( Zhang et al.,, 2018 ) .
### 4.2 Training from Scratch
We start the empirical exploration of MLPs by training them from scratch (i.e. without any extra data) on popular vision benchmarks. All models were trained with the LION optimizer ( Chen et al.,, 2023 ) with a learning rate η = 5 e \eta=5\mathrm{e} - 5 5 . In order to combat overfitting we use strong label smoothing α = 0.3 \alpha=0.3 . We display the resulting test accuracies in Table 1 . We observe that both the standard architecture and the bottleneck without any data augmentation suffer from overfitting, leading to suboptimal performance. When turning it on, data augmentation as a regularizer however really unfolds its full power, significantly pushing the performance by roughly 20 % 20\% across all tasks. As observed in Lin et al., (2015) , the inverted bottleneck architecture leads to an improvement in performance across all datasets. Learning on the other hand significantly slows down with strong augmentations such as MixUp, enabling training for up to 5000 5000 epochs without suffering from overfitting. However, compared to simple modern baselines such as a ResNet18 ( He et al.,, 2015 ) , a large discrepancy in performance remains, highlighting the importance of inductive bias in the small sample regime. We remark that ViTs and MLP-Mixers as well exhibit more learning difficulties if the dataset size is small ( Dosovitskiy et al.,, 2021 ; Tolstikhin et al.,, 2021 ) . We provide more ablation studies in Appendix A.2 .
### 4.3 Transfer Learning
In this section, we aim to analyze how transferable features learnt by MLPs are across different vision tasks. Transferability is one of the hallmark characteristics of modern deep learning, enabling practitioners to fine-tune large models on their specific dataset, leading to superior performance. We are, to the best of our knowledge, the first to measure transferability of MLPs, which is crucial to assess in order to build a theoretical understanding of the process. In this section, we focus on the inverted bottleneck MLP as it generalizes better and is easier to optimize. We provide the dual results for the standard MLP in Appendix B.1 . We restrict to k = 4 k=4 for the expansion factor and denote by B-L/Wi-m a network with L L blocks and width m m . For pre-training we use ImageNet21k, the largest publicly available image dataset with annotated classes. After preprocessing the dataset following Ridnik et al., (2021) , it consists of roughly 12 million images and 11 thousand classes. We then pre-train the MLP with the cross-entropy loss for 800 800 epochs, employing label smoothing and the LION optimizer. To guarantee fast data loading we rely on the FFCV framework ( Leclerc et al.,, 2023 ) for all experiments. In order to measure transferability of the learnt features we fine-tune the network on the new task. We also study training a linear layer on top of the embeddings but defer those results to Appendix A.3 . We again explore the effects of data augmentation during the pre-training stage. For fine-tuning we use SGD with momentum with a learning rate of η head = 0.01 \eta_{\text{head}}=0.01 for the head and η body = 0.001 \eta_{\text{body}}=0.001 for the encoder for 50 50 epochs. We upscale CIFAR images to resolution 64 × 64 × 3 64\times 64\times 3 at fine-tuning time to guarantee compatibility. We display the fine-tuning results in Table 2 . For visualizations of the learnt features, we refer the interested reader to Appendix 12 . We again observe that using data augmentation during the pre-training phase is essential to successful training, boosting performance up to 30 % 30\% in case of CIFAR100. Surprisingly, the learnt features are highly transferable, improving the performances reported previously in Table 1 dramatically. While of course pre-trained on a large quantity of data, we nevertheless want to highlight that such an MLP becomes competitive with a ResNet18 trained from scratch for all the datasets, except for ImageNet1k where performance falls surprisingly short. We hypothesize that MLPs struggle with the more fine-grained distinctions between classes, in combination with the reduced resolution of the images.
#### Test-Time Augmentations.
For ImageNet1k we further notice that objects tend to not be centered, in contrast to datasets like CIFAR10. We suspect that this might lead to the comparatively weaker performance. To test this, we leverage test-time augmentations (TTA). As introduced by Krizhevsky et al., (2012) , for each test image, we produce a fixed number of 100 100 random crops and use the averaged logits for prediction. We observe significant improvements across all datasets, especially for ImageNet we obtain an increase of roughly 8 % 8\% . This indeed indicates that MLPs struggle to localize the object of interest, especially for the more complicated ImageNet1k task. Using a large number of crops alleviates this problem to some degree. This also explains why the gains on tasks like CIFAR10 are smaller as the objects there usually are perfectly centered.
#### ReaL accuary.
As observed in ( Beyer et al.,, 2020 ) , the ImageNet labels do not capture that a single image might contain multiple objects of distinct classes. ImageNet accuracy can thus be misleading in the sense that model classes such as convolutional networks might have implicitly adapted to the particular labeling strategy due to the repeated benchmarking on the same validation set. MLPs most likely lack such an implicit adaptation as this work is to our knowledge the first to evaluate them on ImageNet1k. To address this, Beyer et al., (2020) introduced a novel set of validation labels that better capture the multi-label nature, where a prediction is deemed correct if it matches one of the categories present in the image. We observe further very significant improvements of ≈ 7 % \approx 7\% when employing ImageNet ReaL. Overall, these results underline that a bad inductive bias as exhibited by an MLP can indeed be overcome if subjected to enough scale. For theory, the results are double-edged; while MLPs prove to be a good proxy to understand transfer learning, data augmentation proves to be a crucial component. Also test-time augmentations significantly boost performance. Both these components on the other hand remain rather understudied in theoretical works.
Figure 5: Linear downstream error on CIFAR100 (in % \% ) when pretrained for varying batch-sizes on ImageNet21k, on a log-log scale. Model # \# parameters B-6/Wi-256 9 9 M B-12/Wi-256 12 12 M B-6/Wi-512 24 24 M B-12/Wi-512 37 37 M B-6/Wi-1024 74 74 M B-12/Wi-1024 124 124 M Table 5: The different models and the respective parameter counts in millions.
#### Large batch-sizes.
We further make the counter-intuitive observation that training with larger batch sizes significantly boosts performance both up- and downstream. In Fig. 5 we plot pre-training batch size against resulting linear downstream accuracy on CIFAR100 for different number of pre-training epochs. We observe that across all training times, using a larger batch size leads to significantly better performance. Moreover, we want to highlight that such a plot is even favoring small batch-sizes since those models perform more gradient updates for a fixed number of epochs. This effect is in stark contrast to convolutional architectures where entire lines of works have focused on preserving the performance of the small batch-size regime for larger ones ( Goyal et al.,, 2017 ; You et al.,, 2017 ; Hoffer et al.,, 2017 ; Keskar et al.,, 2017 ) . Training with large batch-sizes without degradation is of high interest as it can lead to potentially more efficient training pipelines since computation can be sharded among more devices. This observation about optimal batch-sizes is in-line with similar recent conclusions in Transformers ( Kaplan et al.,, 2020 ; Touvron et al.,, 2023 ) .
#### Role of augmentations.
The role of data augmentation is very pronounced for MLPs, largely since it provides indirect inductive bias to the model. Remarkably, a model pre-trained on 12 12 million examples without data augmentation shows inferior performance on CIFAR10 compared to a network trained from scratch with augmentations turned on. This emphasizes that augmentations go beyond merely leading to a bigger dataset but provide the model with useful invariances. We investigate the learnt weights in-depth in Appendix 12 , showing that very evidently, more localized features are learnt if data augmentation is employed. The power of augmentations has already been demonstrated previously through the advent of self-supervised learning ( Grill et al.,, 2020 ; Caron et al.,, 2021 ; Chen et al.,, 2020 ) . Even when training on purely random labels, it still provides a powerful learning signal ( Anagnostidis et al.,, 2023 ) .
### 4.4 Scaling Laws
One of the key mysteries in deep learning is that networks tend to improve in terms of generalization when compute, in the form of parameter count and dataset size, is scaled up. Recently it has been observed in several works that the benefits of scale are highly predictable, i.e. generalization performance exhibits a power-law structure when plotted against compute measured in FLOPS ( Rosenfeld et al.,, 2020 ; Hestness et al.,, 2017 ; Hestness et al.,, 2019 ; Kaplan et al.,, 2020 ; Zhai et al.,, 2022 ) . The functional form has recently been further refined ( Caballero et al.,, 2023 ) . The predictable nature of test performance has even been leveraged to estimate the optimal model before training ( Hoffmann et al.,, 2022 ; OpenAI,, 2023 ) . In order to understand this important characteristic of deep learning theoretically, it is important to analyze whether MLPs exhibit similar properties.
#### Compute.